From 5c4b82e29f06335dc56d710fc4fc7cab50f914c5 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 17:48:14 +0800 Subject: [PATCH 01/22] feat: add HunyuanVideo support for disaggregated diffusion pipeline Switch disagg diffusion demo from Wan2.2-TI2V-5B to HunyuanVideo v1 (13B DiT + Llama3-8B + CLIP dual encoders). Key changes: - Auto-detect dual text encoders from model_index.json - Handle multi-element tensor lists (incompatible shapes) in IntermediateOutputStage and inject_tensors_to_req - Patch HunyuanConfig missing task_type default - Patch SGLang triton norm_infer for non-contiguous tensors - Add .contiguous() in DeviceMoveStage for ZMQ-transferred tensors - Save output as mp4 via imageio instead of raw npy - Make prompt configurable via PROMPT env var Tested: 2 concurrent requests, 4 GPUs (enc=GPU0, den=GPU1,2 TP=2, vae=GPU3), HunyuanVideo 544x960 resolution. Co-Authored-By: Claude Opus 4.6 --- .../phase1_workers/partial_gpu_worker.py | 518 ++++++++++++++++++ .../phase1_workers/run_e2e_sglang.py | 447 +++++++++++++++ .../phase1_workers/sglang_utils.py | 341 ++++++++++++ 3 files changed, 1306 insertions(+) create mode 100644 examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py create mode 100644 examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py create mode 100644 examples/disagg_diffusion/phase1_workers/sglang_utils.py diff --git a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py new file mode 100644 index 000000000000..95e8550121a8 --- /dev/null +++ b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py @@ -0,0 +1,518 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PartialGPUWorker, IntermediateOutputStage, and subprocess launcher. + +This module is the single source of truth for disaggregated diffusion's +integration with sglang. Everything that runs **inside the Scheduler +subprocess** lives here — no Dynamo imports, no heavy dependencies beyond +sglang + torch. + +Architecture:: + + Main process Subprocess (spawned) + ─────────── ──────────────────── + launch_partial_server() ──► _run_partial_scheduler_process() + ├─ monkey-patch GPUWorker → PartialGPUWorker + └─ run_scheduler_process() + └─ Scheduler() + └─ PartialGPUWorker() + └─ pipeline.forward() + SchedulerClient ◄────ZMQ────► Scheduler.event_loop() + +``IntermediateOutputStage`` is appended as the last pipeline stage for +encoder / denoiser. It packages ``Req`` tensors into ``OutputBatch`` +so the result travels through ZMQ as a standard ``OutputBatch``. + +``DecodingStage`` (VAE) already returns ``OutputBatch``, so no extra +stage is needed for the VAE worker. +""" + +from __future__ import annotations + +import logging +import multiprocessing as mp +import os +import sys +from typing import Callable, Dict, List, Optional + +import torch +from setproctitle import setproctitle + +from sglang.multimodal_gen.runtime.distributed import ( + get_tp_rank, + get_tp_world_size, + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_ring_parallel_rank, + get_ring_parallel_world_size, + get_ulysses_parallel_rank, + get_ulysses_parallel_world_size, +) +from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req, OutputBatch +from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +# layerwise_offload may not exist in all sglang versions — guard import +try: + from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +except ImportError: + OffloadableDiTMixin = None # type: ignore[misc,assignment] + +logger = logging.getLogger(__name__) + + +# ═══════════════════════════════════════════════════════════════════════ +# IntermediateOutputStage +# ═══════════════════════════════════════════════════════════════════════ + + +class DeviceMoveStage(PipelineStage): + """Move tensor fields on ``Req`` to the current CUDA device. + + Tensors arriving via ZMQ (pickle) land on CPU. Prepend this stage + at the start of denoiser / VAE pipelines that receive tensors from + other processes. + """ + + def __init__(self, tensor_fields: List[str]): + super().__init__() + self._tensor_fields = tensor_fields + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = torch.device("cuda") + for field in self._tensor_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + setattr(batch, field, [ + t.to(device).contiguous() if isinstance(t, torch.Tensor) else t + for t in val + ]) + elif isinstance(val, torch.Tensor): + setattr(batch, field, val.to(device).contiguous()) + return batch + + +class IntermediateOutputStage(PipelineStage): + """Pipeline stage that packages ``Req`` tensors into ``OutputBatch``. + + Append as the **last** stage in encoder / denoiser partial pipelines. + ``GPUWorker.execute_forward`` sees an ``OutputBatch`` (not a ``Req``) + and returns it through ZMQ without any custom override. + + ``output_fields`` lists the ``Req`` attributes to extract. + Single-element lists are unwrapped to bare tensors. + Multi-element lists (e.g. dual-encoder embeddings) are preserved as + lists so the receiver can reconstruct them correctly. + The handler receives ``OutputBatch.output`` as ``dict[str, Tensor|list]``. + """ + + def __init__(self, output_fields: List[str]): + super().__init__() + self._output_fields = output_fields + + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + tensors: Dict[str, object] = {} + for field in self._output_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + if len(val) == 0: + continue + if len(val) == 1: + tensors[field] = val[0] + else: + # Preserve multi-element lists (e.g. dual-encoder outputs + # with incompatible shapes) so they survive ZMQ pickle. + tensors[field] = val + elif isinstance(val, torch.Tensor): + tensors[field] = val + return OutputBatch(output=tensors) + + +# ═══════════════════════════════════════════════════════════════════════ +# Stage builder functions (picklable — used as subprocess args) +# ═══════════════════════════════════════════════════════════════════════ + + +def build_encoder_stages(pipeline, server_args): + """TextEncodingStage → IntermediateOutputStage(prompt_embeds, …). + + Automatically detects all loaded text encoders/tokenizers so that + both single-encoder (Wan) and dual-encoder (HunyuanVideo) models work. + """ + from sglang.multimodal_gen.runtime.pipelines.stages.text_encoding import ( + TextEncodingStage, + ) + from sglang_utils import get_component_backend + + # Collect all available text encoders and tokenizers + text_encoders = [] + tokenizers = [] + for name in ["text_encoder", "text_encoder_2", "text_encoder_3"]: + enc = pipeline.get_module(name) + if enc is not None: + text_encoders.append(enc) + logger.info("%s backend: %s", name, get_component_backend(enc)) + for name in ["tokenizer", "tokenizer_2", "tokenizer_3"]: + tok = pipeline.get_module(name) + if tok is not None: + tokenizers.append(tok) + + assert len(text_encoders) > 0, "No text encoders found in pipeline" + assert len(text_encoders) == len(tokenizers), ( + f"Encoder/tokenizer count mismatch: {len(text_encoders)} vs {len(tokenizers)}" + ) + + # Always list both fields; IntermediateOutputStage skips None values, + # so negative_prompt_embeds is simply omitted when CFG is disabled. + return [ + TextEncodingStage(text_encoders=text_encoders, tokenizers=tokenizers), + IntermediateOutputStage(["prompt_embeds", "negative_prompt_embeds"]), + ] + + +def build_denoiser_stages(pipeline, server_args): + """DeviceMove → LatentPrep → TimestepPrep → Denoising → IntermediateOutput.""" + from sglang.multimodal_gen.runtime.pipelines.stages.latent_preparation import ( + LatentPreparationStage, + ) + from sglang.multimodal_gen.runtime.pipelines.stages.timestep_preparation import ( + TimestepPreparationStage, + ) + from sglang.multimodal_gen.runtime.pipelines.stages.denoising import ( + DenoisingStage, + ) + from sglang_utils import get_component_backend + + transformer = pipeline.get_module("transformer") + scheduler = pipeline.get_module("scheduler") + logger.info("transformer backend: %s", get_component_backend(transformer)) + + # Fields that arrive from the encoder via ZMQ and need GPU placement. + # prompt_embeds may be a single tensor (Wan) or list of tensors + # (HunyuanVideo dual-encoder) — DeviceMoveStage handles both. + device_move_fields = ["prompt_embeds", "negative_prompt_embeds"] + + return [ + DeviceMoveStage(device_move_fields), + LatentPreparationStage(scheduler=scheduler, transformer=transformer), + TimestepPreparationStage(scheduler=scheduler), + DenoisingStage(transformer=transformer, scheduler=scheduler), + IntermediateOutputStage(["latents"]), + ] + + +def build_vae_stages(pipeline, server_args): + """DecodingStage (already returns OutputBatch — no extra stage needed).""" + from sglang.multimodal_gen.runtime.pipelines.stages.decoding import ( + DecodingStage, + ) + from sglang_utils import get_component_backend + + vae = pipeline.get_module("vae") + logger.info("vae backend: %s", get_component_backend(vae)) + return [DecodingStage(vae=vae, pipeline=pipeline)] + + +# ═══════════════════════════════════════════════════════════════════════ +# PartialGPUWorker +# ═══════════════════════════════════════════════════════════════════════ + + +class PartialGPUWorker(GPUWorker): + """GPUWorker that loads only a subset of pipeline modules. + + Only ``init_device_and_model`` is overridden. Everything else + (``execute_forward``, ``do_mem_analysis``, LoRA, etc.) is inherited. + """ + + def __init__( + self, + required_modules: List[str], + custom_stages_fn: Optional[Callable] = None, + **kwargs, + ): + self._required_modules = required_modules + self._custom_stages_fn = custom_stages_fn + super().__init__(**kwargs) + + def init_device_and_model(self) -> None: + torch.get_device_module().set_device(self.local_rank) + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(self.master_port) + os.environ["LOCAL_RANK"] = str(self.local_rank) + os.environ["RANK"] = str(self.rank) + os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus) + + dist_kwargs = dict( + tp_size=self.server_args.tp_size, + enable_cfg_parallel=self.server_args.enable_cfg_parallel, + ulysses_degree=getattr(self.server_args, "ulysses_degree", 1), + ring_degree=getattr(self.server_args, "ring_degree", 1), + sp_size=self.server_args.sp_degree, + dp_size=self.server_args.dp_size, + distributed_init_method=f"tcp://127.0.0.1:{self.master_port}", + ) + # dist_timeout only in newer sglang versions + import inspect + sig = inspect.signature(maybe_init_distributed_environment_and_model_parallel) + if "dist_timeout" in sig.parameters: + dist_kwargs["dist_timeout"] = self.server_args.dist_timeout + maybe_init_distributed_environment_and_model_parallel(**dist_kwargs) + + if model_parallel_is_initialized(): + suffix = "" + if get_tp_world_size() != 1: + suffix += f"_TP{get_tp_rank()}" + if get_ulysses_parallel_world_size() != 1: + suffix += f"_U{get_ulysses_parallel_rank()}" + if get_ring_parallel_world_size() != 1: + suffix += f"_R{get_ring_parallel_rank()}" + if get_classifier_free_guidance_world_size() != 1: + suffix += f"_C{get_classifier_free_guidance_rank()}" + setproctitle(f"sgl_diffusion::partial{suffix}") + else: + setproctitle(f"sgl_diffusion::partial_{self.local_rank}") + + from sglang_utils import build_partial_pipeline + + self.pipeline = build_partial_pipeline( + self.server_args, + required_modules=self._required_modules, + ) + + if self._custom_stages_fn is not None: + stages = self._custom_stages_fn(self.pipeline, self.server_args) + for stage in stages: + name = type(stage).__name__ + self.pipeline.add_stage(name, stage) + + if getattr(self.server_args, "dit_layerwise_offload", False) and OffloadableDiTMixin is not None: + for module_name in [ + "transformer", "transformer_2", + "video_dit", "video_dit_2", "audio_dit", + ]: + dit = self.pipeline.get_module(module_name) + if dit is not None: + if isinstance(dit, OffloadableDiTMixin): + dit.configure_layerwise_offload(self.server_args) + else: + logger.info( + "Module %s does not support layerwise offload.", + type(dit).__name__, + ) + + logger.info( + "PartialGPUWorker %d: modules=%s stages=%s", + self.rank, + self._required_modules, + list(self.pipeline._stage_name_mapping.keys()), + ) + + +# ═══════════════════════════════════════════════════════════════════════ +# Subprocess entry point + launcher +# ═══════════════════════════════════════════════════════════════════════ + + +def _patch_triton_norm_contiguous(): + """Patch SGLang's triton norm_infer to handle non-contiguous tensors. + + HunyuanVideo's transformer produces non-contiguous intermediate tensors + (e.g. from attention reshapes) which trip the triton kernel assertion + ``assert x.stride(-1) == 1``. + + Must patch both ``triton_ops`` and ``layernorm`` modules because + layernorm uses ``from triton_ops import norm_infer`` (direct binding). + """ + try: + import sglang.multimodal_gen.runtime.layers.triton_ops as triton_ops + import sglang.multimodal_gen.runtime.layers.layernorm as layernorm_mod + + _orig = triton_ops.norm_infer + + def _safe_norm_infer(x, *args, **kwargs): + if not x.is_contiguous(): + x = x.contiguous() + return _orig(x, *args, **kwargs) + + triton_ops.norm_infer = _safe_norm_infer + layernorm_mod.norm_infer = _safe_norm_infer + except (ImportError, AttributeError): + pass + + +def _run_partial_scheduler_process( + required_modules: List[str], + custom_stages_fn: Callable, + local_rank: int, + rank: int, + master_port: int, + server_args: ServerArgs, + pipe_writer, + task_pipe_r, + result_pipe_w, + task_pipes_to_slaves: list, + result_pipes_from_slaves: list, +) -> None: + """Subprocess entry point: patch GPUWorker then delegate to sglang.""" + # Ensure sglang_utils etc. are importable in the subprocess + workers_dir = os.path.dirname(os.path.abspath(__file__)) + if workers_dir not in sys.path: + sys.path.insert(0, workers_dir) + + _patch_triton_norm_contiguous() + + import sglang.multimodal_gen.runtime.managers.scheduler as sched_mod + + _rm, _csf = required_modules, custom_stages_fn + + class _PatchedGPUWorker(PartialGPUWorker): + def __init__(self, **kwargs): + super().__init__(required_modules=_rm, custom_stages_fn=_csf, **kwargs) + + sched_mod.GPUWorker = _PatchedGPUWorker + + from sglang.multimodal_gen.runtime.managers.gpu_worker import ( + run_scheduler_process, + ) + + run_scheduler_process( + local_rank, rank, master_port, server_args, pipe_writer, + task_pipe_r, result_pipe_w, + task_pipes_to_slaves, result_pipes_from_slaves, + ) + + +def launch_partial_server( + server_args: ServerArgs, + required_modules: List[str], + custom_stages_fn: Callable, +) -> list[mp.Process]: + """Spawn Scheduler subprocess(es) with ``PartialGPUWorker``. + + Mirrors sglang's ``launch_server()`` logic: spawns ``num_gpus`` + processes with master/slave pipe wiring so that TP, SP, and CFG + parallelism all work out of the box. + + * ``num_gpus == 1`` → single process, no pipes. + * ``num_gpus > 1`` → rank 0 is master (has ZMQ receiver + pipes to + slaves), ranks 1..N are slaves (coordinate via ``torch.distributed`` + broadcast). + + Returns the process list (for later shutdown). + + After this returns, call ``async_scheduler_client.initialize(server_args)`` + to connect from the main process. + """ + from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger + + configure_logger(server_args) + + num_gpus = server_args.num_gpus + master_port = server_args.master_port + processes: list[mp.Process] = [] + + # --- pipes for master ↔ slave coordination (same as launch_server) --- + task_pipes_to_slaves_w: list = [] + task_pipes_to_slaves_r: list = [] + result_pipes_from_slaves_w: list = [] + result_pipes_from_slaves_r: list = [] + + for _ in range(num_gpus - 1): + r, w = mp.Pipe(duplex=False) + task_pipes_to_slaves_r.append(r) + task_pipes_to_slaves_w.append(w) + + for _ in range(num_gpus - 1): + r, w = mp.Pipe(duplex=False) + result_pipes_from_slaves_r.append(r) + result_pipes_from_slaves_w.append(w) + + # --- spawn one process per GPU --- + readiness_readers: list = [] + + for i in range(num_gpus): + reader, writer = mp.Pipe(duplex=False) + readiness_readers.append(reader) + + if i == 0: + # Rank 0 (master): owns ZMQ receiver + write-ends to slaves + args = ( + required_modules, + custom_stages_fn, + i, # local_rank + i, # rank + master_port, + server_args, + writer, # pipe_writer (readiness signal) + None, # task_pipe_r (master doesn't receive tasks) + None, # result_pipe_w (master doesn't send results) + task_pipes_to_slaves_w, + result_pipes_from_slaves_r, + ) + else: + # Rank > 0 (slave): coordinates via torch.distributed broadcast + args = ( + required_modules, + custom_stages_fn, + i, # local_rank + i, # rank + master_port, + server_args, + writer, # pipe_writer (readiness signal) + None, # task_pipe_r + None, # result_pipe_w + task_pipes_to_slaves_r[i - 1], + result_pipes_from_slaves_w[i - 1], + ) + + process = mp.Process( + target=_run_partial_scheduler_process, + args=args, + name=f"sglang-partial-worker-{i}", + daemon=True, + ) + process.start() + writer.close() + processes.append(process) + + # --- close unused pipe ends in the parent (same as launch_server) --- + for p in task_pipes_to_slaves_w: + p.close() + for p in task_pipes_to_slaves_r: + p.close() + for p in result_pipes_from_slaves_w: + p.close() + for p in result_pipes_from_slaves_r: + p.close() + + # --- wait for all workers to report ready --- + for i, reader in enumerate(readiness_readers): + try: + data = reader.recv() + except EOFError: + processes[i].join() + raise RuntimeError( + f"Partial scheduler rank {i} died (exit code {processes[i].exitcode}). " + "Check logs above for errors." + ) + reader.close() + if data.get("status") != "ready": + raise RuntimeError(f"Partial scheduler rank {i} failed to initialize.") + + pids = [p.pid for p in processes] + logger.info( + "Partial scheduler ready: %d GPU(s), pids=%s", num_gpus, pids, + ) + return processes diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py new file mode 100644 index 000000000000..7278e5baab54 --- /dev/null +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -0,0 +1,447 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end disaggregated diffusion pipeline with SGLang backend. + +Launches Encoder, Denoiser (TP=2), and VAE as separate SGLang scheduler +processes on different GPUs, then runs the full pipeline: + Encoder -> Denoiser -> VAE + +Measures per-stage timing. Supports concurrent requests for benchmarking. + +GPU assignment: + Encoder : GPU 0 (1 GPU) + Denoiser: GPU 1,2 (TP=2) + VAE : GPU 3 (1 GPU) + +Usage: + # Single request (correctness check) + python run_e2e_sglang.py + + # Benchmark: 4 requests, 2 concurrent + NUM_REQUESTS=4 CONCURRENCY=2 python run_e2e_sglang.py + +Environment variables: + MODEL_PATH Model to use (default: hunyuanvideo-community/HunyuanVideo) + PROMPT Text prompt (default: A cat walking on green grass) + GPU_ENC GPU(s) for encoder (default: 0) + GPU_DEN GPU(s) for denoiser (default: 1,2) + GPU_VAE GPU(s) for VAE (default: 3) + TP_SIZE Tensor parallelism for denoiser (default: auto from GPU_DEN) + NUM_REQUESTS Number of pipeline runs (default: 1) + CONCURRENCY Max concurrent pipelines (default: 1) + NUM_FRAMES Number of video frames (default: 61) + NUM_STEPS Denoising steps (default: 50) + HEIGHT Frame height (default: 544) + WIDTH Frame width (default: 960) + GUIDANCE Guidance scale (default: 1.0; >1.0 enables CFG) +""" + +from __future__ import annotations + +import asyncio +import logging +import multiprocessing as mp +import os +import statistics +import sys +import time +from typing import List + +# --- ensure workers dir is on sys.path so subprocesses find sglang_utils --- +WORKERS_DIR = os.path.dirname(os.path.abspath(__file__)) +if WORKERS_DIR not in sys.path: + sys.path.insert(0, WORKERS_DIR) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", +) +logger = logging.getLogger("e2e") + +# ── Configuration ──────────────────────────────────────────────────────── +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +PROMPT = os.environ.get("PROMPT", "A cat walking on green grass") +GPU_ENC = os.environ.get("GPU_ENC", "0") +GPU_DEN = os.environ.get("GPU_DEN", "1,2") +GPU_VAE = os.environ.get("GPU_VAE", "3") +TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(","))))) +NUM_REQUESTS = int(os.environ.get("NUM_REQUESTS", "1")) +CONCURRENCY = int(os.environ.get("CONCURRENCY", "1")) +NUM_FRAMES = int(os.environ.get("NUM_FRAMES", "61")) +NUM_STEPS = int(os.environ.get("NUM_STEPS", "50")) +HEIGHT = int(os.environ.get("HEIGHT", "544")) +WIDTH = int(os.environ.get("WIDTH", "960")) +GUIDANCE = float(os.environ.get("GUIDANCE", "1.0")) +SEED = int(os.environ.get("SEED", "42")) +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_e2e") + + +# ── SGLang compatibility patches ───────────────────────────────────────── + +def _patch_hunyuan_config_task_type(): + """HunyuanConfig inherits ``task_type`` from PipelineConfig without a + default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply + ``task_type=T2V`` when omitted. Idempotent. + """ + from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType + try: + from sglang.multimodal_gen.configs.pipelines.hunyuan import ( + HunyuanConfig, FastHunyuanConfig, + ) + except ImportError: + return + + for cls in (HunyuanConfig, FastHunyuanConfig): + if getattr(cls, "_task_type_patched", False): + continue + orig = cls.__init__ + + def _patched(self, *a, task_type=ModelTaskType.T2V, _orig=orig, **kw): + _orig(self, *a, task_type=task_type, **kw) + + cls.__init__ = _patched + cls._task_type_patched = True + + +def _detect_encoder_modules(model_path: str) -> List[str]: + """Return the required_modules list for the encoder stage.""" + try: + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + maybe_download_model_index, verify_model_config_and_directory, + ) + config = (verify_model_config_and_directory(model_path) + if os.path.exists(model_path) + else maybe_download_model_index(model_path)) + modules = ["text_encoder", "tokenizer"] + if "text_encoder_2" in config: + modules += ["text_encoder_2", "tokenizer_2"] + modules.append("scheduler") + return modules + except Exception: + pass + # Fallback: include dual encoders for known models + if "hunyuan" in model_path.lower(): + return ["text_encoder", "text_encoder_2", + "tokenizer", "tokenizer_2", "scheduler"] + return ["text_encoder", "tokenizer", "scheduler"] + + +# ── Non-singleton ZMQ client ──────────────────────────────────────────── + +class StageClient: + """Async ZMQ REQ client that talks to a SGLang Scheduler.""" + + def __init__(self, endpoint: str, name: str = ""): + import zmq.asyncio + self._name = name + self._ctx = zmq.asyncio.Context() + self._sock = self._ctx.socket(zmq.REQ) + self._sock.connect(endpoint) + self._lock = asyncio.Lock() + logger.info("StageClient(%s) connected to %s", name, endpoint) + + async def forward(self, reqs): + """Send request(s) and receive response.""" + async with self._lock: + await self._sock.send_pyobj(reqs) + return await self._sock.recv_pyobj() + + def close(self): + self._sock.close() + self._ctx.term() + + +# ── Stage launchers ───────────────────────────────────────────────────── + +def _launch_stage( + stage_name: str, + cuda_devices: str, + required_modules: List[str], + custom_stages_fn, + tp_size: int = 1, + scheduler_port: int = 15600, +): + """Launch a partial scheduler for one pipeline stage. + + Returns (processes, server_args). + """ + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_devices + num_gpus = len(cuda_devices.split(",")) + + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from partial_gpu_worker import launch_partial_server + + server_args = ServerArgs.from_kwargs( + model_path=MODEL_PATH, + num_gpus=num_gpus, + tp_size=tp_size, + scheduler_port=scheduler_port, + ) + set_global_server_args(server_args) + + logger.info( + "Launching %s: CUDA_VISIBLE_DEVICES=%s num_gpus=%d tp=%d port=%d", + stage_name, cuda_devices, num_gpus, tp_size, server_args.scheduler_port, + ) + + processes = launch_partial_server( + server_args, + required_modules=required_modules, + custom_stages_fn=custom_stages_fn, + ) + + logger.info("%s ready (%d processes)", stage_name, len(processes)) + return processes, server_args + + +def terminate_processes(processes, name=""): + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) + logger.info("Terminated %s processes", name) + + +# ── Pipeline execution ────────────────────────────────────────────────── + +async def run_single_pipeline( + req_id: int, + encoder_client: StageClient, + denoiser_client: StageClient, + vae_client: StageClient, + seed: int, + save_output: bool = False, +) -> dict: + """Run one Encoder -> Denoiser -> VAE pipeline, return timing dict.""" + import torch + from sglang_utils import build_req, inject_tensors_to_req + + timings = {"req_id": req_id} + negative_prompt = "bad quality" if GUIDANCE > 1.0 else "" + req_kwargs = dict( + prompt=PROMPT, negative_prompt=negative_prompt, + height=HEIGHT, width=WIDTH, num_frames=NUM_FRAMES, + num_inference_steps=NUM_STEPS, guidance_scale=GUIDANCE, seed=seed, + ) + t_pipeline = time.monotonic() + + # ── Encoder ────────────────────────────────────────────────────── + t0 = time.monotonic() + enc_output = await encoder_client.forward([build_req(**req_kwargs)]) + timings["encoder_s"] = time.monotonic() - t0 + if enc_output.error: + raise RuntimeError(f"Encoder error: {enc_output.error}") + enc_tensors = enc_output.output + logger.info( + "req %d | Encoder done %.2fs — keys: %s", req_id, timings["encoder_s"], + {k: (list(v.shape) if hasattr(v, "shape") else f"list[{len(v)}]") + for k, v in enc_tensors.items()}, + ) + + # ── Denoiser ───────────────────────────────────────────────────── + t0 = time.monotonic() + den_req = build_req(**req_kwargs) + inject_tensors_to_req(den_req, enc_tensors) + den_req.do_classifier_free_guidance = (GUIDANCE > 1.0) + den_output = await denoiser_client.forward([den_req]) + timings["denoiser_s"] = time.monotonic() - t0 + if den_output.error: + raise RuntimeError(f"Denoiser error: {den_output.error}") + latents = den_output.output["latents"] + logger.info( + "req %d | Denoiser done %.2fs — latents: %s", + req_id, timings["denoiser_s"], list(latents.shape), + ) + + # ── VAE ────────────────────────────────────────────────────────── + t0 = time.monotonic() + vae_req = build_req(prompt="", height=HEIGHT, width=WIDTH, + num_frames=NUM_FRAMES, num_inference_steps=NUM_STEPS, + guidance_scale=0.0, seed=seed) + vae_req.latents = latents.cpu() + vae_output = await vae_client.forward([vae_req]) + timings["vae_s"] = time.monotonic() - t0 + if vae_output.error: + raise RuntimeError(f"VAE error: {vae_output.error}") + + timings["total_s"] = time.monotonic() - t_pipeline + logger.info( + "req %d | VAE done %.2fs — total pipeline: %.2fs", + req_id, timings["vae_s"], timings["total_s"], + ) + + # Save output as mp4 + if save_output and vae_output.output is not None: + _save_video(vae_output.output, req_id) + + return timings + + +def _save_video(frames_tensor, req_id: int): + """Save decoded video tensor [B,C,T,H,W] as mp4.""" + try: + import numpy as np + import imageio + + os.makedirs(OUTPUT_DIR, exist_ok=True) + if hasattr(frames_tensor, "cpu"): + frames_tensor = frames_tensor.cpu().float().numpy() + # [B, C, T, H, W] -> [T, H, W, C] + frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) + out_path = os.path.join(OUTPUT_DIR, f"output_{req_id}.mp4") + imageio.mimwrite(out_path, frames, fps=24, codec="libx264") + logger.info("req %d | Saved %d frames to %s", req_id, frames.shape[0], out_path) + except Exception as e: + logger.warning("req %d | Could not save video: %s", req_id, e) + + +def print_timing_report(all_timings: list, wall_elapsed: float): + """Print per-stage timing statistics.""" + stages = ["encoder_s", "denoiser_s", "vae_s", "total_s"] + n = len(all_timings) + + logger.info("") + logger.info("=" * 72) + logger.info(" Timing Report (%d requests, concurrency=%d)", n, CONCURRENCY) + logger.info("=" * 72) + + for t in all_timings: + logger.info( + " req %2d | enc=%6.2fs den=%6.2fs vae=%6.2fs total=%6.2fs", + t["req_id"], t["encoder_s"], t["denoiser_s"], t["vae_s"], t["total_s"], + ) + + if n > 1: + logger.info("-" * 72) + for stage in stages: + vals = [t[stage] for t in all_timings] + mean = statistics.mean(vals) + med = statistics.median(vals) + mn, mx = min(vals), max(vals) + std = statistics.stdev(vals) if n >= 2 else 0.0 + logger.info( + " %-10s mean=%6.2fs median=%6.2fs min=%6.2fs max=%6.2fs std=%5.2fs", + stage, mean, med, mn, mx, std, + ) + + logger.info("-" * 72) + throughput = n / wall_elapsed if wall_elapsed > 0 else 0 + logger.info(" Wall time: %.2fs | Throughput: %.2f req/s", wall_elapsed, throughput) + logger.info("=" * 72) + + +# ── Main ──────────────────────────────────────────────────────────────── + +async def main(): + from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages + + # Apply patches once before any SGLang config is created + _patch_hunyuan_config_task_type() + + logger.info("=" * 72) + logger.info(" Disaggregated Diffusion E2E — SGLang Backend") + logger.info(" Model: %s", MODEL_PATH) + logger.info(" Prompt: %s", PROMPT) + logger.info(" Encoder: GPU %s", GPU_ENC) + logger.info(" Denoiser: GPU %s (TP=%d)", GPU_DEN, TP_SIZE) + logger.info(" VAE: GPU %s", GPU_VAE) + logger.info(" Requests: %d Concurrency: %d", NUM_REQUESTS, CONCURRENCY) + logger.info(" Frames: %d Steps: %d Size: %dx%d Guidance: %.1f", + NUM_FRAMES, NUM_STEPS, WIDTH, HEIGHT, GUIDANCE) + logger.info("=" * 72) + + enc_procs = den_procs = vae_procs = None + enc_client = den_client = vae_client = None + + try: + # ── Launch all 3 stages ────────────────────────────────────── + t_launch = time.monotonic() + + enc_procs, enc_args = _launch_stage( + "Encoder", GPU_ENC, + required_modules=_detect_encoder_modules(MODEL_PATH), + custom_stages_fn=build_encoder_stages, + tp_size=1, + scheduler_port=15600, + ) + + den_procs, den_args = _launch_stage( + "Denoiser", GPU_DEN, + required_modules=["transformer", "scheduler"], + custom_stages_fn=build_denoiser_stages, + tp_size=TP_SIZE, + scheduler_port=15700, + ) + + vae_procs, vae_args = _launch_stage( + "VAE", GPU_VAE, + required_modules=["vae", "scheduler"], + custom_stages_fn=build_vae_stages, + tp_size=1, + scheduler_port=15800, + ) + + logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) + + # ── Connect clients ────────────────────────────────────────── + enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") + den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") + vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") + + # ── Warmup ─────────────────────────────────────────────────── + logger.info("Warmup request …") + warmup = await run_single_pipeline( + -1, enc_client, den_client, vae_client, SEED, save_output=False, + ) + logger.info( + "Warmup done — enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", + warmup["encoder_s"], warmup["denoiser_s"], + warmup["vae_s"], warmup["total_s"], + ) + + # ── Run pipeline(s) ────────────────────────────────────────── + if NUM_REQUESTS <= 1: + t_wall = time.monotonic() + timings = await run_single_pipeline( + 0, enc_client, den_client, vae_client, SEED, save_output=True, + ) + wall_elapsed = time.monotonic() - t_wall + print_timing_report([timings], wall_elapsed) + else: + logger.info("Firing %d requests (concurrency=%d) …", NUM_REQUESTS, CONCURRENCY) + sem = asyncio.Semaphore(CONCURRENCY) + + async def _run_one(i): + async with sem: + return await run_single_pipeline( + i, enc_client, den_client, vae_client, + SEED + i, save_output=(i == 0), + ) + + t_wall = time.monotonic() + tasks = [asyncio.create_task(_run_one(i)) for i in range(NUM_REQUESTS)] + all_timings = list(await asyncio.gather(*tasks)) + wall_elapsed = time.monotonic() - t_wall + print_timing_report(all_timings, wall_elapsed) + + finally: + for client in [enc_client, den_client, vae_client]: + if client is not None: + try: + client.close() + except Exception: + pass + for procs, name in [ + (enc_procs, "encoder"), (den_procs, "denoiser"), (vae_procs, "vae"), + ]: + if procs is not None: + terminate_processes(procs, name) + + +if __name__ == "__main__": + mp.set_start_method("spawn", force=True) + asyncio.run(main()) diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py new file mode 100644 index 000000000000..099710172b94 --- /dev/null +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SGLang PipelineStage utilities for disaggregated diffusion workers. + +Provides helpers to construct ServerArgs, load partial pipelines (only the +modules each worker needs), and convert between Dynamo protocol types and +SGLang's Req dataclass. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Dict, List, Optional + +import torch + +logger = logging.getLogger(__name__) + + +def build_server_args(model_path: str, **overrides): + """Create a ServerArgs, initialize torch.distributed, and set the global singleton.""" + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, + set_global_server_args, + ) + + defaults = dict( + model_path=model_path, + num_gpus=1, + ) + defaults.update(overrides) + server_args = ServerArgs.from_kwargs(**defaults) + set_global_server_args(server_args) + + _ensure_distributed_init(server_args) + + return server_args + + +def _ensure_distributed_init(server_args): + """Initialize torch.distributed and model-parallel groups via SGLang.""" + import inspect + from sglang.multimodal_gen.runtime.distributed import ( + model_parallel_is_initialized, + maybe_init_distributed_environment_and_model_parallel, + ) + + if model_parallel_is_initialized(): + return + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", str(server_args.master_port)) + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + + kwargs = dict( + tp_size=server_args.tp_size, + enable_cfg_parallel=server_args.enable_cfg_parallel, + ulysses_degree=server_args.ulysses_degree, + ring_degree=server_args.ring_degree, + sp_size=server_args.sp_degree, + dp_size=server_args.dp_size, + distributed_init_method=f"tcp://127.0.0.1:{server_args.master_port}", + ) + sig = inspect.signature(maybe_init_distributed_environment_and_model_parallel) + if "dist_timeout" in sig.parameters: + kwargs["dist_timeout"] = server_args.dist_timeout + + maybe_init_distributed_environment_and_model_parallel(**kwargs) + + +def build_partial_pipeline( + server_args, + required_modules: List[str], +): + """Load a pipeline with only *required_modules* populated. + + Auto-detects pipeline class from model_index.json, suppresses automatic + stage creation, and syncs all component configs (even unloaded ones). + """ + from sglang.multimodal_gen.runtime.pipelines import get_model_info + + model_info = get_model_info(server_args.model_path) + base_pipeline_cls = model_info.pipeline_cls + + # Build a partial pipeline class that: + # 1. Suppresses automatic stage creation + # 2. Skips LoRA init (which assumes 'transformer' always exists) + def _noop_create_stages(self, server_args): + return None + + def _safe_init(self, **kwargs): + # Call ComposedPipelineBase.__init__ directly, skipping LoRAPipeline + # which tries to access self.modules['transformer'] + from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + ComposedPipelineBase, + ) + ComposedPipelineBase.__init__(self, **kwargs) + + partial_cls = type( + f"_Partial{base_pipeline_cls.__name__}", + (base_pipeline_cls,), + { + "create_pipeline_stages": _noop_create_stages, + "__init__": _safe_init, + }, + ) + + pipeline = partial_cls( + model_path=server_args.model_path, + server_args=server_args, + required_config_modules=required_modules, + ) + + _sync_all_component_configs(server_args, pipeline) + + return pipeline + + +def _sync_all_component_configs(server_args, pipeline): + """Read config.json for every component in model_index.json and update + the corresponding arch_config in ``server_args.pipeline_config``, ensuring + correct parameters even for components whose weights are not loaded. + """ + import json + + CONFIG_ATTR_MAP = { + "vae": ("vae_config", "update_model_arch"), + "video_vae": ("vae_config", "update_model_arch"), + "transformer": ("dit_config", "update_model_arch"), + "video_dit": ("dit_config", "update_model_arch"), + "audio_dit": ("audio_dit_config", "update_model_arch"), + "audio_vae": ("audio_vae_config", "update_model_arch"), + } + + # pipeline.model_path is resolved to a local path by _load_config() + # (may differ from server_args.model_path if the original was a hub ID). + model_path = pipeline.model_path + model_index_path = os.path.join(model_path, "model_index.json") + if not os.path.isfile(model_index_path): + return + + with open(model_index_path, "r") as f: + model_index = json.load(f) + + pipeline_config = server_args.pipeline_config + for component_name, mapping in CONFIG_ATTR_MAP.items(): + config_attr, update_method_name = mapping + cfg = getattr(pipeline_config, config_attr, None) + if cfg is None: + continue + + if component_name not in model_index: + continue + + config_json_path = os.path.join(model_path, component_name, "config.json") + if not os.path.isfile(config_json_path): + continue + + with open(config_json_path, "r") as f: + hf_config = json.load(f) + + hf_config.pop("_class_name", None) + hf_config.pop("_diffusers_version", None) + + update_fn = getattr(cfg, update_method_name, None) + if update_fn is not None: + update_fn(hf_config) + logger.info("Synced %s config from %s (e.g. z_dim=%s)", + config_attr, + config_json_path, + getattr(getattr(cfg, "arch_config", cfg), "z_dim", "N/A")) + + if hasattr(cfg, "post_init"): + cfg.post_init() + + +def get_component_backend(module) -> str: + """Return a human-readable string indicating which backend loaded *module*.""" + mod = type(module).__module__ or "" + cls = type(module).__qualname__ + if mod.startswith("sglang."): + return f"sglang-optimized ({cls})" + if mod.startswith("diffusers."): + return f"native-diffusers ({cls})" + if mod.startswith("transformers."): + return f"native-transformers ({cls})" + return f"unknown ({mod}.{cls})" + + +def build_config(server_args): + """Construct a minimal Config from a diffusion ServerArgs.""" + import types + + try: + from dynamo.sglang.args import Config, DynamoConfig + dynamo_args = DynamoConfig.__new__(DynamoConfig) + except ImportError: + dynamo_args = types.SimpleNamespace() + Config = None + + dynamo_args.component = "disagg_diffusion" + dynamo_args.namespace = "disagg_diffusion" + dynamo_args.diffusion_worker = True + dynamo_args.use_kv_events = False + dynamo_args.media_output_fs_url = "file:///tmp/disagg_videos" + dynamo_args.media_output_http_url = None + dynamo_args.use_sglang_tokenizer = False + dynamo_args.multimodal_processor = False + dynamo_args.multimodal_encode_worker = False + dynamo_args.multimodal_worker = False + dynamo_args.embedding_worker = False + dynamo_args.image_diffusion_worker = False + dynamo_args.video_generation_worker = False + dynamo_args.disagg_config = None + dynamo_args.disagg_config_key = None + dynamo_args.endpoint = "generate" + dynamo_args.discovery_backend = "etcd" + dynamo_args.request_plane = "tcp" + dynamo_args.event_plane = "tcp" + dynamo_args.connector = [] + dynamo_args.enable_local_indexer = False + dynamo_args.durable_kv_events = False + dynamo_args.endpoint_types = "generate" + dynamo_args.dump_config_to = None + dynamo_args.multimodal_embedding_cache_capacity_gb = 0.0 + dynamo_args.output_modalities = ["video"] + dynamo_args.dyn_tool_call_parser = None + dynamo_args.dyn_reasoning_parser = None + dynamo_args.custom_jinja_template = None + + if not hasattr(server_args, "disaggregation_mode"): + server_args.disaggregation_mode = "null" + + if Config is not None: + return Config(server_args, dynamo_args) + + cfg = types.SimpleNamespace() + cfg.server_args = server_args + cfg.dynamo_args = dynamo_args + cfg.serving_mode = getattr(server_args, "disaggregation_mode", "null") + return cfg + + +def build_req( + prompt: str, + negative_prompt: Optional[str] = "", + height: int = 480, + width: int = 832, + num_frames: int = 17, + num_inference_steps: int = 20, + guidance_scale: float = 5.0, + seed: int = 42, + device: str = "cuda", + **extra_fields, +) -> "Req": + """Construct a minimal SGLang ``Req`` for running pipeline stages.""" + from sglang.multimodal_gen.configs.sample.base import DataType + from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req + + req = Req( + data_type=DataType.VIDEO, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + generator=torch.Generator(device="cpu").manual_seed(seed), + do_classifier_free_guidance=(guidance_scale > 1.0), + ) + + for k, v in extra_fields.items(): + setattr(req, k, v) + + return req + + +def extract_tensors_from_req( + req, + keys: List[str], +) -> Dict[str, object]: + """Pull named tensor fields out of a ``Req`` for NIXL transfer. + + Single-element lists are unwrapped to bare tensors. + Multi-element lists (dual-encoder outputs) are preserved as lists. + """ + result: Dict[str, object] = {} + for key in keys: + val = getattr(req, key, None) + if val is None: + continue + if isinstance(val, list): + if len(val) == 0: + continue + if len(val) == 1: + result[key] = val[0] + else: + # Keep multi-element lists intact (e.g. dual-encoder outputs + # with incompatible shapes) + result[key] = val + elif isinstance(val, torch.Tensor): + result[key] = val + return result + + +def inject_tensors_to_req( + req, + tensors: Dict[str, object], + list_fields: Optional[List[str]] = None, +): + """Inject received tensors back into a ``Req``. + + Values may be bare tensors (single-encoder outputs) or lists of tensors + (dual-encoder outputs like HunyuanVideo's Llama + CLIP embeddings). + ``list_fields`` specifies Req attributes that expect a list value; + bare tensors are auto-wrapped in a single-element list. + """ + list_fields = set(list_fields or [ + "prompt_embeds", + "negative_prompt_embeds", + "pooled_embeds", + "neg_pooled_embeds", + "image_embeds", + ]) + for key, value in tensors.items(): + if key in list_fields: + if isinstance(value, list): + # Multi-encoder: already a list of tensors + setattr(req, key, value) + else: + # Single encoder: wrap in list + setattr(req, key, [value]) + else: + setattr(req, key, value) + return req From 189bbbbfaa220850e59bf7c18687580e04f02c46 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 17:53:42 +0800 Subject: [PATCH 02/22] docs: update disagg diffusion design doc and README for HunyuanVideo - Rewrite design doc with high-level architecture diagrams and pipeline parallelism illustration - Add benchmark results (native vs disagg, HunyuanVideo 544x960) - Update milestone tracker: mark SGLang integration as completed - Add scaling roadmap (NIXL transfer, independent stage scaling, encoder caching, sequence parallelism, heterogeneous hardware) - Update README with quick-start commands and env var reference Co-Authored-By: Claude Opus 4.6 --- docs/design/disaggregated_diffusion.md | 313 +++++++++++++------------ examples/disagg_diffusion/README.md | 97 ++++---- 2 files changed, 215 insertions(+), 195 deletions(-) diff --git a/docs/design/disaggregated_diffusion.md b/docs/design/disaggregated_diffusion.md index 8703a9e0d9d0..a8dd4edc9638 100644 --- a/docs/design/disaggregated_diffusion.md +++ b/docs/design/disaggregated_diffusion.md @@ -1,203 +1,210 @@ -# Design Doc: Disaggregated Diffusion Inference (Diff-Disagg) in Dynamo +# Design Doc: Disaggregated Diffusion Inference in Dynamo ## 1. Motivation -Current diffusion inference in Dynamo (e.g., SGLang backend) is monolithic: a single worker loads all model components (Text Encoder, Transformer/UNet, VAE). This design faces several challenges with modern large-scale diffusion models (Flux, SD3, Video Models): +Modern video diffusion models (HunyuanVideo 13B, Wan2.2-14B, etc.) comprise +heterogeneous components with vastly different compute profiles: -1. **Huge Text Encoders**: Models like Flux use T5-XXL, consuming 10-20GB VRAM just for the encoder, which is only used once at the beginning. -2. **Resource Inefficiency**: During the long denoising loop (20-50 steps), the text encoder weights occupy precious H100/H200 memory without being used. -3. **Heterogeneous Compute**: - * **Text Encoding**: Compute-intensive, single pass. Can run on lower-end GPUs or even CPU. - * **Denoising**: Memory-bandwidth and compute-intensive, iterative. Requires high-end GPUs (H100/H200). - * **VAE Decoding**: Memory-intensive, single pass. Can be offloaded. +| Component | Params | Compute Pattern | VRAM (bf16) | +|-----------|--------|----------------|-------------| +| Text Encoder (e.g. Llama3-8B) | 8B | Single forward pass | ~16 GB | +| DiT Denoiser | 5-13B | 30-50 iterative steps | 10-26 GB | +| 3D VAE Decoder | ~200M | Single forward pass | ~2 GB | -**Goal**: Decompose the diffusion pipeline into flexible, independent stages (Encoder, Denoiser, VAE) that can be deployed on different hardware and scaled independently. +In a monolithic deployment, **all components occupy a single GPU** throughout +the entire request lifetime, even though the encoder is idle during denoising +(96%+ of wall time) and the VAE is idle during encoding+denoising. -## 2. Architecture: Router-Orchestrated Multi-Stage Pipeline +**Disaggregated diffusion** decomposes the pipeline into independently +deployable stages, enabling: -We adopt a **Router-Centric** architecture similar to Dynamo's existing EPD (Encoder-Prefill-Decode) flow for LLMs. The Frontend/Router acts as the orchestrator, chaining calls between specialized workers. +- **Independent scaling** per stage (e.g. 1 encoder : 4 denoisers : 1 VAE) +- **Heterogeneous hardware** (encoder on cost-efficient GPUs, denoiser on high-end) +- **Pipeline parallelism** across concurrent requests +- **Memory efficiency** (each GPU loads only its stage's weights) -### 2.1 Component Roles +## 2. Architecture -1. **Frontend / Global Router**: - * Acts as the conductor. - * Receives the initial user request. - * Orchestrates the sequence: `Encoder -> Denoiser -> VAE`. - * Maintains the request state and context. +### 2.1 Three-Stage Pipeline -2. **Stage Workers**: - * **Encoder Worker**: Loads CLIP/T5. Input: Text. Output: Embeddings (Tensor Bytes). - * **Denoiser Worker**: Loads Transformer/UNet. Input: Embeddings + Latents (optional). Output: Denoised Latents. - * **VAE Worker**: Loads VAE. Input: Denoised Latents. Output: Image/Video (Pixel Data). +``` + ZMQ ZMQ ZMQ + Client ──► [ Encoder ] ─── embeds ──► [ Denoiser ] ─── latents ──► [ VAE ] ──► Video + GPU 0 GPU 1,2 (TP=2) GPU 3 + Llama3-8B HunyuanVideo DiT 3D Causal VAE + + CLIP 13B params ~200M params +``` -### 2.2 Data Flow (Request/Response) +Each stage runs as a separate SGLang `Scheduler` subprocess with its own CUDA +context. Inter-stage communication uses ZMQ (pickle serialization). + +### 2.2 Request Flow ```mermaid sequenceDiagram participant Client - participant Frontend - participant Router - participant EncoderWorker - participant DenoiserWorker - participant VAEWorker - - Client->>Frontend: POST /v1/images/generations (Prompt) - - Note over Frontend: Step 1: Text Encoding - Frontend->>Router: Route(ModelType.Encoder, Prompt) - Router->>EncoderWorker: GenerateEmbeddings(Prompt) - EncoderWorker-->>Frontend: Return {prompt_embeds, pooled_embeds} - - Note over Frontend: Step 2: Denoising - Frontend->>Router: Route(ModelType.Denoiser, Embeddings) - Router->>DenoiserWorker: Denoise(Embeddings, Latents=None) - DenoiserWorker-->>Frontend: Return {denoised_latents} - - Note over Frontend: Step 3: VAE Decoding - Frontend->>Router: Route(ModelType.VAE, Latents) - Router->>VAEWorker: Decode(Latents) - VAEWorker-->>Frontend: Return {image_data / url} - - Frontend-->>Client: Final Response -``` + participant Orchestrator + participant Encoder as Encoder (GPU 0) + participant Denoiser as Denoiser (GPU 1,2) + participant VAE as VAE (GPU 3) -*Optimization Note*: For large data (like Video Latents), we can use **Reference Passing** via a shared storage (Object Store / Shared Memory) instead of passing raw bytes through the Router/Frontend. + Client->>Orchestrator: generate(prompt, params) -## 3. Detailed Design + Orchestrator->>Encoder: Req(prompt) + Encoder-->>Orchestrator: {prompt_embeds: [Llama, CLIP]} -### 3.1 Protocol Extensions (`dynamo/common/protocols`) + Orchestrator->>Denoiser: Req(prompt_embeds, params) + Note over Denoiser: 50 denoising steps (TP=2) + Denoiser-->>Orchestrator: {latents: [B,C,T,H,W]} -We need to define data structures for intermediate results. + Orchestrator->>VAE: Req(latents) + VAE-->>Orchestrator: decoded video [B,C,T,H,W] -**New Protocol Types:** - -```python -class DiffusionEmbeddingData(BaseModel): - """Output from Encoder Stage""" - prompt_embeds: bytes # Serialized Tensor (e.g., safetensors/numpy) - pooled_prompt_embeds: Optional[bytes] = None - negative_prompt_embeds: Optional[bytes] = None - negative_pooled_prompt_embeds: Optional[bytes] = None - -class DiffusionLatentData(BaseModel): - """Output from Denoiser Stage""" - latents: bytes # Serialized Tensor - shape: List[int] - dtype: str - -class StageRequest(BaseModel): - """Generic Request for a specific stage""" - stage: str # "encoder", "denoiser", "vae" - input_data: Union[str, DiffusionEmbeddingData, DiffusionLatentData] - params: Dict[str, Any] # Generation params (steps, cfg, etc.) + Orchestrator-->>Client: video.mp4 ``` -### 3.2 ModelType Expansion (`dynamo/llm/src/model.rs` & Python Enums) +### 2.3 Concurrent Request Pipelining -Extend `ModelType` to support fine-grained stages. +With disaggregation, stages process different requests simultaneously: -```python -class ModelType(IntFlag): - # Existing - Tokens = auto() - # ... - - # New Diffusion Stages - DiffusionEncoder = auto() # Text Encoder only - DiffusionDenoiser = auto() # Transformer/UNet only - DiffusionVAE = auto() # VAE only ``` +Time ──────────────────────────────────────────────────────► -### 3.3 SGLang Backend Adaptation (`components/src/dynamo/sglang`) - -We need to modify `init_diffusion.py` and handlers to support partial loading. - -**Configuration:** -Add `--diffusion-stage` argument to `sglang` worker. - -* `--diffusion-stage full` (Default): Loads everything (current behavior). -* `--diffusion-stage encoder`: Loads only Text Encoders. -* `--diffusion-stage denoiser`: Loads only Transformer, accepts Embeddings. -* `--diffusion-stage vae`: Loads only VAE, accepts Latents. +Encoder: |req0 enc|req1 enc| |req2 enc|req3 enc| +Denoiser: | req0 denoise (50 steps) | req1 denoise ... +VAE: |req0 vae| +``` -**Handler Implementation:** +The encoder and VAE are freed immediately after their stage completes, +allowing them to serve the next request while the denoiser is busy. -1. **`EncoderHandler`**: - * Uses `pipe.encode_prompt()`. - * Returns serialized embeddings. +## 3. Implementation -2. **`DenoiserHandler`**: - * Initializes pipeline with `text_encoder=None`, `vae=None`. - * Implements `generate(prompt_embeds=...)`. - * Returns latents (skips VAE decode). +### 3.1 Core Components -3. **`VAEHandler`**: - * Initializes pipeline with `text_encoder=None`, `transformer=None`. - * Implements `decode(latents=...)`. +| File | Role | +|------|------| +| `partial_gpu_worker.py` | `PartialGPUWorker` — loads a subset of pipeline modules; `IntermediateOutputStage` / `DeviceMoveStage` for inter-stage tensor transfer | +| `sglang_utils.py` | `build_partial_pipeline()` — suppresses default stage creation, syncs component configs; tensor injection/extraction helpers | +| `run_e2e_sglang.py` | Orchestrator — launches 3 stage schedulers, connects ZMQ clients, runs E2E pipeline with timing | -### 3.4 Router Logic (`components/src/dynamo/global_router`) +### 3.2 Partial Pipeline Loading -The Global Router needs to be aware of these new `ModelType`s. +Each stage loads only its required modules via `build_partial_pipeline()`: -* **Registration**: Workers register with specific types (e.g., `ModelType.DiffusionEncoder`). -* **Routing**: - * `handle_encoder_request`: Routes to Encoder Pool. - * `handle_denoiser_request`: Routes to Denoiser Pool. - * `handle_vae_request`: Routes to VAE Pool. +```python +# Encoder: loads text encoders + tokenizers only (~16 GB) +pipeline = build_partial_pipeline(server_args, + required_modules=["text_encoder", "text_encoder_2", + "tokenizer", "tokenizer_2", "scheduler"]) + +# Denoiser: loads transformer only (~26 GB, or ~13 GB/GPU with TP=2) +pipeline = build_partial_pipeline(server_args, + required_modules=["transformer", "scheduler"]) + +# VAE: loads VAE only (~2 GB) +pipeline = build_partial_pipeline(server_args, + required_modules=["vae", "scheduler"]) +``` -## 4. Implementation Plan +A dynamic subclass of the model's pipeline is created at runtime to: +1. Suppress automatic stage creation (`create_pipeline_stages` → no-op) +2. Skip LoRA init (which assumes `transformer` always exists) +3. Sync all component configs (VAE z_dim, DiT patch_size, etc.) even for unloaded modules -### POC (Proof of Concept) +### 3.3 Multi-Encoder Support -Full POC code lives in [`examples/disagg_diffusion/`](../../examples/disagg_diffusion/). +Models with dual text encoders (e.g. HunyuanVideo: Llama3-8B + CLIP) produce +embedding lists with incompatible shapes. The pipeline handles this transparently: -#### Phase 0: Offline Validation (`phase0_validate/`) +- `IntermediateOutputStage`: preserves multi-element lists as-is (no `torch.cat`) +- `inject_tensors_to_req()`: accepts both bare tensors and lists +- `DeviceMoveStage`: moves each tensor in the list to GPU with `.contiguous()` -Single-GPU script that proves diffusers supports split execution. -Runs Encoder → Denoiser → VAE as three separate stages with serialized -intermediate tensors, then compares the result against a monolithic run. +### 3.4 SGLang Compatibility Patches -**Key risk validated**: Can diffusers pipelines accept `prompt_embeds` and -return `output_type="latent"` to bypass text encoder / VAE respectively? +Two runtime patches are applied for HunyuanVideo compatibility: -#### Phase 1: Dynamo Stage Workers (`phase1_workers/`) +1. **`HunyuanConfig` task_type**: The upstream config class lacks a default + `task_type`, causing instantiation to fail. We wrap `__init__` to supply + `ModelTaskType.T2V` when omitted. -Three independent Dynamo workers, each loading only its model component: +2. **Triton norm contiguous**: HunyuanVideo's transformer produces + non-contiguous intermediate tensors from attention reshapes. SGLang's triton + layernorm kernel asserts `x.stride(-1) == 1`. We patch `norm_infer` to call + `.contiguous()` when needed. -| Worker | Loads | Endpoint | Input | Output | -|--------|-------|----------|-------|--------| -| `encoder_worker.py` | CLIP + T5 | `disagg_diffusion.encoder.encode` | text prompt | embeddings (b64) | -| `denoiser_worker.py` | Transformer | `disagg_diffusion.denoiser.denoise` | embeddings + params | latents (b64) | -| `vae_worker.py` | VAE | `disagg_diffusion.vae.decode` | latents | image (b64 PNG) | +## 4. Supported Models -Intermediate data is serialized as base64-encoded `torch.save` bytes. -Protocol types are defined in `protocol.py`. +| Model | Pipeline Class | Text Encoders | DiT | Status | +|-------|---------------|---------------|-----|--------| +| **HunyuanVideo v1** | `HunyuanVideoPipeline` | Llama3-8B + CLIP | 13B | **Validated** | +| Wan2.2-TI2V-5B | `WanPipeline` | T5 | 5B | Validated | +| Wan2.1-T2V-14B | `WanPipeline` | T5 | 14B | Untested | +| FLUX.1 (image) | `FluxPipeline` | CLIP + T5 | 12B | Phase 0 only | -#### Phase 2: Orchestrator Client (`phase2_orchestrator/`) +The encoder module detection is automatic via `model_index.json` — any SGLang-supported +diffusion model with the standard 3-stage structure should work without code changes. -A lightweight Python client that connects to the Dynamo runtime, calls the -three stage endpoints in sequence, and produces the final image. +## 5. Benchmarks (HunyuanVideo v1, 544x960, 50 steps, 61 frames) -### Production Roadmap +### Single Request Latency -#### Phase 3: Protocol & Types -1. Define `DiffusionEmbeddingData` and `DiffusionLatentData` in `dynamo/common/protocols`. -2. Add `DiffusionEncoder`, `DiffusionDenoiser`, `DiffusionVAE` to `ModelType` enum (Rust + Python). +| | Native SGLang (1 GPU) | Disagg (4 GPUs, TP=2) | +|---|---|---| +| Encoder | 0.31s | 0.35s | +| Denoiser | 512.19s (10.24s/step) | 477.14s (9.42s/step) | +| VAE | 24.53s | 18.32s | +| **Total** | **537.12s** | **495.81s** | -#### Phase 4: SGLang Modularization -1. Modify `init_diffusion.py` to accept `--diffusion-stage {full,encoder,denoiser,vae}`. -2. Implement `EncoderHandler`, `DenoiserHandler`, `VAEHandler` in `sglang/request_handlers/disagg_diffusion/`. -3. Add logic to conditionally load model components based on stage. +### 2 Concurrent Requests (9 frames, 3 steps — smoke test) -#### Phase 5: Frontend / Router Orchestration -1. Update `Frontend` to detect if the backend is disaggregated diffusion. -2. Implement the multi-step orchestration logic (`Encoder → Denoiser → VAE`) - in `Frontend` or a specialized `DiffusionOrchestrator` component. -3. Support configurable stage DAGs (e.g., skip VAE for latent-only output). +``` +req 0 | enc=0.34s den=2.92s vae=2.26s total=5.52s +req 1 | enc=0.66s den=5.50s vae=2.25s total=8.41s +Wall time: 8.44s | Throughput: 0.24 req/s +``` -#### Phase 6: Optimization -1. **Reference Passing**: Replace base64 tensor payloads with shared-memory - handles or object-store URLs (e.g., `media_output_fs_url`). -2. **Pipelining**: Overlap encoding of request N+1 with denoising of request N. -3. **Encoder Caching**: LRU cache for repeated prompts at the Encoder stage. -4. **Video Extension**: Extend to video models (larger latents, more frames). +Pipeline parallelism observed: req 1's encoder overlapped with req 0's denoiser. + +## 6. Milestone Tracker + +### Completed + +- [x] **Phase 0**: Offline split validation (FLUX.1-schnell, diffusers) +- [x] **Phase 1a**: Dynamo stage workers (encoder/denoiser/VAE as Dynamo services, FLUX) +- [x] **Phase 1b**: SGLang backend integration + - [x] `PartialGPUWorker` with monkey-patched `GPUWorker` + - [x] `build_partial_pipeline()` with auto-detected pipeline class + - [x] `IntermediateOutputStage` for ZMQ tensor transfer + - [x] TP=2 denoiser support (multi-process pipe wiring) + - [x] HunyuanVideo v1 dual-encoder support (Llama + CLIP) + - [x] E2E orchestrator with per-stage timing and mp4 output + - [x] Concurrent request support with pipeline parallelism +- [x] **Phase 2**: Orchestrator client (ZMQ-based, async) + +### In Progress + +- [ ] **Phase 3**: Dynamo Router integration + - [ ] Register stage workers as typed endpoints (`DiffusionEncoder`, `DiffusionDenoiser`, `DiffusionVAE`) + - [ ] Router-orchestrated multi-stage chaining (replace manual ZMQ orchestrator) + - [ ] Frontend API: `POST /v1/video/generations` + +### Scaling Roadmap + +- [ ] **NIXL tensor transfer**: Replace ZMQ pickle with RDMA/GPU-direct for + inter-stage latent transfer (critical for video — latents are O(100 MB)) +- [ ] **Independent stage scaling**: N:M:K ratio (e.g. 1 encoder : 4 denoisers : 1 VAE) + with load-balanced routing per stage pool +- [ ] **Encoder caching**: LRU cache for repeated prompts at the encoder stage + (common in batch generation / prompt engineering workflows) +- [ ] **Sequence parallelism**: Ulysses/Ring SP for denoiser to scale beyond TP + (especially for long videos with high temporal resolution) +- [ ] **VAE tiling**: Enable tiled 3D VAE decoding for high-resolution / long videos + that exceed single-GPU VRAM +- [ ] **Heterogeneous hardware**: Deploy encoder on cost-efficient GPUs (e.g. L4/T4), + denoiser on compute-optimized GPUs (H100/H200), VAE on memory-optimized nodes +- [ ] **HunyuanVideo v1.5 support**: Requires SGLang to implement + `HunyuanVideo15Pipeline` (Qwen2.5-VL + ByT5 encoders, 8.3B DiT) +- [ ] **Prefill-decode analogy**: Continuous batching within the denoiser stage + (batch multiple requests at different denoising steps) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 72a8793cc506..872fbbc9f77e 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,60 +1,73 @@ -# Disaggregated Diffusion Inference POC +# Disaggregated Diffusion Inference -Split a monolithic diffusion pipeline (Text Encoder → Transformer → VAE) into -independent stages that can run on separate GPUs and scale independently. +Split a video diffusion pipeline (Text Encoder -> DiT Denoiser -> 3D VAE) into +independent stages on separate GPUs for pipeline parallelism and independent scaling. Design doc: [docs/design/disaggregated_diffusion.md](../../docs/design/disaggregated_diffusion.md) -## Phases +## Quick Start (HunyuanVideo, 4 GPUs) -### Phase 0: Offline Validation (no Dynamo) +```bash +# Single request — generates a 2.5s video at 544x960 +python phase1_workers/run_e2e_sglang.py -Proves that diffusers supports split execution: encode, denoise, and VAE decode -can run independently with serialized intermediate tensors. +# 4 requests, 2 concurrent +NUM_REQUESTS=4 CONCURRENCY=2 python phase1_workers/run_e2e_sglang.py -```bash -# Single GPU, ~24 GB VRAM for FLUX.1-schnell -python phase0_validate/validate_split.py \ - --model black-forest-labs/FLUX.1-schnell \ - --prompt "A photo of a cat sitting on a windowsill" \ - --output-dir /tmp/disagg_validate +# Custom prompt and parameters +PROMPT="A rocket launching into space" NUM_FRAMES=33 NUM_STEPS=30 \ + python phase1_workers/run_e2e_sglang.py ``` -### Phase 1: Dynamo Stage Workers - -Three independent Dynamo workers, each loading only its model component: +Default GPU layout: Encoder (GPU 0) | Denoiser TP=2 (GPU 1,2) | VAE (GPU 3). -```bash -# Terminal 1: Encoder Worker (loads CLIP + T5, ~12 GB) -python phase1_workers/encoder_worker.py --model black-forest-labs/FLUX.1-schnell +Override with `GPU_ENC`, `GPU_DEN`, `GPU_VAE` environment variables. -# Terminal 2: Denoiser Worker (loads Transformer, ~24 GB) -python phase1_workers/denoiser_worker.py --model black-forest-labs/FLUX.1-schnell +## Architecture -# Terminal 3: VAE Worker (loads VAE, ~1 GB) -python phase1_workers/vae_worker.py --model black-forest-labs/FLUX.1-schnell ``` - -### Phase 2: Orchestrator - -Chains the three stage endpoints into an end-to-end generation pipeline: - -```bash -python phase2_orchestrator/run_disagg.py \ - --prompt "A photo of a cat sitting on a windowsill" \ - --output /tmp/disagg_output.png + Encoder (GPU 0) Denoiser (GPU 1,2) VAE (GPU 3) + ┌──────────────┐ ┌──────────────────┐ ┌─────────────┐ + │ Llama3-8B │ embeds │ HunyuanVideo DiT │ lats │ 3D Causal │ + │ + CLIP │───ZMQ──│ 13B (TP=2) │──ZMQ──│ VAE │──► video + │ ~16 GB │ │ ~13 GB/GPU │ │ ~2 GB │ + └──────────────┘ └──────────────────┘ └─────────────┘ ``` -Or use the all-in-one launch script: - -```bash -bash launch/run_all.sh black-forest-labs/FLUX.1-schnell "A photo of a cat" -``` +Each stage runs as a separate SGLang scheduler subprocess with its own CUDA context. ## Supported Models -Any diffusers pipeline that exposes `encode_prompt()` and supports -`prompt_embeds` / `output_type="latent"`. Tested with: - -- `black-forest-labs/FLUX.1-schnell` (recommended, 4 steps) -- `stabilityai/stable-diffusion-3.5-medium` +| Model | Encoders | DiT | Status | +|-------|----------|-----|--------| +| **HunyuanVideo v1** (`hunyuanvideo-community/HunyuanVideo`) | Llama3-8B + CLIP | 13B | Default | +| Wan2.2-TI2V-5B (`Wan-AI/Wan2.2-TI2V-5B-Diffusers`) | T5 | 5B | Supported | + +To use Wan: `MODEL_PATH=Wan-AI/Wan2.2-TI2V-5B-Diffusers GUIDANCE=5.0 python phase1_workers/run_e2e_sglang.py` + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | +| `PROMPT` | `A cat walking on green grass` | Text prompt | +| `GPU_ENC` | `0` | GPU(s) for encoder | +| `GPU_DEN` | `1,2` | GPU(s) for denoiser | +| `GPU_VAE` | `3` | GPU(s) for VAE | +| `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism degree | +| `NUM_FRAMES` | `61` | Video frames (~2.5s at 24fps) | +| `NUM_STEPS` | `50` | Denoising steps | +| `HEIGHT` | `544` | Frame height | +| `WIDTH` | `960` | Frame width | +| `GUIDANCE` | `1.0` | Guidance scale (>1.0 enables CFG) | +| `NUM_REQUESTS` | `1` | Number of pipeline runs | +| `CONCURRENCY` | `1` | Max concurrent pipelines | + +## Files + +| File | Purpose | +|------|---------| +| `phase1_workers/run_e2e_sglang.py` | Orchestrator: launches stages, runs E2E pipeline, reports timing | +| `phase1_workers/partial_gpu_worker.py` | `PartialGPUWorker`, `IntermediateOutputStage`, subprocess launcher | +| `phase1_workers/sglang_utils.py` | Partial pipeline loading, tensor injection/extraction, config sync | +| `phase0_validate/` | Offline split validation (FLUX.1-schnell, diffusers) | From df4296691d2c7d32e8887110cd14d3a7e6de763a Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 19:56:59 +0800 Subject: [PATCH 03/22] feat: add NIXL RDMA GPU-direct tensor transfer between disagg stages Replace ZMQ pickle tensor transfer with NIXL RDMA for inter-stage communication. Only small metadata (~1.5 KB) now travels over ZMQ; actual tensor data (embeddings, latents) transfers GPU-to-GPU without CPU round-trips. - New nixl_transfer.py: NixlTensorSender/Receiver with persistent Connector and async-to-sync bridge for PipelineStage.forward() - NixlSendStage: registers GPU tensors as NIXL-readable, returns metadata in OutputBatch; handles dual-encoder indexed fields - NixlReceiveStage: RDMA-pulls tensors from sender GPU, reconstructs indexed fields back into lists; falls back to CPU device-move when NIXL metadata is absent - Orchestrator routes NIXL metadata between stages instead of tensors - New serve.py: FastAPI HTTP server for curl-based requests - Update design doc: remove benchmarks, reflect NIXL architecture Co-Authored-By: Claude Opus 4.6 --- docs/design/disaggregated_diffusion.md | 160 ++++------- .../phase1_workers/nixl_transfer.py | 135 +++++++++ .../phase1_workers/partial_gpu_worker.py | 144 +++++++--- .../phase1_workers/run_e2e_sglang.py | 31 ++- .../disagg_diffusion/phase1_workers/serve.py | 262 ++++++++++++++++++ 5 files changed, 587 insertions(+), 145 deletions(-) create mode 100644 examples/disagg_diffusion/phase1_workers/nixl_transfer.py create mode 100644 examples/disagg_diffusion/phase1_workers/serve.py diff --git a/docs/design/disaggregated_diffusion.md b/docs/design/disaggregated_diffusion.md index a8dd4edc9638..a703f9759f15 100644 --- a/docs/design/disaggregated_diffusion.md +++ b/docs/design/disaggregated_diffusion.md @@ -28,15 +28,18 @@ deployable stages, enabling: ### 2.1 Three-Stage Pipeline ``` - ZMQ ZMQ ZMQ - Client ──► [ Encoder ] ─── embeds ──► [ Denoiser ] ─── latents ──► [ VAE ] ──► Video - GPU 0 GPU 1,2 (TP=2) GPU 3 - Llama3-8B HunyuanVideo DiT 3D Causal VAE - + CLIP 13B params ~200M params + NIXL RDMA NIXL RDMA + Client ──► [ Encoder ] ── embeds ──► [ Denoiser ] ── latents ──► [ VAE ] ──► Video + GPU 0 GPU 1,2 (TP=2) GPU 3 + Llama3-8B HunyuanVideo DiT 3D Causal VAE + + CLIP 13B params ~200M params ``` Each stage runs as a separate SGLang `Scheduler` subprocess with its own CUDA -context. Inter-stage communication uses ZMQ (pickle serialization). +context. The data plane uses **NIXL RDMA** for GPU-direct tensor transfer +between stages — only small metadata (~1.5 KB) travels over the ZMQ control +plane; actual tensor data (embeddings ~3.6 MB, latents ~1.4 MB) transfers +GPU-to-GPU without CPU round-trips. ### 2.2 Request Flow @@ -50,14 +53,16 @@ sequenceDiagram Client->>Orchestrator: generate(prompt, params) - Orchestrator->>Encoder: Req(prompt) - Encoder-->>Orchestrator: {prompt_embeds: [Llama, CLIP]} + Orchestrator->>Encoder: Req(prompt) [ZMQ] + Encoder-->>Orchestrator: {nixl_metadata} [ZMQ] - Orchestrator->>Denoiser: Req(prompt_embeds, params) + Orchestrator->>Denoiser: Req(nixl_metadata, params) [ZMQ] + Note over Denoiser: NIXL RDMA pull embeddings from Encoder GPU Note over Denoiser: 50 denoising steps (TP=2) - Denoiser-->>Orchestrator: {latents: [B,C,T,H,W]} + Denoiser-->>Orchestrator: {nixl_metadata} [ZMQ] - Orchestrator->>VAE: Req(latents) + Orchestrator->>VAE: Req(nixl_metadata) [ZMQ] + Note over VAE: NIXL RDMA pull latents from Denoiser GPU VAE-->>Orchestrator: decoded video [B,C,T,H,W] Orchestrator-->>Client: video.mp4 @@ -75,64 +80,51 @@ Denoiser: | req0 denoise (50 steps) | req1 denoise ... VAE: |req0 vae| ``` -The encoder and VAE are freed immediately after their stage completes, -allowing them to serve the next request while the denoiser is busy. - ## 3. Implementation ### 3.1 Core Components | File | Role | |------|------| -| `partial_gpu_worker.py` | `PartialGPUWorker` — loads a subset of pipeline modules; `IntermediateOutputStage` / `DeviceMoveStage` for inter-stage tensor transfer | -| `sglang_utils.py` | `build_partial_pipeline()` — suppresses default stage creation, syncs component configs; tensor injection/extraction helpers | -| `run_e2e_sglang.py` | Orchestrator — launches 3 stage schedulers, connects ZMQ clients, runs E2E pipeline with timing | +| `nixl_transfer.py` | `NixlTensorSender` / `NixlTensorReceiver` — GPU-direct RDMA transfer via NIXL | +| `partial_gpu_worker.py` | `PartialGPUWorker` — loads subset of pipeline modules; `NixlSendStage` / `NixlReceiveStage` for inter-stage transfer | +| `sglang_utils.py` | `build_partial_pipeline()` — suppresses default stage creation, syncs component configs | +| `run_e2e_sglang.py` | Orchestrator — launches 3 stage schedulers, connects ZMQ clients, routes NIXL metadata | +| `serve.py` | HTTP serving mode — FastAPI wrapper for curl-based requests | -### 3.2 Partial Pipeline Loading +### 3.2 NIXL Tensor Transfer -Each stage loads only its required modules via `build_partial_pipeline()`: +Inter-stage tensor data flows GPU-to-GPU via NIXL RDMA: -```python -# Encoder: loads text encoders + tokenizers only (~16 GB) -pipeline = build_partial_pipeline(server_args, - required_modules=["text_encoder", "text_encoder_2", - "tokenizer", "tokenizer_2", "scheduler"]) - -# Denoiser: loads transformer only (~26 GB, or ~13 GB/GPU with TP=2) -pipeline = build_partial_pipeline(server_args, - required_modules=["transformer", "scheduler"]) - -# VAE: loads VAE only (~2 GB) -pipeline = build_partial_pipeline(server_args, - required_modules=["vae", "scheduler"]) -``` +1. **Sender** (`NixlSendStage`): flattens output tensors into a contiguous buffer, + registers it as NIXL-readable, returns metadata (shapes, dtypes, NIXL descriptor) +2. **Orchestrator**: forwards only the small metadata dict via ZMQ (~1.5 KB) +3. **Receiver** (`NixlReceiveStage`): allocates a GPU buffer, uses NIXL to RDMA-pull + the tensor data directly from the sender's GPU, reconstructs individual tensors -A dynamic subclass of the model's pipeline is created at runtime to: -1. Suppress automatic stage creation (`create_pipeline_stages` → no-op) -2. Skip LoRA init (which assumes `transformer` always exists) -3. Sync all component configs (VAE z_dim, DiT patch_size, etc.) even for unloaded modules +Falls back to ZMQ pickle transfer when NIXL is unavailable. -### 3.3 Multi-Encoder Support +### 3.3 Partial Pipeline Loading -Models with dual text encoders (e.g. HunyuanVideo: Llama3-8B + CLIP) produce -embedding lists with incompatible shapes. The pipeline handles this transparently: +Each stage loads only its required modules via `build_partial_pipeline()`: -- `IntermediateOutputStage`: preserves multi-element lists as-is (no `torch.cat`) -- `inject_tensors_to_req()`: accepts both bare tensors and lists -- `DeviceMoveStage`: moves each tensor in the list to GPU with `.contiguous()` +```python +# Encoder: text encoders + tokenizers only (~16 GB) +required_modules=["text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler"] -### 3.4 SGLang Compatibility Patches +# Denoiser: transformer only (~26 GB, or ~13 GB/GPU with TP=2) +required_modules=["transformer", "scheduler"] -Two runtime patches are applied for HunyuanVideo compatibility: +# VAE: VAE only (~2 GB) +required_modules=["vae", "scheduler"] +``` -1. **`HunyuanConfig` task_type**: The upstream config class lacks a default - `task_type`, causing instantiation to fail. We wrap `__init__` to supply - `ModelTaskType.T2V` when omitted. +### 3.4 Multi-Encoder Support -2. **Triton norm contiguous**: HunyuanVideo's transformer produces - non-contiguous intermediate tensors from attention reshapes. SGLang's triton - layernorm kernel asserts `x.stride(-1) == 1`. We patch `norm_infer` to call - `.contiguous()` when needed. +Models with dual text encoders (e.g. HunyuanVideo: Llama3-8B + CLIP) produce +embedding lists with incompatible shapes. The NIXL transfer handles this by +indexing each element separately (`prompt_embeds_0`, `prompt_embeds_1`) during +send, and reconstructing the list on receive. ## 4. Supported Models @@ -141,70 +133,32 @@ Two runtime patches are applied for HunyuanVideo compatibility: | **HunyuanVideo v1** | `HunyuanVideoPipeline` | Llama3-8B + CLIP | 13B | **Validated** | | Wan2.2-TI2V-5B | `WanPipeline` | T5 | 5B | Validated | | Wan2.1-T2V-14B | `WanPipeline` | T5 | 14B | Untested | -| FLUX.1 (image) | `FluxPipeline` | CLIP + T5 | 12B | Phase 0 only | The encoder module detection is automatic via `model_index.json` — any SGLang-supported diffusion model with the standard 3-stage structure should work without code changes. -## 5. Benchmarks (HunyuanVideo v1, 544x960, 50 steps, 61 frames) - -### Single Request Latency - -| | Native SGLang (1 GPU) | Disagg (4 GPUs, TP=2) | -|---|---|---| -| Encoder | 0.31s | 0.35s | -| Denoiser | 512.19s (10.24s/step) | 477.14s (9.42s/step) | -| VAE | 24.53s | 18.32s | -| **Total** | **537.12s** | **495.81s** | - -### 2 Concurrent Requests (9 frames, 3 steps — smoke test) - -``` -req 0 | enc=0.34s den=2.92s vae=2.26s total=5.52s -req 1 | enc=0.66s den=5.50s vae=2.25s total=8.41s -Wall time: 8.44s | Throughput: 0.24 req/s -``` - -Pipeline parallelism observed: req 1's encoder overlapped with req 0's denoiser. - -## 6. Milestone Tracker +## 5. Milestone Tracker ### Completed - [x] **Phase 0**: Offline split validation (FLUX.1-schnell, diffusers) -- [x] **Phase 1a**: Dynamo stage workers (encoder/denoiser/VAE as Dynamo services, FLUX) -- [x] **Phase 1b**: SGLang backend integration - - [x] `PartialGPUWorker` with monkey-patched `GPUWorker` - - [x] `build_partial_pipeline()` with auto-detected pipeline class - - [x] `IntermediateOutputStage` for ZMQ tensor transfer - - [x] TP=2 denoiser support (multi-process pipe wiring) - - [x] HunyuanVideo v1 dual-encoder support (Llama + CLIP) - - [x] E2E orchestrator with per-stage timing and mp4 output - - [x] Concurrent request support with pipeline parallelism -- [x] **Phase 2**: Orchestrator client (ZMQ-based, async) - -### In Progress - -- [ ] **Phase 3**: Dynamo Router integration - - [ ] Register stage workers as typed endpoints (`DiffusionEncoder`, `DiffusionDenoiser`, `DiffusionVAE`) - - [ ] Router-orchestrated multi-stage chaining (replace manual ZMQ orchestrator) - - [ ] Frontend API: `POST /v1/video/generations` +- [x] **Phase 1a**: Dynamo stage workers (encoder/denoiser/VAE as Dynamo services) +- [x] **Phase 1b**: SGLang backend integration with `PartialGPUWorker` +- [x] **NIXL RDMA transfer**: GPU-direct tensor transfer between stages +- [x] **HunyuanVideo v1**: Dual-encoder (Llama + CLIP), 13B DiT, TP=2 +- [x] **HTTP serving**: FastAPI-based `/generate` endpoint +- [x] **Concurrent requests**: Pipeline parallelism with async orchestrator ### Scaling Roadmap -- [ ] **NIXL tensor transfer**: Replace ZMQ pickle with RDMA/GPU-direct for - inter-stage latent transfer (critical for video — latents are O(100 MB)) +- [ ] **Dynamo Router integration**: Register stage workers as typed endpoints, + replace manual orchestrator with Router-orchestrated multi-stage chaining - [ ] **Independent stage scaling**: N:M:K ratio (e.g. 1 encoder : 4 denoisers : 1 VAE) with load-balanced routing per stage pool - [ ] **Encoder caching**: LRU cache for repeated prompts at the encoder stage - (common in batch generation / prompt engineering workflows) - [ ] **Sequence parallelism**: Ulysses/Ring SP for denoiser to scale beyond TP - (especially for long videos with high temporal resolution) - [ ] **VAE tiling**: Enable tiled 3D VAE decoding for high-resolution / long videos - that exceed single-GPU VRAM -- [ ] **Heterogeneous hardware**: Deploy encoder on cost-efficient GPUs (e.g. L4/T4), - denoiser on compute-optimized GPUs (H100/H200), VAE on memory-optimized nodes -- [ ] **HunyuanVideo v1.5 support**: Requires SGLang to implement - `HunyuanVideo15Pipeline` (Qwen2.5-VL + ByT5 encoders, 8.3B DiT) -- [ ] **Prefill-decode analogy**: Continuous batching within the denoiser stage - (batch multiple requests at different denoising steps) +- [ ] **Heterogeneous hardware**: Deploy encoder on cost-efficient GPUs (L4/T4), + denoiser on compute-optimized GPUs (H100/H200) +- [ ] **Continuous batching**: Batch multiple requests at different denoising steps + within the denoiser stage (prefill-decode analogy) diff --git a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py new file mode 100644 index 000000000000..234606ba2f47 --- /dev/null +++ b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NIXL RDMA tensor transfer for disaggregated diffusion stages. + +Provides GPU-direct tensor transfer between stage workers. Only small +metadata (shapes, dtypes, NIXL descriptor ~1.5 KB) travels over the ZMQ +control plane; actual tensor data (embeddings, latents) transfers +GPU→GPU via NIXL RDMA. + +Usage inside PipelineStage.forward() (synchronous context):: + + sender = NixlTensorSender() + meta = sender.send({"latents": tensor}) # registers & returns metadata + # ... pass meta via ZMQ ... + + receiver = NixlTensorReceiver() + tensors = receiver.recv(meta, device="cuda") # RDMA pull +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any, Dict, List, Optional + +import torch + +logger = logging.getLogger(__name__) + +try: + import dynamo.nixl_connect as nixl_connect + NIXL_AVAILABLE = True +except ImportError: + NIXL_AVAILABLE = False + logger.info("NIXL not available — falling back to ZMQ tensor transfer") + + +class _PersistentConnector: + """Lazily-initialized NIXL Connector singleton per process.""" + + _instance: Optional[nixl_connect.Connector] = None if NIXL_AVAILABLE else None + + @classmethod + async def get(cls) -> nixl_connect.Connector: + if cls._instance is None: + cls._instance = nixl_connect.Connector() + await cls._instance.initialize() + return cls._instance + + +class NixlTensorSender: + """Register GPU tensors as NIXL-readable. Returns metadata for the receiver. + + The readable is kept alive via a background task so the sender process + can return immediately after yielding metadata. + """ + + def __init__(self): + self._pending: list = [] + + def send(self, tensors: Dict[str, torch.Tensor]) -> dict: + """Register tensors and return metadata dict (synchronous wrapper).""" + return asyncio.run(self._async_send(tensors)) + + async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: + # Clean completed tasks + self._pending = [t for t in self._pending if not t.done()] + + connector = await _PersistentConnector.get() + + # Flatten all tensors into a single contiguous buffer + flat = torch.cat([t.contiguous().view(-1) for t in tensors.values()]) + descriptor = nixl_connect.Descriptor(flat) + readable = await connector.create_readable(descriptor) + raw_meta = readable.metadata() + + meta = { + "tensor_keys": list(tensors.keys()), + "shapes": {k: list(t.shape) for k, t in tensors.items()}, + "dtypes": {k: str(t.dtype).removeprefix("torch.") for k, t in tensors.items()}, + "nixl_metadata": raw_meta.model_dump() if hasattr(raw_meta, "model_dump") else raw_meta, + } + + # Keep readable alive until the receiver has pulled the data + async def _keep_alive(): + try: + await readable.wait_for_completion() + except Exception as e: + logger.warning("NIXL readable wait failed: %s", e) + + task = asyncio.ensure_future(_keep_alive()) + self._pending.append(task) + return meta + + +class NixlTensorReceiver: + """Pull tensors from a remote sender via NIXL RDMA.""" + + def recv(self, meta: dict, device: str = "cuda") -> Dict[str, torch.Tensor]: + """Pull tensors described by metadata. Returns {name: tensor}.""" + return asyncio.run(self._async_recv(meta, device)) + + async def _async_recv(self, meta: dict, device: str) -> Dict[str, torch.Tensor]: + connector = await _PersistentConnector.get() + + # Calculate total size and per-tensor specs + specs = [] + total_bytes = 0 + for key in meta["tensor_keys"]: + shape = meta["shapes"][key] + dtype = getattr(torch, meta["dtypes"][key]) + numel = 1 + for s in shape: + numel *= s + size = numel * dtype.itemsize + specs.append((key, shape, dtype, size)) + total_bytes += size + + # Allocate receive buffer directly on target device (GPU-direct) + flat = torch.empty(total_bytes, dtype=torch.uint8, device=device) + descriptor = nixl_connect.Descriptor(flat) + + rdma_meta = nixl_connect.RdmaMetadata.model_validate(meta["nixl_metadata"]) + read_op = await connector.begin_read(rdma_meta, descriptor) + await read_op.wait_for_completion() + + # Slice the flat buffer into individual tensors + result = {} + offset = 0 + for key, shape, dtype, size in specs: + result[key] = flat[offset:offset + size].view(dtype=dtype).reshape(shape) + offset += size + + return result diff --git a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py index 95e8550121a8..546dbdf69f11 100644 --- a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py @@ -72,19 +72,52 @@ # ═══════════════════════════════════════════════════════════════════════ -class DeviceMoveStage(PipelineStage): - """Move tensor fields on ``Req`` to the current CUDA device. +class NixlReceiveStage(PipelineStage): + """Pull tensor fields onto the current GPU via NIXL RDMA. - Tensors arriving via ZMQ (pickle) land on CPU. Prepend this stage - at the start of denoiser / VAE pipelines that receive tensors from - other processes. + Prepend at the start of denoiser / VAE pipelines. The ``Req`` + carries NIXL metadata (set by the orchestrator); this stage uses + it to pull the actual tensor data directly from the sender's GPU + without a CPU round-trip. """ def __init__(self, tensor_fields: List[str]): super().__init__() self._tensor_fields = tensor_fields + self._receiver = None def forward(self, batch: Req, server_args: ServerArgs) -> Req: + nixl_meta = getattr(batch, "_nixl_transfer_meta", None) + if nixl_meta is not None: + return self._nixl_pull(batch, nixl_meta) + # Fallback: tensors arrived via ZMQ pickle — just move to GPU + return self._device_move(batch) + + def _nixl_pull(self, batch: Req, meta: dict) -> Req: + from nixl_transfer import NixlTensorReceiver + if self._receiver is None: + self._receiver = NixlTensorReceiver() + tensors = self._receiver.recv(meta, device="cuda") + # Reconstruct indexed fields (e.g. prompt_embeds_0, prompt_embeds_1 + # + __prompt_embeds_count → prompt_embeds list) + reconstructed = {} + indexed = {} # base_name → {idx: tensor} + for k, v in tensors.items(): + if k.startswith("__") and k.endswith("_count"): + continue + parts = k.rsplit("_", 1) + if len(parts) == 2 and parts[1].isdigit(): + indexed.setdefault(parts[0], {})[int(parts[1])] = v + else: + reconstructed[k] = v + for base, idx_map in indexed.items(): + reconstructed[base] = [idx_map[i] for i in sorted(idx_map)] + from sglang_utils import inject_tensors_to_req + inject_tensors_to_req(batch, reconstructed) + return batch + + def _device_move(self, batch: Req) -> Req: + """Fallback: move CPU tensors to GPU (ZMQ pickle path).""" device = torch.device("cuda") for field in self._tensor_fields: val = getattr(batch, field, None) @@ -100,26 +133,36 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: return batch -class IntermediateOutputStage(PipelineStage): - """Pipeline stage that packages ``Req`` tensors into ``OutputBatch``. +class NixlSendStage(PipelineStage): + """Register ``Req`` tensors as NIXL-readable and return metadata. Append as the **last** stage in encoder / denoiser partial pipelines. - ``GPUWorker.execute_forward`` sees an ``OutputBatch`` (not a ``Req``) - and returns it through ZMQ without any custom override. - - ``output_fields`` lists the ``Req`` attributes to extract. - Single-element lists are unwrapped to bare tensors. - Multi-element lists (e.g. dual-encoder embeddings) are preserved as - lists so the receiver can reconstruct them correctly. - The handler receives ``OutputBatch.output`` as ``dict[str, Tensor|list]``. + Returns an ``OutputBatch`` whose ``output`` dict contains: + - ``_nixl_transfer_meta``: NIXL metadata for the receiver (~1.5 KB) + - tensor shapes/dtypes for logging + + Actual tensor data stays on GPU — only metadata travels over ZMQ. + Falls back to sending raw tensors when NIXL is unavailable. """ def __init__(self, output_fields: List[str]): super().__init__() self._output_fields = output_fields + self._sender = None def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: - tensors: Dict[str, object] = {} + tensors = self._extract_tensors(batch) + if not tensors: + return OutputBatch(output={}) + + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE: + return self._nixl_send(tensors) + return self._fallback_send(tensors) + + def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: + """Flatten list-valued fields into individual tensors.""" + result: Dict[str, torch.Tensor] = {} for field in self._output_fields: val = getattr(batch, field, None) if val is None: @@ -128,14 +171,51 @@ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: if len(val) == 0: continue if len(val) == 1: - tensors[field] = val[0] + result[field] = val[0] else: - # Preserve multi-element lists (e.g. dual-encoder outputs - # with incompatible shapes) so they survive ZMQ pickle. - tensors[field] = val + # Dual-encoder: store each element separately for NIXL + for i, t in enumerate(val): + result[f"{field}_{i}"] = t + result[f"__{field}_count"] = torch.tensor(len(val)) elif isinstance(val, torch.Tensor): - tensors[field] = val - return OutputBatch(output=tensors) + result[field] = val + return result + + def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + from nixl_transfer import NixlTensorSender + if self._sender is None: + self._sender = NixlTensorSender() + # Filter out non-GPU metadata tensors for NIXL + gpu_tensors = {k: v for k, v in tensors.items() + if isinstance(v, torch.Tensor) and v.is_cuda} + cpu_tensors = {k: v for k, v in tensors.items() + if isinstance(v, torch.Tensor) and not v.is_cuda} + meta = self._sender.send(gpu_tensors) + # Include CPU metadata tensors directly (e.g. __count fields) + meta["cpu_tensors"] = cpu_tensors + return OutputBatch(output={"_nixl_transfer_meta": meta}) + + def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + """Fallback: send raw tensors via ZMQ pickle.""" + # Reconstruct list-valued fields for backward compat + output: Dict[str, object] = {} + counts = {} + for k, v in tensors.items(): + if k.startswith("__") and k.endswith("_count"): + base = k[2:-6] + counts[base] = int(v.item()) + elif "_" in k and k.rsplit("_", 1)[1].isdigit(): + base, idx = k.rsplit("_", 1) + output.setdefault(f"_list_{base}", {})[int(idx)] = v + else: + output[k] = v + # Reassemble lists + for base, idx_map in list(output.items()): + if base.startswith("_list_"): + real_key = base[6:] + output[real_key] = [idx_map[i] for i in sorted(idx_map)] + del output[base] + return OutputBatch(output=output) # ═══════════════════════════════════════════════════════════════════════ @@ -172,11 +252,11 @@ def build_encoder_stages(pipeline, server_args): f"Encoder/tokenizer count mismatch: {len(text_encoders)} vs {len(tokenizers)}" ) - # Always list both fields; IntermediateOutputStage skips None values, + # Always list both fields; NixlSendStage skips None values, # so negative_prompt_embeds is simply omitted when CFG is disabled. return [ TextEncodingStage(text_encoders=text_encoders, tokenizers=tokenizers), - IntermediateOutputStage(["prompt_embeds", "negative_prompt_embeds"]), + NixlSendStage(["prompt_embeds", "negative_prompt_embeds"]), ] @@ -197,22 +277,17 @@ def build_denoiser_stages(pipeline, server_args): scheduler = pipeline.get_module("scheduler") logger.info("transformer backend: %s", get_component_backend(transformer)) - # Fields that arrive from the encoder via ZMQ and need GPU placement. - # prompt_embeds may be a single tensor (Wan) or list of tensors - # (HunyuanVideo dual-encoder) — DeviceMoveStage handles both. - device_move_fields = ["prompt_embeds", "negative_prompt_embeds"] - return [ - DeviceMoveStage(device_move_fields), + NixlReceiveStage(["prompt_embeds", "negative_prompt_embeds"]), LatentPreparationStage(scheduler=scheduler, transformer=transformer), TimestepPreparationStage(scheduler=scheduler), DenoisingStage(transformer=transformer, scheduler=scheduler), - IntermediateOutputStage(["latents"]), + NixlSendStage(["latents"]), ] def build_vae_stages(pipeline, server_args): - """DecodingStage (already returns OutputBatch — no extra stage needed).""" + """NixlReceive → DecodingStage.""" from sglang.multimodal_gen.runtime.pipelines.stages.decoding import ( DecodingStage, ) @@ -220,7 +295,10 @@ def build_vae_stages(pipeline, server_args): vae = pipeline.get_module("vae") logger.info("vae backend: %s", get_component_backend(vae)) - return [DecodingStage(vae=vae, pipeline=pipeline)] + return [ + NixlReceiveStage(["latents"]), + DecodingStage(vae=vae, pipeline=pipeline), + ] # ═══════════════════════════════════════════════════════════════════════ diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index 7278e5baab54..92d047333a7d 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -235,26 +235,36 @@ async def run_single_pipeline( timings["encoder_s"] = time.monotonic() - t0 if enc_output.error: raise RuntimeError(f"Encoder error: {enc_output.error}") - enc_tensors = enc_output.output + enc_result = enc_output.output + # enc_result is either {"_nixl_transfer_meta": {...}} (NIXL) or + # {"prompt_embeds": tensor, ...} (ZMQ fallback) + nixl_mode = "_nixl_transfer_meta" in enc_result logger.info( - "req %d | Encoder done %.2fs — keys: %s", req_id, timings["encoder_s"], - {k: (list(v.shape) if hasattr(v, "shape") else f"list[{len(v)}]") - for k, v in enc_tensors.items()}, + "req %d | Encoder done %.2fs — transfer: %s", + req_id, timings["encoder_s"], + "NIXL RDMA" if nixl_mode else f"ZMQ (keys: {list(enc_result.keys())})", ) # ── Denoiser ───────────────────────────────────────────────────── t0 = time.monotonic() den_req = build_req(**req_kwargs) - inject_tensors_to_req(den_req, enc_tensors) den_req.do_classifier_free_guidance = (GUIDANCE > 1.0) + if nixl_mode: + # Pass NIXL metadata — NixlReceiveStage will RDMA-pull the tensors + den_req._nixl_transfer_meta = enc_result["_nixl_transfer_meta"] + else: + # ZMQ fallback — tensors already in enc_result + inject_tensors_to_req(den_req, enc_result) den_output = await denoiser_client.forward([den_req]) timings["denoiser_s"] = time.monotonic() - t0 if den_output.error: raise RuntimeError(f"Denoiser error: {den_output.error}") - latents = den_output.output["latents"] + den_result = den_output.output + nixl_mode_den = "_nixl_transfer_meta" in den_result logger.info( - "req %d | Denoiser done %.2fs — latents: %s", - req_id, timings["denoiser_s"], list(latents.shape), + "req %d | Denoiser done %.2fs — transfer: %s", + req_id, timings["denoiser_s"], + "NIXL RDMA" if nixl_mode_den else "ZMQ", ) # ── VAE ────────────────────────────────────────────────────────── @@ -262,7 +272,10 @@ async def run_single_pipeline( vae_req = build_req(prompt="", height=HEIGHT, width=WIDTH, num_frames=NUM_FRAMES, num_inference_steps=NUM_STEPS, guidance_scale=0.0, seed=seed) - vae_req.latents = latents.cpu() + if nixl_mode_den: + vae_req._nixl_transfer_meta = den_result["_nixl_transfer_meta"] + else: + vae_req.latents = den_result["latents"].cpu() vae_output = await vae_client.forward([vae_req]) timings["vae_s"] = time.monotonic() - t0 if vae_output.error: diff --git a/examples/disagg_diffusion/phase1_workers/serve.py b/examples/disagg_diffusion/phase1_workers/serve.py new file mode 100644 index 000000000000..7181855503b8 --- /dev/null +++ b/examples/disagg_diffusion/phase1_workers/serve.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""HTTP serving mode for disaggregated diffusion pipeline. + +Launches Encoder, Denoiser (TP=2), and VAE stages, then serves +requests via a FastAPI HTTP server. + +Usage: + python serve.py + + # Then send requests: + curl -X POST http://localhost:8090/generate \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \ + --output output.mp4 + + # Check server health: + curl http://localhost:8090/health +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import logging +import multiprocessing as mp +import os +import sys +import time +from typing import Optional + +# Ensure workers dir on path +WORKERS_DIR = os.path.dirname(os.path.abspath(__file__)) +if WORKERS_DIR not in sys.path: + sys.path.insert(0, WORKERS_DIR) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", +) +logger = logging.getLogger("serve") + +# ── Configuration ──────────────────────────────────────────────────────── +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +GPU_ENC = os.environ.get("GPU_ENC", "0") +GPU_DEN = os.environ.get("GPU_DEN", "1,2") +GPU_VAE = os.environ.get("GPU_VAE", "3") +TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(","))))) +SERVE_HOST = os.environ.get("SERVE_HOST", "0.0.0.0") +SERVE_PORT = int(os.environ.get("SERVE_PORT", "8090")) +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_e2e") + + +# ── Reuse run_e2e components ───────────────────────────────────────────── +from run_e2e_sglang import ( + _patch_hunyuan_config_task_type, + _detect_encoder_modules, + _launch_stage, + StageClient, + terminate_processes, +) + +# Global state +enc_client: Optional[StageClient] = None +den_client: Optional[StageClient] = None +vae_client: Optional[StageClient] = None +_all_procs = [] +_request_counter = 0 + + +async def run_pipeline( + prompt: str, + negative_prompt: str = "", + height: int = 544, + width: int = 960, + num_frames: int = 61, + num_steps: int = 50, + guidance_scale: float = 1.0, + seed: int = 42, +) -> bytes: + """Run the full E2E pipeline and return mp4 bytes.""" + global _request_counter + import torch + from sglang_utils import build_req, inject_tensors_to_req + + req_id = _request_counter + _request_counter += 1 + + req_kwargs = dict( + prompt=prompt, negative_prompt=negative_prompt, + height=height, width=width, num_frames=num_frames, + num_inference_steps=num_steps, guidance_scale=guidance_scale, seed=seed, + ) + + t_start = time.monotonic() + + # Encoder + t0 = time.monotonic() + enc_output = await enc_client.forward([build_req(**req_kwargs)]) + enc_time = time.monotonic() - t0 + if enc_output.error: + raise RuntimeError(f"Encoder error: {enc_output.error}") + logger.info("req %d | Encoder %.2fs", req_id, enc_time) + + # Denoiser + t0 = time.monotonic() + den_req = build_req(**req_kwargs) + inject_tensors_to_req(den_req, enc_output.output) + den_req.do_classifier_free_guidance = (guidance_scale > 1.0) + den_output = await den_client.forward([den_req]) + den_time = time.monotonic() - t0 + if den_output.error: + raise RuntimeError(f"Denoiser error: {den_output.error}") + logger.info("req %d | Denoiser %.2fs", req_id, den_time) + + # VAE + t0 = time.monotonic() + vae_req = build_req(prompt="", height=height, width=width, + num_frames=num_frames, num_inference_steps=num_steps, + guidance_scale=0.0, seed=seed) + vae_req.latents = den_output.output["latents"].cpu() + vae_output = await vae_client.forward([vae_req]) + vae_time = time.monotonic() - t0 + if vae_output.error: + raise RuntimeError(f"VAE error: {vae_output.error}") + + total = time.monotonic() - t_start + logger.info("req %d | VAE %.2fs | Total %.2fs", req_id, vae_time, total) + + # Convert to mp4 bytes + import numpy as np + import imageio + + frames_tensor = vae_output.output + if hasattr(frames_tensor, "cpu"): + frames_tensor = frames_tensor.cpu().float().numpy() + frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) + + buf = io.BytesIO() + imageio.mimwrite(buf, frames, format="mp4", fps=24, codec="libx264") + mp4_bytes = buf.getvalue() + + logger.info( + "req %d | Done: %d frames, %dx%d, %.1f KB, enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", + req_id, frames.shape[0], width, height, len(mp4_bytes) / 1024, + enc_time, den_time, vae_time, total, + ) + return mp4_bytes + + +# ── FastAPI app ────────────────────────────────────────────────────────── + +from fastapi import FastAPI, HTTPException +from fastapi.responses import Response +from pydantic import BaseModel + +app = FastAPI(title="Disaggregated Diffusion Server") + + +class GenerateRequest(BaseModel): + prompt: str = "A cat walking on green grass" + negative_prompt: str = "" + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_steps: int = 50 + guidance_scale: float = 1.0 + seed: int = 42 + + +@app.get("/health") +async def health(): + return {"status": "ok", "model": MODEL_PATH} + + +@app.post("/generate") +async def generate(req: GenerateRequest): + try: + mp4_bytes = await run_pipeline( + prompt=req.prompt, + negative_prompt=req.negative_prompt, + height=req.height, + width=req.width, + num_frames=req.num_frames, + num_steps=req.num_steps, + guidance_scale=req.guidance_scale, + seed=req.seed, + ) + return Response(content=mp4_bytes, media_type="video/mp4") + except Exception as e: + logger.exception("Generation failed") + raise HTTPException(status_code=500, detail=str(e)) + + +# ── Startup / Shutdown ─────────────────────────────────────────────────── + +def launch_all_stages(): + """Launch all 3 stage schedulers and return clients.""" + global enc_client, den_client, vae_client, _all_procs + + from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages + + _patch_hunyuan_config_task_type() + + logger.info("Launching stages: Encoder=%s Denoiser=%s (TP=%d) VAE=%s", + GPU_ENC, GPU_DEN, TP_SIZE, GPU_VAE) + + enc_procs, enc_args = _launch_stage( + "Encoder", GPU_ENC, + required_modules=_detect_encoder_modules(MODEL_PATH), + custom_stages_fn=build_encoder_stages, + tp_size=1, scheduler_port=15600, + ) + den_procs, den_args = _launch_stage( + "Denoiser", GPU_DEN, + required_modules=["transformer", "scheduler"], + custom_stages_fn=build_denoiser_stages, + tp_size=TP_SIZE, scheduler_port=15700, + ) + vae_procs, vae_args = _launch_stage( + "VAE", GPU_VAE, + required_modules=["vae", "scheduler"], + custom_stages_fn=build_vae_stages, + tp_size=1, scheduler_port=15800, + ) + + _all_procs.extend(enc_procs + den_procs + vae_procs) + + enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") + den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") + vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") + + logger.info("All stages ready. Serving at http://%s:%d", SERVE_HOST, SERVE_PORT) + + +def shutdown_stages(): + for client in [enc_client, den_client, vae_client]: + if client: + try: + client.close() + except Exception: + pass + for p in _all_procs: + p.terminate() + for p in _all_procs: + p.join(timeout=10) + logger.info("All stages terminated.") + + +if __name__ == "__main__": + mp.set_start_method("spawn", force=True) + + launch_all_stages() + + import uvicorn + try: + uvicorn.run(app, host=SERVE_HOST, port=SERVE_PORT, log_level="info") + finally: + shutdown_stages() From 350f26084d704d7b73119c88d51cab8c9704dfae Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 20:07:49 +0800 Subject: [PATCH 04/22] docs: consolidate design into README, improve serve.py clarity - Rewrite README with full design, architecture diagram, data flow, implementation details, quick start, API reference, and roadmap - Remove separate design doc (all content now in README) - Improve serve.py: clear process architecture in docstring, detailed health endpoint showing GPU layout and worker count - Verified: 3 sequential HTTP requests via curl all succeed with NIXL RDMA transfer on 4 GPUs Co-Authored-By: Claude Opus 4.6 --- docs/design/disaggregated_diffusion.md | 164 ---------------- examples/disagg_diffusion/README.md | 181 +++++++++++++----- .../disagg_diffusion/phase1_workers/serve.py | 158 ++++++++++----- 3 files changed, 244 insertions(+), 259 deletions(-) delete mode 100644 docs/design/disaggregated_diffusion.md diff --git a/docs/design/disaggregated_diffusion.md b/docs/design/disaggregated_diffusion.md deleted file mode 100644 index a703f9759f15..000000000000 --- a/docs/design/disaggregated_diffusion.md +++ /dev/null @@ -1,164 +0,0 @@ -# Design Doc: Disaggregated Diffusion Inference in Dynamo - -## 1. Motivation - -Modern video diffusion models (HunyuanVideo 13B, Wan2.2-14B, etc.) comprise -heterogeneous components with vastly different compute profiles: - -| Component | Params | Compute Pattern | VRAM (bf16) | -|-----------|--------|----------------|-------------| -| Text Encoder (e.g. Llama3-8B) | 8B | Single forward pass | ~16 GB | -| DiT Denoiser | 5-13B | 30-50 iterative steps | 10-26 GB | -| 3D VAE Decoder | ~200M | Single forward pass | ~2 GB | - -In a monolithic deployment, **all components occupy a single GPU** throughout -the entire request lifetime, even though the encoder is idle during denoising -(96%+ of wall time) and the VAE is idle during encoding+denoising. - -**Disaggregated diffusion** decomposes the pipeline into independently -deployable stages, enabling: - -- **Independent scaling** per stage (e.g. 1 encoder : 4 denoisers : 1 VAE) -- **Heterogeneous hardware** (encoder on cost-efficient GPUs, denoiser on high-end) -- **Pipeline parallelism** across concurrent requests -- **Memory efficiency** (each GPU loads only its stage's weights) - -## 2. Architecture - -### 2.1 Three-Stage Pipeline - -``` - NIXL RDMA NIXL RDMA - Client ──► [ Encoder ] ── embeds ──► [ Denoiser ] ── latents ──► [ VAE ] ──► Video - GPU 0 GPU 1,2 (TP=2) GPU 3 - Llama3-8B HunyuanVideo DiT 3D Causal VAE - + CLIP 13B params ~200M params -``` - -Each stage runs as a separate SGLang `Scheduler` subprocess with its own CUDA -context. The data plane uses **NIXL RDMA** for GPU-direct tensor transfer -between stages — only small metadata (~1.5 KB) travels over the ZMQ control -plane; actual tensor data (embeddings ~3.6 MB, latents ~1.4 MB) transfers -GPU-to-GPU without CPU round-trips. - -### 2.2 Request Flow - -```mermaid -sequenceDiagram - participant Client - participant Orchestrator - participant Encoder as Encoder (GPU 0) - participant Denoiser as Denoiser (GPU 1,2) - participant VAE as VAE (GPU 3) - - Client->>Orchestrator: generate(prompt, params) - - Orchestrator->>Encoder: Req(prompt) [ZMQ] - Encoder-->>Orchestrator: {nixl_metadata} [ZMQ] - - Orchestrator->>Denoiser: Req(nixl_metadata, params) [ZMQ] - Note over Denoiser: NIXL RDMA pull embeddings from Encoder GPU - Note over Denoiser: 50 denoising steps (TP=2) - Denoiser-->>Orchestrator: {nixl_metadata} [ZMQ] - - Orchestrator->>VAE: Req(nixl_metadata) [ZMQ] - Note over VAE: NIXL RDMA pull latents from Denoiser GPU - VAE-->>Orchestrator: decoded video [B,C,T,H,W] - - Orchestrator-->>Client: video.mp4 -``` - -### 2.3 Concurrent Request Pipelining - -With disaggregation, stages process different requests simultaneously: - -``` -Time ──────────────────────────────────────────────────────► - -Encoder: |req0 enc|req1 enc| |req2 enc|req3 enc| -Denoiser: | req0 denoise (50 steps) | req1 denoise ... -VAE: |req0 vae| -``` - -## 3. Implementation - -### 3.1 Core Components - -| File | Role | -|------|------| -| `nixl_transfer.py` | `NixlTensorSender` / `NixlTensorReceiver` — GPU-direct RDMA transfer via NIXL | -| `partial_gpu_worker.py` | `PartialGPUWorker` — loads subset of pipeline modules; `NixlSendStage` / `NixlReceiveStage` for inter-stage transfer | -| `sglang_utils.py` | `build_partial_pipeline()` — suppresses default stage creation, syncs component configs | -| `run_e2e_sglang.py` | Orchestrator — launches 3 stage schedulers, connects ZMQ clients, routes NIXL metadata | -| `serve.py` | HTTP serving mode — FastAPI wrapper for curl-based requests | - -### 3.2 NIXL Tensor Transfer - -Inter-stage tensor data flows GPU-to-GPU via NIXL RDMA: - -1. **Sender** (`NixlSendStage`): flattens output tensors into a contiguous buffer, - registers it as NIXL-readable, returns metadata (shapes, dtypes, NIXL descriptor) -2. **Orchestrator**: forwards only the small metadata dict via ZMQ (~1.5 KB) -3. **Receiver** (`NixlReceiveStage`): allocates a GPU buffer, uses NIXL to RDMA-pull - the tensor data directly from the sender's GPU, reconstructs individual tensors - -Falls back to ZMQ pickle transfer when NIXL is unavailable. - -### 3.3 Partial Pipeline Loading - -Each stage loads only its required modules via `build_partial_pipeline()`: - -```python -# Encoder: text encoders + tokenizers only (~16 GB) -required_modules=["text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler"] - -# Denoiser: transformer only (~26 GB, or ~13 GB/GPU with TP=2) -required_modules=["transformer", "scheduler"] - -# VAE: VAE only (~2 GB) -required_modules=["vae", "scheduler"] -``` - -### 3.4 Multi-Encoder Support - -Models with dual text encoders (e.g. HunyuanVideo: Llama3-8B + CLIP) produce -embedding lists with incompatible shapes. The NIXL transfer handles this by -indexing each element separately (`prompt_embeds_0`, `prompt_embeds_1`) during -send, and reconstructing the list on receive. - -## 4. Supported Models - -| Model | Pipeline Class | Text Encoders | DiT | Status | -|-------|---------------|---------------|-----|--------| -| **HunyuanVideo v1** | `HunyuanVideoPipeline` | Llama3-8B + CLIP | 13B | **Validated** | -| Wan2.2-TI2V-5B | `WanPipeline` | T5 | 5B | Validated | -| Wan2.1-T2V-14B | `WanPipeline` | T5 | 14B | Untested | - -The encoder module detection is automatic via `model_index.json` — any SGLang-supported -diffusion model with the standard 3-stage structure should work without code changes. - -## 5. Milestone Tracker - -### Completed - -- [x] **Phase 0**: Offline split validation (FLUX.1-schnell, diffusers) -- [x] **Phase 1a**: Dynamo stage workers (encoder/denoiser/VAE as Dynamo services) -- [x] **Phase 1b**: SGLang backend integration with `PartialGPUWorker` -- [x] **NIXL RDMA transfer**: GPU-direct tensor transfer between stages -- [x] **HunyuanVideo v1**: Dual-encoder (Llama + CLIP), 13B DiT, TP=2 -- [x] **HTTP serving**: FastAPI-based `/generate` endpoint -- [x] **Concurrent requests**: Pipeline parallelism with async orchestrator - -### Scaling Roadmap - -- [ ] **Dynamo Router integration**: Register stage workers as typed endpoints, - replace manual orchestrator with Router-orchestrated multi-stage chaining -- [ ] **Independent stage scaling**: N:M:K ratio (e.g. 1 encoder : 4 denoisers : 1 VAE) - with load-balanced routing per stage pool -- [ ] **Encoder caching**: LRU cache for repeated prompts at the encoder stage -- [ ] **Sequence parallelism**: Ulysses/Ring SP for denoiser to scale beyond TP -- [ ] **VAE tiling**: Enable tiled 3D VAE decoding for high-resolution / long videos -- [ ] **Heterogeneous hardware**: Deploy encoder on cost-efficient GPUs (L4/T4), - denoiser on compute-optimized GPUs (H100/H200) -- [ ] **Continuous batching**: Batch multiple requests at different denoising steps - within the denoiser stage (prefill-decode analogy) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 872fbbc9f77e..30529c886b5b 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,73 +1,160 @@ # Disaggregated Diffusion Inference -Split a video diffusion pipeline (Text Encoder -> DiT Denoiser -> 3D VAE) into -independent stages on separate GPUs for pipeline parallelism and independent scaling. +Decomposes a monolithic video diffusion pipeline into three independently +deployable stages — **Text Encoder**, **DiT Denoiser**, and **3D VAE** — +each running on separate GPUs. Inter-stage tensor data transfers use +**NIXL RDMA** (GPU-direct); only small metadata (~1.5 KB) travels over +the ZMQ control plane. -Design doc: [docs/design/disaggregated_diffusion.md](../../docs/design/disaggregated_diffusion.md) +## 1. Design -## Quick Start (HunyuanVideo, 4 GPUs) +### 1.1 Motivation -```bash -# Single request — generates a 2.5s video at 544x960 -python phase1_workers/run_e2e_sglang.py +In monolithic diffusion inference, all model components occupy a single GPU +for the entire request lifetime. The encoder (~16 GB) sits idle during the +denoising loop (96%+ of wall time), and the VAE (~2 GB) sits idle during +encoding + denoising. Disaggregation enables: + +- **Independent scaling** per stage (e.g. 1 encoder : N denoisers : 1 VAE) +- **Pipeline parallelism** — while request N is denoising, request N+1 encodes +- **Memory efficiency** — each GPU loads only its stage's weights +- **Heterogeneous hardware** — encoder on cost-efficient GPUs, denoiser on H100/H200 -# 4 requests, 2 concurrent -NUM_REQUESTS=4 CONCURRENCY=2 python phase1_workers/run_e2e_sglang.py +### 1.2 Architecture -# Custom prompt and parameters -PROMPT="A rocket launching into space" NUM_FRAMES=33 NUM_STEPS=30 \ - python phase1_workers/run_e2e_sglang.py ``` +serve.py (orchestrator, no GPU) +│ +├── Encoder worker GPU 0 Llama3-8B + CLIP (~16 GB) +│ │ +│ │ ── NIXL RDMA (embeddings, ~3.6 MB) ──► +│ ▼ +├── Denoiser worker 0 GPU 1 ┐ +├── Denoiser worker 1 GPU 2 ┘── HunyuanVideo DiT 13B, TP=2 (~13 GB/GPU) +│ │ +│ │ ── NIXL RDMA (latents, ~1.4 MB) ──► +│ ▼ +├── VAE worker GPU 3 3D Causal VAE (~2 GB) +│ +└── FastAPI HTTP server :8090 routes ZMQ metadata between stages +``` + +**Data flow per request:** + +1. Orchestrator sends `Req(prompt)` → Encoder via ZMQ +2. Encoder runs `TextEncodingStage`, registers embeddings as NIXL-readable, + returns metadata via ZMQ (~1.5 KB, no tensor data) +3. Orchestrator forwards NIXL metadata → Denoiser via ZMQ +4. Denoiser's `NixlReceiveStage` RDMA-pulls embeddings directly from Encoder GPU, + runs 50 denoising steps, registers latents as NIXL-readable +5. Orchestrator forwards NIXL metadata → VAE via ZMQ +6. VAE's `NixlReceiveStage` RDMA-pulls latents from Denoiser GPU, + runs 3D VAE decode, returns video frames + +### 1.3 Key Implementation Details -Default GPU layout: Encoder (GPU 0) | Denoiser TP=2 (GPU 1,2) | VAE (GPU 3). +**Partial pipeline loading** — Each stage loads only its required modules via +`build_partial_pipeline()`, which creates a dynamic subclass that suppresses +automatic stage creation and syncs component configs for unloaded modules. -Override with `GPU_ENC`, `GPU_DEN`, `GPU_VAE` environment variables. +**NIXL transfer** (`nixl_transfer.py`) — `NixlTensorSender` flattens GPU +tensors into a contiguous buffer, registers it as NIXL-readable, and returns +a metadata dict. `NixlTensorReceiver` allocates a GPU buffer and RDMA-pulls +the data. Falls back to ZMQ pickle when NIXL is unavailable. -## Architecture +**Dual-encoder support** — HunyuanVideo uses Llama3-8B + CLIP (incompatible +tensor shapes). `NixlSendStage` indexes each element separately +(`prompt_embeds_0`, `prompt_embeds_1`); `NixlReceiveStage` reconstructs the list. +**SGLang compatibility patches** — HunyuanConfig missing `task_type` default +(wrapped `__init__`); triton `norm_infer` non-contiguous assertion (patched +in subprocess entry point). + +## 2. Quick Start + +### HTTP Server (recommended) + +```bash +# Start server — 4 GPU worker processes + HTTP endpoint +python phase1_workers/serve.py + +# Health check +curl http://localhost:8090/health +# → {"status":"ok", "total_gpu_workers":4, "transfer":"NIXL RDMA", ...} + +# Quick test (9 frames, 3 steps, ~10s) +curl -X POST http://localhost:8090/generate \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \ + --output test.mp4 + +# Full quality (61 frames ≈ 2.5s video, 50 steps, ~8 min on H20) +curl -X POST http://localhost:8090/generate \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A rocket launching into space", "num_frames": 61, "num_steps": 50}' \ + --output rocket.mp4 ``` - Encoder (GPU 0) Denoiser (GPU 1,2) VAE (GPU 3) - ┌──────────────┐ ┌──────────────────┐ ┌─────────────┐ - │ Llama3-8B │ embeds │ HunyuanVideo DiT │ lats │ 3D Causal │ - │ + CLIP │───ZMQ──│ 13B (TP=2) │──ZMQ──│ VAE │──► video - │ ~16 GB │ │ ~13 GB/GPU │ │ ~2 GB │ - └──────────────┘ └──────────────────┘ └─────────────┘ + +### Batch / Benchmark Mode + +```bash +python phase1_workers/run_e2e_sglang.py # single request +NUM_REQUESTS=4 CONCURRENCY=2 python phase1_workers/run_e2e_sglang.py # concurrent ``` -Each stage runs as a separate SGLang scheduler subprocess with its own CUDA context. +### Custom GPU Layout + +```bash +# 6 GPUs: encoder=0, denoiser TP=4 on GPUs 1-4, VAE=5 +GPU_ENC=0 GPU_DEN=1,2,3,4 GPU_VAE=5 TP_SIZE=4 python phase1_workers/serve.py +``` -## Supported Models +## 3. Supported Models | Model | Encoders | DiT | Status | |-------|----------|-----|--------| -| **HunyuanVideo v1** (`hunyuanvideo-community/HunyuanVideo`) | Llama3-8B + CLIP | 13B | Default | +| **HunyuanVideo v1** (`hunyuanvideo-community/HunyuanVideo`) | Llama3-8B + CLIP | 13B | Default, validated | | Wan2.2-TI2V-5B (`Wan-AI/Wan2.2-TI2V-5B-Diffusers`) | T5 | 5B | Supported | -To use Wan: `MODEL_PATH=Wan-AI/Wan2.2-TI2V-5B-Diffusers GUIDANCE=5.0 python phase1_workers/run_e2e_sglang.py` +Encoder module detection is automatic via `model_index.json`. Any SGLang-supported +diffusion model with the standard 3-stage structure should work. -## Environment Variables +## 4. Code Structure + +``` +phase1_workers/ +├── serve.py # HTTP server (FastAPI), launches stage workers +├── run_e2e_sglang.py # Batch mode orchestrator with timing report +├── nixl_transfer.py # NixlTensorSender / NixlTensorReceiver (GPU-direct RDMA) +├── partial_gpu_worker.py # PartialGPUWorker, NixlSendStage, NixlReceiveStage +└── sglang_utils.py # build_partial_pipeline(), tensor inject/extract helpers +``` + +## 5. Environment Variables | Variable | Default | Description | |----------|---------|-------------| | `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | -| `PROMPT` | `A cat walking on green grass` | Text prompt | -| `GPU_ENC` | `0` | GPU(s) for encoder | -| `GPU_DEN` | `1,2` | GPU(s) for denoiser | -| `GPU_VAE` | `3` | GPU(s) for VAE | -| `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism degree | -| `NUM_FRAMES` | `61` | Video frames (~2.5s at 24fps) | -| `NUM_STEPS` | `50` | Denoising steps | -| `HEIGHT` | `544` | Frame height | -| `WIDTH` | `960` | Frame width | -| `GUIDANCE` | `1.0` | Guidance scale (>1.0 enables CFG) | -| `NUM_REQUESTS` | `1` | Number of pipeline runs | -| `CONCURRENCY` | `1` | Max concurrent pipelines | - -## Files - -| File | Purpose | -|------|---------| -| `phase1_workers/run_e2e_sglang.py` | Orchestrator: launches stages, runs E2E pipeline, reports timing | -| `phase1_workers/partial_gpu_worker.py` | `PartialGPUWorker`, `IntermediateOutputStage`, subprocess launcher | -| `phase1_workers/sglang_utils.py` | Partial pipeline loading, tensor injection/extraction, config sync | -| `phase0_validate/` | Offline split validation (FLUX.1-schnell, diffusers) | +| `GPU_ENC` | `0` | GPU for encoder | +| `GPU_DEN` | `1,2` | GPUs for denoiser (comma-separated) | +| `GPU_VAE` | `3` | GPU for VAE | +| `TP_SIZE` | auto | Tensor parallelism degree for denoiser | +| `SERVE_PORT` | `8090` | HTTP server port | +| `NUM_FRAMES` | `61` | Video frames (batch mode) | +| `NUM_STEPS` | `50` | Denoising steps (batch mode) | + +## 6. Roadmap + +- [ ] **Dynamo Router integration** — replace manual ZMQ orchestrator with + Router-orchestrated multi-stage chaining, typed stage endpoints +- [ ] **Independent stage scaling** — N:M:K ratio with load-balanced routing +- [ ] **Encoder caching** — LRU cache for repeated prompts +- [ ] **Sequence parallelism** — Ulysses/Ring SP for denoiser beyond TP +- [ ] **Continuous batching** — batch requests at different denoising steps +- [ ] **Heterogeneous hardware** — encoder on L4/T4, denoiser on H100/H200 + +## 7. Dependencies + +```bash +pip install sglang nixl ai-dynamo-runtime imageio imageio-ffmpeg fastapi uvicorn +``` diff --git a/examples/disagg_diffusion/phase1_workers/serve.py b/examples/disagg_diffusion/phase1_workers/serve.py index 7181855503b8..35b55ffe984d 100644 --- a/examples/disagg_diffusion/phase1_workers/serve.py +++ b/examples/disagg_diffusion/phase1_workers/serve.py @@ -4,26 +4,55 @@ """HTTP serving mode for disaggregated diffusion pipeline. -Launches Encoder, Denoiser (TP=2), and VAE stages, then serves -requests via a FastAPI HTTP server. +Launches 3 stage worker groups (Encoder, Denoiser, VAE) on separate GPUs, +then serves video generation requests via HTTP. -Usage: +Process architecture:: + + serve.py (main) + ├── Encoder worker subprocess (GPU 0) ← SGLang Scheduler + Llama + CLIP + ├── Denoiser worker subprocess 0 (GPU 1) ┐ + ├── Denoiser worker subprocess 1 (GPU 2) ┘── SGLang Scheduler + DiT (TP=2) + ├── VAE worker subprocess (GPU 3) ← SGLang Scheduler + 3D VAE + └── FastAPI HTTP server (no GPU) ← orchestrates stages via ZMQ + +Inter-stage tensor data transfers use NIXL RDMA (GPU-direct). +Only metadata (~1.5 KB) travels over ZMQ. + +Usage:: + + # Start server (default: 4 GPUs) python serve.py - # Then send requests: - curl -X POST http://localhost:8090/generate \ - -H "Content-Type: application/json" \ - -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \ - --output output.mp4 + # Custom GPU layout + GPU_ENC=0 GPU_DEN=1,2,3,4 GPU_VAE=5 TP_SIZE=4 python serve.py - # Check server health: + # Send requests curl http://localhost:8090/health + + curl -X POST http://localhost:8090/generate \\ + -H "Content-Type: application/json" \\ + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \\ + --output output.mp4 + + # Longer video (61 frames ≈ 2.5s at 24fps, ~8 min on H20) + curl -X POST http://localhost:8090/generate \\ + -H "Content-Type: application/json" \\ + -d '{"prompt": "A rocket launching into space", "num_frames": 61, "num_steps": 50}' \\ + --output rocket.mp4 + +Environment variables: + MODEL_PATH HuggingFace model (default: hunyuanvideo-community/HunyuanVideo) + GPU_ENC GPU for encoder (default: 0) + GPU_DEN GPUs for denoiser, comma-separated (default: 1,2) + GPU_VAE GPU for VAE (default: 3) + TP_SIZE Tensor parallelism degree (default: auto from GPU_DEN count) + SERVE_HOST HTTP bind address (default: 0.0.0.0) + SERVE_PORT HTTP port (default: 8090) """ from __future__ import annotations -import asyncio -import base64 import io import logging import multiprocessing as mp @@ -32,7 +61,6 @@ import time from typing import Optional -# Ensure workers dir on path WORKERS_DIR = os.path.dirname(os.path.abspath(__file__)) if WORKERS_DIR not in sys.path: sys.path.insert(0, WORKERS_DIR) @@ -51,10 +79,8 @@ TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(","))))) SERVE_HOST = os.environ.get("SERVE_HOST", "0.0.0.0") SERVE_PORT = int(os.environ.get("SERVE_PORT", "8090")) -OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_e2e") - -# ── Reuse run_e2e components ───────────────────────────────────────────── +# ── Reuse core components ──────────────────────────────────────────────── from run_e2e_sglang import ( _patch_hunyuan_config_task_type, _detect_encoder_modules, @@ -71,6 +97,8 @@ _request_counter = 0 +# ── Pipeline execution ─────────────────────────────────────────────────── + async def run_pipeline( prompt: str, negative_prompt: str = "", @@ -81,7 +109,7 @@ async def run_pipeline( guidance_scale: float = 1.0, seed: int = 42, ) -> bytes: - """Run the full E2E pipeline and return mp4 bytes.""" + """Run full E2E pipeline, return mp4 bytes.""" global _request_counter import torch from sglang_utils import build_req, inject_tensors_to_req @@ -95,42 +123,49 @@ async def run_pipeline( num_inference_steps=num_steps, guidance_scale=guidance_scale, seed=seed, ) - t_start = time.monotonic() + t_total = time.monotonic() - # Encoder + # ── Encoder ── t0 = time.monotonic() enc_output = await enc_client.forward([build_req(**req_kwargs)]) - enc_time = time.monotonic() - t0 + t_enc = time.monotonic() - t0 if enc_output.error: raise RuntimeError(f"Encoder error: {enc_output.error}") - logger.info("req %d | Encoder %.2fs", req_id, enc_time) + enc_result = enc_output.output + nixl_enc = "_nixl_transfer_meta" in enc_result - # Denoiser + # ── Denoiser ── t0 = time.monotonic() den_req = build_req(**req_kwargs) - inject_tensors_to_req(den_req, enc_output.output) den_req.do_classifier_free_guidance = (guidance_scale > 1.0) + if nixl_enc: + den_req._nixl_transfer_meta = enc_result["_nixl_transfer_meta"] + else: + inject_tensors_to_req(den_req, enc_result) den_output = await den_client.forward([den_req]) - den_time = time.monotonic() - t0 + t_den = time.monotonic() - t0 if den_output.error: raise RuntimeError(f"Denoiser error: {den_output.error}") - logger.info("req %d | Denoiser %.2fs", req_id, den_time) + den_result = den_output.output + nixl_den = "_nixl_transfer_meta" in den_result - # VAE + # ── VAE ── t0 = time.monotonic() vae_req = build_req(prompt="", height=height, width=width, num_frames=num_frames, num_inference_steps=num_steps, guidance_scale=0.0, seed=seed) - vae_req.latents = den_output.output["latents"].cpu() + if nixl_den: + vae_req._nixl_transfer_meta = den_result["_nixl_transfer_meta"] + else: + vae_req.latents = den_result["latents"].cpu() vae_output = await vae_client.forward([vae_req]) - vae_time = time.monotonic() - t0 + t_vae = time.monotonic() - t0 if vae_output.error: raise RuntimeError(f"VAE error: {vae_output.error}") - total = time.monotonic() - t_start - logger.info("req %d | VAE %.2fs | Total %.2fs", req_id, vae_time, total) + t_elapsed = time.monotonic() - t_total - # Convert to mp4 bytes + # ── Encode as mp4 ── import numpy as np import imageio @@ -143,10 +178,11 @@ async def run_pipeline( imageio.mimwrite(buf, frames, format="mp4", fps=24, codec="libx264") mp4_bytes = buf.getvalue() + transfer = "NIXL" if (nixl_enc and nixl_den) else "ZMQ" logger.info( - "req %d | Done: %d frames, %dx%d, %.1f KB, enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", - req_id, frames.shape[0], width, height, len(mp4_bytes) / 1024, - enc_time, den_time, vae_time, total, + "req %d | %d frames %dx%d | enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs | %s | %.1f KB", + req_id, frames.shape[0], width, height, + t_enc, t_den, t_vae, t_elapsed, transfer, len(mp4_bytes) / 1024, ) return mp4_bytes @@ -173,20 +209,30 @@ class GenerateRequest(BaseModel): @app.get("/health") async def health(): - return {"status": "ok", "model": MODEL_PATH} + n_enc = len(GPU_ENC.split(",")) + n_den = len(GPU_DEN.split(",")) + n_vae = len(GPU_VAE.split(",")) + return { + "status": "ok", + "model": MODEL_PATH, + "gpus": { + "encoder": f"GPU {GPU_ENC} ({n_enc} process{'es' if n_enc > 1 else ''})", + "denoiser": f"GPU {GPU_DEN} ({n_den} processes, TP={TP_SIZE})", + "vae": f"GPU {GPU_VAE} ({n_vae} process{'es' if n_vae > 1 else ''})", + }, + "total_gpu_workers": n_enc + n_den + n_vae, + "transfer": "NIXL RDMA", + "requests_served": _request_counter, + } @app.post("/generate") async def generate(req: GenerateRequest): try: mp4_bytes = await run_pipeline( - prompt=req.prompt, - negative_prompt=req.negative_prompt, - height=req.height, - width=req.width, - num_frames=req.num_frames, - num_steps=req.num_steps, - guidance_scale=req.guidance_scale, + prompt=req.prompt, negative_prompt=req.negative_prompt, + height=req.height, width=req.width, num_frames=req.num_frames, + num_steps=req.num_steps, guidance_scale=req.guidance_scale, seed=req.seed, ) return Response(content=mp4_bytes, media_type="video/mp4") @@ -195,18 +241,24 @@ async def generate(req: GenerateRequest): raise HTTPException(status_code=500, detail=str(e)) -# ── Startup / Shutdown ─────────────────────────────────────────────────── +# ── Stage lifecycle ────────────────────────────────────────────────────── def launch_all_stages(): - """Launch all 3 stage schedulers and return clients.""" global enc_client, den_client, vae_client, _all_procs from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages _patch_hunyuan_config_task_type() - logger.info("Launching stages: Encoder=%s Denoiser=%s (TP=%d) VAE=%s", - GPU_ENC, GPU_DEN, TP_SIZE, GPU_VAE) + t0 = time.monotonic() + logger.info("=" * 60) + logger.info(" Launching disaggregated diffusion stages") + logger.info(" Model: %s", MODEL_PATH) + logger.info(" Encoder: GPU %s (1 process)", GPU_ENC) + logger.info(" Denoiser: GPU %s (%d processes, TP=%d)", GPU_DEN, len(GPU_DEN.split(",")), TP_SIZE) + logger.info(" VAE: GPU %s (1 process)", GPU_VAE) + logger.info(" Total: %d GPU worker processes", len(GPU_ENC.split(",")) + len(GPU_DEN.split(",")) + len(GPU_VAE.split(","))) + logger.info("=" * 60) enc_procs, enc_args = _launch_stage( "Encoder", GPU_ENC, @@ -233,7 +285,18 @@ def launch_all_stages(): den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") - logger.info("All stages ready. Serving at http://%s:%d", SERVE_HOST, SERVE_PORT) + logger.info("=" * 60) + logger.info(" All %d workers ready in %.1fs", len(_all_procs), time.monotonic() - t0) + logger.info(" HTTP server: http://%s:%d", SERVE_HOST, SERVE_PORT) + logger.info(" Transfer: NIXL RDMA (GPU-direct)") + logger.info("") + logger.info(" Try:") + logger.info(" curl http://localhost:%d/health", SERVE_PORT) + logger.info(' curl -X POST http://localhost:%d/generate \\', SERVE_PORT) + logger.info(' -H "Content-Type: application/json" \\') + logger.info(' -d \'{"prompt": "A cat on grass", "num_frames": 9, "num_steps": 3}\' \\') + logger.info(" --output test.mp4") + logger.info("=" * 60) def shutdown_stages(): @@ -247,12 +310,11 @@ def shutdown_stages(): p.terminate() for p in _all_procs: p.join(timeout=10) - logger.info("All stages terminated.") + logger.info("All workers terminated.") if __name__ == "__main__": mp.set_start_method("spawn", force=True) - launch_all_stages() import uvicorn From 0148f344c26d5dce347fcdcbc4497b48c42fb9d0 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 22:13:38 +0800 Subject: [PATCH 05/22] feat: add Dynamo RPC endpoints for disagg diffusion with HunyuanVideo Replace ZMQ orchestrator (serve.py) with Dynamo RPC workers. Each worker spawns SGLang Scheduler subprocess(es) via launch_partial_server() and bridges Dynamo RPC <-> ZMQ, enabling TP support and NIXL RDMA GPU-direct tensor transfer between stages. Key changes: - Workers use launch_partial_server() + StageClient instead of serve.py - NixlSendStage/NixlReceiveStage for GPU-direct inter-stage transfer - Dual encoder auto-detection (Llama 8B + CLIP for HunyuanVideo) - Triton norm contiguous patch for HunyuanVideo transformer - HunyuanVideo defaults (544x960, 50 steps, guidance_scale=1.0) - Removed serve.py (ZMQ orchestrator, replaced by Dynamo RPC workers) - Added Dynamo RPC endpoint wrappers in encoder/denoiser/vae_worker.py - Simplified protocol.py (removed NixlTensorSender/Receiver) - Updated run_disagg.py orchestrator for HunyuanVideo defaults Tested: HunyuanVideo 13B, 4 GPUs (Enc=1, Den TP=2, VAE=1), NIXL RDMA, 30 steps 544x960 -> high quality output verified. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 213 ++++----- .../phase1_workers/denoiser_worker.py | 229 +++++----- .../phase1_workers/encoder_worker.py | 173 +++++--- .../phase1_workers/nixl_transfer.py | 6 +- .../phase1_workers/partial_gpu_worker.py | 48 +- .../phase1_workers/protocol.py | 118 ++--- .../phase1_workers/run_e2e_sglang.py | 10 +- .../disagg_diffusion/phase1_workers/serve.py | 324 -------------- .../phase1_workers/sglang_utils.py | 18 +- .../phase1_workers/vae_worker.py | 230 ++++++---- .../phase2_orchestrator/run_disagg.py | 417 +++++++++++++----- 11 files changed, 846 insertions(+), 940 deletions(-) delete mode 100644 examples/disagg_diffusion/phase1_workers/serve.py diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 30529c886b5b..9d5a1b75e0c1 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,160 +1,133 @@ -# Disaggregated Diffusion Inference +# Disaggregated Diffusion Inference (HunyuanVideo) -Decomposes a monolithic video diffusion pipeline into three independently -deployable stages — **Text Encoder**, **DiT Denoiser**, and **3D VAE** — -each running on separate GPUs. Inter-stage tensor data transfers use -**NIXL RDMA** (GPU-direct); only small metadata (~1.5 KB) travels over -the ZMQ control plane. +Split a monolithic video diffusion pipeline (Text Encoder -> Transformer -> VAE) into +independent stages on separate GPUs. Tensor data transfers between stages use +**NIXL RDMA** (GPU-direct); only small metadata travels over Dynamo RPC. -## 1. Design +Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. -### 1.1 Motivation - -In monolithic diffusion inference, all model components occupy a single GPU -for the entire request lifetime. The encoder (~16 GB) sits idle during the -denoising loop (96%+ of wall time), and the VAE (~2 GB) sits idle during -encoding + denoising. Disaggregation enables: - -- **Independent scaling** per stage (e.g. 1 encoder : N denoisers : 1 VAE) -- **Pipeline parallelism** — while request N is denoising, request N+1 encodes -- **Memory efficiency** — each GPU loads only its stage's weights -- **Heterogeneous hardware** — encoder on cost-efficient GPUs, denoiser on H100/H200 - -### 1.2 Architecture +## Architecture ``` -serve.py (orchestrator, no GPU) -│ -├── Encoder worker GPU 0 Llama3-8B + CLIP (~16 GB) -│ │ -│ │ ── NIXL RDMA (embeddings, ~3.6 MB) ──► -│ ▼ -├── Denoiser worker 0 GPU 1 ┐ -├── Denoiser worker 1 GPU 2 ┘── HunyuanVideo DiT 13B, TP=2 (~13 GB/GPU) -│ │ -│ │ ── NIXL RDMA (latents, ~1.4 MB) ──► -│ ▼ -├── VAE worker GPU 3 3D Causal VAE (~2 GB) -│ -└── FastAPI HTTP server :8090 routes ZMQ metadata between stages + Dynamo RPC (metadata only) + ┌────────────┬────────────────────────┐ + │ │ │ + ▼ ▼ ▼ +GPU 0: Encoder Worker GPU 1,2: Denoiser Worker GPU 3: VAE Worker + │ (Llama + CLIP) │ (DiT, TP=2) │ (3D VAE) + │ │ │ + └── NIXL RDMA ──────────┘── NIXL RDMA ──────────────┘ + (embeddings) (latents) + ▲ + │ + Orchestrator (HTTP, no GPU) ``` -**Data flow per request:** +Each worker is a Dynamo `@dynamo_worker` that: +1. Spawns SGLang Scheduler subprocess(es) via `launch_partial_server()` +2. Connects a ZMQ `StageClient` to its local Scheduler +3. Exposes `serve_endpoint("generate")` for Dynamo RPC +4. Bridges RPC requests to the Scheduler, which runs the model + NIXL transfer -1. Orchestrator sends `Req(prompt)` → Encoder via ZMQ -2. Encoder runs `TextEncodingStage`, registers embeddings as NIXL-readable, - returns metadata via ZMQ (~1.5 KB, no tensor data) -3. Orchestrator forwards NIXL metadata → Denoiser via ZMQ -4. Denoiser's `NixlReceiveStage` RDMA-pulls embeddings directly from Encoder GPU, - runs 50 denoising steps, registers latents as NIXL-readable -5. Orchestrator forwards NIXL metadata → VAE via ZMQ -6. VAE's `NixlReceiveStage` RDMA-pulls latents from Denoiser GPU, - runs 3D VAE decode, returns video frames +## Quick Start -### 1.3 Key Implementation Details - -**Partial pipeline loading** — Each stage loads only its required modules via -`build_partial_pipeline()`, which creates a dynamic subclass that suppresses -automatic stage creation and syncs component configs for unloaded modules. +```bash +conda activate omni +export HF_HUB_CACHE=/path/to/huggingface/hub -**NIXL transfer** (`nixl_transfer.py`) — `NixlTensorSender` flattens GPU -tensors into a contiguous buffer, registers it as NIXL-readable, and returns -a metadata dict. `NixlTensorReceiver` allocates a GPU buffer and RDMA-pulls -the data. Falls back to ZMQ pickle when NIXL is unavailable. +# Terminal 0: etcd (service discovery) +etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 -**Dual-encoder support** — HunyuanVideo uses Llama3-8B + CLIP (incompatible -tensor shapes). `NixlSendStage` indexes each element separately -(`prompt_embeds_0`, `prompt_embeds_1`); `NixlReceiveStage` reconstructs the list. +# Terminal 1: Encoder Worker (Llama 8B + CLIP, ~18 GB VRAM) +CUDA_VISIBLE_DEVICES=0 python phase1_workers/encoder_worker.py -**SGLang compatibility patches** — HunyuanConfig missing `task_type` default -(wrapped `__init__`); triton `norm_infer` non-contiguous assertion (patched -in subprocess entry point). +# Terminal 2: Denoiser Worker (DiT TP=2, ~24 GB VRAM per GPU) +CUDA_VISIBLE_DEVICES=1,2 python phase1_workers/denoiser_worker.py -## 2. Quick Start +# Terminal 3: VAE Worker (~6 GB VRAM) +CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py -### HTTP Server (recommended) +# Terminal 4: Orchestrator (no GPU) +python phase2_orchestrator/run_disagg.py +``` +Test: ```bash -# Start server — 4 GPU worker processes + HTTP endpoint -python phase1_workers/serve.py - -# Health check -curl http://localhost:8090/health -# → {"status":"ok", "total_gpu_workers":4, "transfer":"NIXL RDMA", ...} - -# Quick test (9 frames, 3 steps, ~10s) -curl -X POST http://localhost:8090/generate \ +curl -X POST http://localhost:8080/v1/videos/generations \ -H "Content-Type: application/json" \ - -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \ - --output test.mp4 - -# Full quality (61 frames ≈ 2.5s video, 50 steps, ~8 min on H20) -curl -X POST http://localhost:8090/generate \ - -H "Content-Type: application/json" \ - -d '{"prompt": "A rocket launching into space", "num_frames": 61, "num_steps": 50}' \ - --output rocket.mp4 + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_inference_steps": 3}' ``` -### Batch / Benchmark Mode +### Standalone E2E (no Dynamo, no etcd) + +Single-process launcher that starts all 3 stages and runs the pipeline: ```bash -python phase1_workers/run_e2e_sglang.py # single request -NUM_REQUESTS=4 CONCURRENCY=2 python phase1_workers/run_e2e_sglang.py # concurrent +python phase1_workers/run_e2e_sglang.py ``` -### Custom GPU Layout +## Phases + +### Phase 0: Offline Validation (no Dynamo) + +Single-GPU script proving diffusers supports split execution. ```bash -# 6 GPUs: encoder=0, denoiser TP=4 on GPUs 1-4, VAE=5 -GPU_ENC=0 GPU_DEN=1,2,3,4 GPU_VAE=5 TP_SIZE=4 python phase1_workers/serve.py +python phase0_validate/validate_split.py \ + --model hunyuanvideo-community/HunyuanVideo \ + --prompt "A cat walking on grass" \ + --num-steps 3 --num-frames 9 \ + --output-dir /tmp/disagg_validate ``` -## 3. Supported Models +### Phase 1: Dynamo Stage Workers (NIXL + SGLang Scheduler) + +Three Dynamo workers, each wrapping an SGLang Scheduler subprocess: -| Model | Encoders | DiT | Status | -|-------|----------|-----|--------| -| **HunyuanVideo v1** (`hunyuanvideo-community/HunyuanVideo`) | Llama3-8B + CLIP | 13B | Default, validated | -| Wan2.2-TI2V-5B (`Wan-AI/Wan2.2-TI2V-5B-Diffusers`) | T5 | 5B | Supported | +| Worker | Model Component | VRAM | Endpoint | +|--------|----------------|------|----------| +| `encoder_worker.py` | Llama 8B + CLIP text encoders | ~18 GB | `disagg_diffusion.encoder.generate` | +| `denoiser_worker.py` | HunyuanVideo DiT (TP=2) | ~24 GB/GPU | `disagg_diffusion.denoiser.generate` | +| `vae_worker.py` | 3D VAE decoder | ~6 GB | `disagg_diffusion.vae.generate` | -Encoder module detection is automatic via `model_index.json`. Any SGLang-supported -diffusion model with the standard 3-stage structure should work. +### Phase 2: Orchestrator (Pipeline Parallel) -## 4. Code Structure +Chains three stage endpoints with pipeline parallelism. Multiple concurrent +requests overlap across stages: ``` -phase1_workers/ -├── serve.py # HTTP server (FastAPI), launches stage workers -├── run_e2e_sglang.py # Batch mode orchestrator with timing report -├── nixl_transfer.py # NixlTensorSender / NixlTensorReceiver (GPU-direct RDMA) -├── partial_gpu_worker.py # PartialGPUWorker, NixlSendStage, NixlReceiveStage -└── sglang_utils.py # build_partial_pipeline(), tensor inject/extract helpers +Request 1: [Encoder] → [Denoiser] → [ VAE ] +Request 2: [Encoder] → [Denoiser] → [ VAE ] +Request 3: [Encoder] → [Denoiser] → ... ``` -## 5. Environment Variables +## Environment Variables + +### Worker Configuration | Variable | Default | Description | |----------|---------|-------------| -| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | -| `GPU_ENC` | `0` | GPU for encoder | -| `GPU_DEN` | `1,2` | GPUs for denoiser (comma-separated) | -| `GPU_VAE` | `3` | GPU for VAE | -| `TP_SIZE` | auto | Tensor parallelism degree for denoiser | -| `SERVE_PORT` | `8090` | HTTP server port | -| `NUM_FRAMES` | `61` | Video frames (batch mode) | -| `NUM_STEPS` | `50` | Denoising steps (batch mode) | - -## 6. Roadmap - -- [ ] **Dynamo Router integration** — replace manual ZMQ orchestrator with - Router-orchestrated multi-stage chaining, typed stage endpoints -- [ ] **Independent stage scaling** — N:M:K ratio with load-balanced routing -- [ ] **Encoder caching** — LRU cache for repeated prompts -- [ ] **Sequence parallelism** — Ulysses/Ring SP for denoiser beyond TP -- [ ] **Continuous batching** — batch requests at different denoising steps -- [ ] **Heterogeneous hardware** — encoder on L4/T4, denoiser on H100/H200 - -## 7. Dependencies +| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model path | +| `SCHEDULER_PORT` | `15600/15700/15800` | ZMQ port for local Scheduler | +| `TP_SIZE` | auto from `CUDA_VISIBLE_DEVICES` | Tensor parallelism (denoiser) | +| `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory (VAE) | + +### Orchestrator Configuration + +| Variable | Default | Description | +|----------|---------|-------------| +| `PORT` | `8080` | HTTP server port | +| `MAX_PIPELINE_DEPTH` | `4` | Max concurrent requests in pipeline | +| `OUTPUT_DIR` | `/tmp/disagg_videos` | Shared directory for video output | + +## Supported Models + +- `hunyuanvideo-community/HunyuanVideo` (13B, dual encoder, recommended) +- `Wan-AI/Wan2.2-TI2V-5B-Diffusers` (5B, single encoder) +- Any SGLang-supported diffusion model with encoder/denoiser/VAE stages + +## Dependencies ```bash -pip install sglang nixl ai-dynamo-runtime imageio imageio-ffmpeg fastapi uvicorn +pip install ai-dynamo-runtime sglang imageio imageio-ffmpeg pyzmq setproctitle etcd-distro ``` diff --git a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py index ebc29696ba20..b214830ebeb3 100755 --- a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py +++ b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py @@ -1,136 +1,150 @@ #!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — Denoiser Worker - -Loads only the Transformer (DiT) and scheduler. Accepts pre-computed -embeddings and returns denoised latents (skips VAE decode). - -Usage: - python denoiser_worker.py --model black-forest-labs/FLUX.1-schnell +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — Denoiser Worker (Dynamo RPC) + +Wraps SGLang Scheduler subprocess(es) running NixlReceiveStage + Denoising + +NixlSendStage. Supports TP via launch_partial_server(tp_size=N). + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess(es) (spawned by launch_partial_server) + +-- PartialGPUWorker (TP=N) + |-- NixlReceiveStage <-- RDMA-pull embeddings from encoder + |-- LatentPreparationStage + |-- TimestepPreparationStage + |-- DenoisingStage + +-- NixlSendStage <-- register latents as NIXL-readable """ import asyncio +import json import logging +import multiprocessing as mp import os import sys -import torch import uvloop sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from protocol import ( # noqa: E402 - DenoiserRequest, - DenoiserResponse, - b64_to_tensors, - tensors_to_b64, -) - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 logger = logging.getLogger(__name__) -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15700")) -class DenoiserStage: - """Denoiser stage: embeddings → latents (no text encoder, no VAE).""" +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from run_e2e_sglang import ( + _patch_hunyuan_config_task_type, + StageClient, + ) + from partial_gpu_worker import build_denoiser_stages, launch_partial_server + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from sglang_utils import build_req - def __init__(self): - self.pipe = None - self.vae_scaling_factor = 1.0 - self.vae_shift_factor = None + _patch_hunyuan_config_task_type() - def load_model(self): - from diffusers import FluxPipeline + # Auto-detect GPU count from CUDA_VISIBLE_DEVICES + num_gpus = len(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")) + tp_size = int(os.environ.get("TP_SIZE", str(num_gpus))) - logger.info("Loading transformer from %s …", MODEL_PATH) - self.pipe = FluxPipeline.from_pretrained( - MODEL_PATH, torch_dtype=torch.bfloat16 - ) - self.pipe.to(DEVICE) - - # Capture VAE config before discarding - self.vae_scaling_factor = self.pipe.vae.config.scaling_factor - self.vae_shift_factor = getattr(self.pipe.vae.config, "shift_factor", None) - - # Free text encoders + VAE - self.pipe.text_encoder = None - self.pipe.text_encoder_2 = None - self.pipe.tokenizer = None - self.pipe.tokenizer_2 = None - self.pipe.vae = None - torch.cuda.empty_cache() - - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("Denoiser ready — VRAM: %.0f MB (transformer only)", vram) - - @dynamo_endpoint(DenoiserRequest, DenoiserResponse) - async def generate(self, request: DenoiserRequest): - logger.info( - "Denoising %dx%d, %d steps, seed=%d", - request.width, request.height, - request.num_inference_steps, request.seed, - ) + server_args = ServerArgs.from_kwargs( + model_path=MODEL_PATH, + num_gpus=num_gpus, + tp_size=tp_size, + scheduler_port=SCHEDULER_PORT, + ) + set_global_server_args(server_args) - embeddings = b64_to_tensors(request.embeddings_b64, DEVICE) - generator = torch.Generator(device=DEVICE).manual_seed(request.seed) - - loop = asyncio.get_event_loop() - result = await loop.run_in_executor( - None, - lambda: self.pipe( - prompt_embeds=embeddings["prompt_embeds"], - pooled_prompt_embeds=embeddings["pooled_prompt_embeds"], - num_inference_steps=request.num_inference_steps, - guidance_scale=request.guidance_scale, - generator=generator, - height=request.height, - width=request.width, - output_type="latent", - ), - ) - latents = result.images + logger.info( + "Launching denoiser Scheduler: num_gpus=%d, tp=%d, port=%d", + num_gpus, tp_size, SCHEDULER_PORT, + ) + processes = launch_partial_server( + server_args, + required_modules=["transformer", "scheduler"], + custom_stages_fn=build_denoiser_stages, + ) - latent_payload = { - "latents": latents, - "scaling_factor": torch.tensor(self.vae_scaling_factor), + # Connect ZMQ client to local Scheduler + client = StageClient(server_args.scheduler_endpoint, "denoiser") + + # ── Dynamo RPC handlers ────────────────────────────────────────── + + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) + + req = build_req( + prompt="(embeddings via NIXL)", + negative_prompt="", + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=request.get("num_inference_steps", 50), + guidance_scale=request.get("guidance_scale", 1.0), + seed=request.get("seed", 42), + ) + req.do_classifier_free_guidance = (req.guidance_scale > 1.0) + + # Pass NIXL metadata for NixlReceiveStage to RDMA-pull embeddings + transfer_meta = request.get("transfer_meta", {}) + if transfer_meta: + req._nixl_transfer_meta = transfer_meta + + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "transfer_meta": {}, "shape": []} + return + + result = output.output + transfer_meta_out = result.get("_nixl_transfer_meta", {}) + logger.info("Denoised — NIXL latent metadata ready") + yield {"transfer_meta": transfer_meta_out, "shape": []} + + except Exception as e: + logger.error("Denoiser generate failed: %s", e, exc_info=True) + yield {"error": str(e), "transfer_meta": {}, "shape": []} + + async def handle_health(request, context): + yield { + "status": "ok", "stage": "denoiser", + "model": MODEL_PATH, "tp_size": tp_size, } - if self.vae_shift_factor is not None: - latent_payload["shift_factor"] = torch.tensor(self.vae_shift_factor) - response = DenoiserResponse( - latents_b64=tensors_to_b64(latent_payload), - shape=list(latents.shape), - ) - logger.info("Denoised — latents %s", list(latents.shape)) - yield response.model_dump() + # ── Serve Dynamo endpoints ─────────────────────────────────────── + ns = runtime.namespace("disagg_diffusion") + gen_ep = ns.component("denoiser").endpoint("generate") + health_ep = ns.component("denoiser").endpoint("health") -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.denoiser.generate") - - stage = DenoiserStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) - - logger.info("Serving denoiser endpoint: disagg_diffusion.denoiser.generate") - await endpoint.serve_endpoint(stage.generate) + logger.info("Serving: disagg_diffusion.denoiser.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) if __name__ == "__main__": @@ -138,5 +152,6 @@ async def worker(runtime: DistributedRuntime): level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s", ) + mp.set_start_method("spawn", force=True) uvloop.install() asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase1_workers/encoder_worker.py b/examples/disagg_diffusion/phase1_workers/encoder_worker.py index f52c2fabfd14..7f678f5a6320 100755 --- a/examples/disagg_diffusion/phase1_workers/encoder_worker.py +++ b/examples/disagg_diffusion/phase1_workers/encoder_worker.py @@ -1,105 +1,133 @@ #!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — Encoder Worker - -Loads only the text encoders (CLIP + T5) and serves an endpoint that -converts text prompts into serialized embeddings. - -Usage: - python encoder_worker.py --model black-forest-labs/FLUX.1-schnell +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — Encoder Worker (Dynamo RPC) + +Wraps an SGLang Scheduler subprocess running TextEncodingStage + NixlSendStage. +Exposes a Dynamo RPC endpoint. The Scheduler handles model loading, GPU +execution, and NIXL tensor registration. This worker bridges Dynamo RPC <-> ZMQ. + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess (spawned by launch_partial_server) + +-- PartialGPUWorker + |-- TextEncodingStage + +-- NixlSendStage <-- registers embeddings as NIXL-readable """ import asyncio +import json import logging +import multiprocessing as mp import os import sys -import torch import uvloop -# Allow importing protocol.py from the same directory sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from protocol import EncoderRequest, EncoderResponse, tensors_to_b64 # noqa: E402 - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 logger = logging.getLogger(__name__) -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15600")) -class EncoderStage: - """Text encoder stage: CLIP + T5 → embeddings.""" +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from run_e2e_sglang import ( + _patch_hunyuan_config_task_type, + _detect_encoder_modules, + StageClient, + ) + from partial_gpu_worker import build_encoder_stages, launch_partial_server + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from sglang_utils import build_req - def __init__(self): - self.pipe = None + _patch_hunyuan_config_task_type() - def load_model(self): - from diffusers import FluxPipeline + # Launch SGLang Scheduler subprocess with text encoder stages + enc_modules = _detect_encoder_modules(MODEL_PATH) + server_args = ServerArgs.from_kwargs( + model_path=MODEL_PATH, + num_gpus=1, + tp_size=1, + scheduler_port=SCHEDULER_PORT, + ) + set_global_server_args(server_args) - logger.info("Loading text encoders from %s …", MODEL_PATH) - self.pipe = FluxPipeline.from_pretrained( - MODEL_PATH, torch_dtype=torch.bfloat16 - ) - self.pipe.to(DEVICE) + logger.info( + "Launching encoder Scheduler: modules=%s, port=%d", + enc_modules, SCHEDULER_PORT, + ) + processes = launch_partial_server( + server_args, + required_modules=enc_modules, + custom_stages_fn=build_encoder_stages, + ) - # Free transformer + VAE — we only need text encoders - self.pipe.transformer = None - self.pipe.vae = None - torch.cuda.empty_cache() + # Connect ZMQ client to local Scheduler + client = StageClient(server_args.scheduler_endpoint, "encoder") - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("Encoder ready — VRAM: %.0f MB (text encoders only)", vram) + # ── Dynamo RPC handlers ────────────────────────────────────────── - @dynamo_endpoint(EncoderRequest, EncoderResponse) - async def generate(self, request: EncoderRequest): - logger.info("Encoding prompt: %.80s…", request.prompt) + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) - loop = asyncio.get_event_loop() - prompt_embeds, pooled_prompt_embeds, text_ids = await loop.run_in_executor( - None, - lambda: self.pipe.encode_prompt(prompt=request.prompt, prompt_2=None), - ) + req = build_req( + prompt=request.get("prompt", ""), + negative_prompt=request.get("negative_prompt", ""), + guidance_scale=request.get("guidance_scale", 1.0), + ) - embeddings = { - "prompt_embeds": prompt_embeds, - "pooled_prompt_embeds": pooled_prompt_embeds, - "text_ids": text_ids, - } + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "transfer_meta": {}, "shapes": {}} + return - response = EncoderResponse( - embeddings_b64=tensors_to_b64(embeddings), - shapes={k: list(v.shape) for k, v in embeddings.items()}, - ) - logger.info("Encoded — prompt_embeds %s", list(prompt_embeds.shape)) - yield response.model_dump() + result = output.output + transfer_meta = result.get("_nixl_transfer_meta", {}) + logger.info("Encoded prompt — NIXL metadata ready") + yield {"transfer_meta": transfer_meta, "shapes": {}} + except Exception as e: + logger.error("Encoder generate failed: %s", e, exc_info=True) + yield {"error": str(e), "transfer_meta": {}, "shapes": {}} -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.encoder.generate") + async def handle_health(request, context): + yield {"status": "ok", "stage": "encoder", "model": MODEL_PATH} + + # ── Serve Dynamo endpoints ─────────────────────────────────────── - stage = EncoderStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) + ns = runtime.namespace("disagg_diffusion") + gen_ep = ns.component("encoder").endpoint("generate") + health_ep = ns.component("encoder").endpoint("health") - logger.info("Serving encoder endpoint: disagg_diffusion.encoder.generate") - await endpoint.serve_endpoint(stage.generate) + logger.info("Serving: disagg_diffusion.encoder.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) if __name__ == "__main__": @@ -107,5 +135,6 @@ async def worker(runtime: DistributedRuntime): level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s", ) + mp.set_start_method("spawn", force=True) uvloop.install() asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py index 234606ba2f47..52bfe9fe404c 100644 --- a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py +++ b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py @@ -6,7 +6,7 @@ Provides GPU-direct tensor transfer between stage workers. Only small metadata (shapes, dtypes, NIXL descriptor ~1.5 KB) travels over the ZMQ control plane; actual tensor data (embeddings, latents) transfers -GPU→GPU via NIXL RDMA. +GPU->GPU via NIXL RDMA. Usage inside PipelineStage.forward() (synchronous context):: @@ -39,10 +39,10 @@ class _PersistentConnector: """Lazily-initialized NIXL Connector singleton per process.""" - _instance: Optional[nixl_connect.Connector] = None if NIXL_AVAILABLE else None + _instance = None @classmethod - async def get(cls) -> nixl_connect.Connector: + async def get(cls): if cls._instance is None: cls._instance = nixl_connect.Connector() await cls._instance.initialize() diff --git a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py index 546dbdf69f11..72dac1566830 100644 --- a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""PartialGPUWorker, IntermediateOutputStage, and subprocess launcher. +"""PartialGPUWorker, NixlSendStage, NixlReceiveStage, and subprocess launcher. This module is the single source of truth for disaggregated diffusion's integration with sglang. Everything that runs **inside the Scheduler @@ -20,12 +20,15 @@ └─ pipeline.forward() SchedulerClient ◄────ZMQ────► Scheduler.event_loop() -``IntermediateOutputStage`` is appended as the last pipeline stage for -encoder / denoiser. It packages ``Req`` tensors into ``OutputBatch`` -so the result travels through ZMQ as a standard ``OutputBatch``. +``NixlSendStage`` is appended as the last pipeline stage for +encoder / denoiser. It registers ``Req`` tensors as NIXL-readable and +returns an ``OutputBatch`` with NIXL metadata. + +``NixlReceiveStage`` is prepended at the start of denoiser / VAE +pipelines. It RDMA-pulls tensors from the previous stage. ``DecodingStage`` (VAE) already returns ``OutputBatch``, so no extra -stage is needed for the VAE worker. +send stage is needed for the VAE worker. """ from __future__ import annotations @@ -54,8 +57,8 @@ get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker -from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req, OutputBatch -from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs # layerwise_offload may not exist in all sglang versions — guard import @@ -68,7 +71,7 @@ # ═══════════════════════════════════════════════════════════════════════ -# IntermediateOutputStage +# NIXL Pipeline Stages (run inside Scheduler subprocess) # ═══════════════════════════════════════════════════════════════════════ @@ -150,10 +153,19 @@ def __init__(self, output_fields: List[str]): self._output_fields = output_fields self._sender = None + @staticmethod + def _make_timings(): + """Create a RequestTimings so gpu_worker.execute_forward doesn't crash.""" + try: + from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings + return RequestTimings(request_id="nixl") + except Exception: + return None + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: tensors = self._extract_tensors(batch) if not tensors: - return OutputBatch(output={}) + return OutputBatch(output={}, timings=self._make_timings()) from nixl_transfer import NIXL_AVAILABLE if NIXL_AVAILABLE: @@ -193,7 +205,7 @@ def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: meta = self._sender.send(gpu_tensors) # Include CPU metadata tensors directly (e.g. __count fields) meta["cpu_tensors"] = cpu_tensors - return OutputBatch(output={"_nixl_transfer_meta": meta}) + return OutputBatch(output={"_nixl_transfer_meta": meta}, timings=self._make_timings()) def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: """Fallback: send raw tensors via ZMQ pickle.""" @@ -215,7 +227,7 @@ def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: real_key = base[6:] output[real_key] = [idx_map[i] for i in sorted(idx_map)] del output[base] - return OutputBatch(output=output) + return OutputBatch(output=output, timings=self._make_timings()) # ═══════════════════════════════════════════════════════════════════════ @@ -224,12 +236,12 @@ def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: def build_encoder_stages(pipeline, server_args): - """TextEncodingStage → IntermediateOutputStage(prompt_embeds, …). + """TextEncodingStage → NixlSendStage(prompt_embeds, …). Automatically detects all loaded text encoders/tokenizers so that both single-encoder (Wan) and dual-encoder (HunyuanVideo) models work. """ - from sglang.multimodal_gen.runtime.pipelines.stages.text_encoding import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( TextEncodingStage, ) from sglang_utils import get_component_backend @@ -261,14 +273,14 @@ def build_encoder_stages(pipeline, server_args): def build_denoiser_stages(pipeline, server_args): - """DeviceMove → LatentPrep → TimestepPrep → Denoising → IntermediateOutput.""" - from sglang.multimodal_gen.runtime.pipelines.stages.latent_preparation import ( + """NixlReceive → LatentPrep → TimestepPrep → Denoising → NixlSend.""" + from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines.stages.timestep_preparation import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( TimestepPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines.stages.denoising import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( DenoisingStage, ) from sglang_utils import get_component_backend @@ -288,7 +300,7 @@ def build_denoiser_stages(pipeline, server_args): def build_vae_stages(pipeline, server_args): """NixlReceive → DecodingStage.""" - from sglang.multimodal_gen.runtime.pipelines.stages.decoding import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( DecodingStage, ) from sglang_utils import get_component_backend diff --git a/examples/disagg_diffusion/phase1_workers/protocol.py b/examples/disagg_diffusion/phase1_workers/protocol.py index b4aa6c1a19c6..ce935db9cbef 100644 --- a/examples/disagg_diffusion/phase1_workers/protocol.py +++ b/examples/disagg_diffusion/phase1_workers/protocol.py @@ -1,74 +1,32 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Protocol types for disaggregated diffusion stages. -These Pydantic models define the request/response contracts between -Encoder, Denoiser, and VAE workers. Tensor payloads are serialized -as base64-encoded torch.save bytes. For production use, replace with -shared-memory or object-store references. +Uses NIXL RDMA for GPU-direct tensor transfer between stage workers. +Only small metadata (shapes, dtypes, NIXL descriptor) travels over Dynamo RPC. +Tensor transfer is handled by NixlSendStage/NixlReceiveStage in +partial_gpu_worker.py and nixl_transfer.py. """ -import base64 -import io from typing import Any, Dict, List, Optional -import torch from pydantic import BaseModel -# --------------------------------------------------------------------------- -# Tensor serialization helpers -# --------------------------------------------------------------------------- - -def tensor_to_b64(t: torch.Tensor) -> str: - buf = io.BytesIO() - torch.save(t.cpu(), buf) - return base64.b64encode(buf.getvalue()).decode("ascii") - - -def b64_to_tensor(s: str, device: str = "cuda") -> torch.Tensor: - raw = base64.b64decode(s) - return torch.load(io.BytesIO(raw), weights_only=True).to(device) - - -def tensors_to_b64(tensors: Dict[str, torch.Tensor]) -> str: - buf = io.BytesIO() - cpu_tensors = {k: v.cpu() for k, v in tensors.items()} - torch.save(cpu_tensors, buf) - return base64.b64encode(buf.getvalue()).decode("ascii") - - -def b64_to_tensors(s: str, device: str = "cuda") -> Dict[str, torch.Tensor]: - raw = base64.b64decode(s) - data = torch.load(io.BytesIO(raw), weights_only=True) - return {k: v.to(device) for k, v in data.items()} - - # --------------------------------------------------------------------------- # Stage 1: Encoder # --------------------------------------------------------------------------- class EncoderRequest(BaseModel): prompt: str - model: str = "black-forest-labs/FLUX.1-schnell" + negative_prompt: str = "" + guidance_scale: float = 1.0 class EncoderResponse(BaseModel): - """Serialized text embeddings.""" - embeddings_b64: str # base64(torch.save({prompt_embeds, pooled_prompt_embeds, text_ids})) - shapes: Dict[str, List[int]] + transfer_meta: Dict[str, Any] = {} + shapes: Dict[str, List[int]] = {} # --------------------------------------------------------------------------- @@ -76,19 +34,18 @@ class EncoderResponse(BaseModel): # --------------------------------------------------------------------------- class DenoiserRequest(BaseModel): - embeddings_b64: str - model: str = "black-forest-labs/FLUX.1-schnell" - height: int = 512 - width: int = 512 - num_inference_steps: int = 4 - guidance_scale: float = 0.0 + transfer_meta: Dict[str, Any] + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_inference_steps: int = 50 + guidance_scale: float = 1.0 seed: int = 42 class DenoiserResponse(BaseModel): - """Serialized denoised latents.""" - latents_b64: str # base64(torch.save({latents, scaling_factor, shift_factor})) - shape: List[int] + transfer_meta: Dict[str, Any] = {} + shape: List[int] = [] # --------------------------------------------------------------------------- @@ -96,32 +53,39 @@ class DenoiserResponse(BaseModel): # --------------------------------------------------------------------------- class VAEDecodeRequest(BaseModel): - latents_b64: str - model: str = "black-forest-labs/FLUX.1-schnell" + transfer_meta: Dict[str, Any] + request_id: str = "" class VAEDecodeResponse(BaseModel): - """Final output image.""" - image_b64: Optional[str] = None # base64 PNG - url: Optional[str] = None + video_path: str = "" + num_frames: int = 0 # --------------------------------------------------------------------------- -# End-to-end (for orchestrator convenience) +# End-to-end (orchestrator convenience) # --------------------------------------------------------------------------- class GenerateRequest(BaseModel): prompt: str - model: str = "black-forest-labs/FLUX.1-schnell" - height: int = 512 - width: int = 512 - num_inference_steps: int = 4 - guidance_scale: float = 0.0 + negative_prompt: str = "" + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_inference_steps: int = 50 + guidance_scale: float = 1.0 seed: int = 42 - response_format: str = "b64_json" -class GenerateResponse(BaseModel): - image_b64: Optional[str] = None - url: Optional[str] = None - timings: Optional[Dict[str, float]] = None +# --------------------------------------------------------------------------- +# Health (per-stage) +# --------------------------------------------------------------------------- + +class HealthRequest(BaseModel): + pass + + +class HealthResponse(BaseModel): + status: str = "ok" + stage: str = "" + model: str = "" diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index 92d047333a7d..6fb17b605bbc 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -85,9 +85,9 @@ def _patch_hunyuan_config_task_type(): default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply ``task_type=T2V`` when omitted. Idempotent. """ - from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType + from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType try: - from sglang.multimodal_gen.configs.pipelines.hunyuan import ( + from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( HunyuanConfig, FastHunyuanConfig, ) except ImportError: @@ -401,9 +401,9 @@ async def main(): logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) # ── Connect clients ────────────────────────────────────────── - enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") - den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") - vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") + enc_client = StageClient(enc_args.scheduler_endpoint, "encoder") + den_client = StageClient(den_args.scheduler_endpoint, "denoiser") + vae_client = StageClient(vae_args.scheduler_endpoint, "vae") # ── Warmup ─────────────────────────────────────────────────── logger.info("Warmup request …") diff --git a/examples/disagg_diffusion/phase1_workers/serve.py b/examples/disagg_diffusion/phase1_workers/serve.py deleted file mode 100644 index 35b55ffe984d..000000000000 --- a/examples/disagg_diffusion/phase1_workers/serve.py +++ /dev/null @@ -1,324 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""HTTP serving mode for disaggregated diffusion pipeline. - -Launches 3 stage worker groups (Encoder, Denoiser, VAE) on separate GPUs, -then serves video generation requests via HTTP. - -Process architecture:: - - serve.py (main) - ├── Encoder worker subprocess (GPU 0) ← SGLang Scheduler + Llama + CLIP - ├── Denoiser worker subprocess 0 (GPU 1) ┐ - ├── Denoiser worker subprocess 1 (GPU 2) ┘── SGLang Scheduler + DiT (TP=2) - ├── VAE worker subprocess (GPU 3) ← SGLang Scheduler + 3D VAE - └── FastAPI HTTP server (no GPU) ← orchestrates stages via ZMQ - -Inter-stage tensor data transfers use NIXL RDMA (GPU-direct). -Only metadata (~1.5 KB) travels over ZMQ. - -Usage:: - - # Start server (default: 4 GPUs) - python serve.py - - # Custom GPU layout - GPU_ENC=0 GPU_DEN=1,2,3,4 GPU_VAE=5 TP_SIZE=4 python serve.py - - # Send requests - curl http://localhost:8090/health - - curl -X POST http://localhost:8090/generate \\ - -H "Content-Type: application/json" \\ - -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_steps": 3}' \\ - --output output.mp4 - - # Longer video (61 frames ≈ 2.5s at 24fps, ~8 min on H20) - curl -X POST http://localhost:8090/generate \\ - -H "Content-Type: application/json" \\ - -d '{"prompt": "A rocket launching into space", "num_frames": 61, "num_steps": 50}' \\ - --output rocket.mp4 - -Environment variables: - MODEL_PATH HuggingFace model (default: hunyuanvideo-community/HunyuanVideo) - GPU_ENC GPU for encoder (default: 0) - GPU_DEN GPUs for denoiser, comma-separated (default: 1,2) - GPU_VAE GPU for VAE (default: 3) - TP_SIZE Tensor parallelism degree (default: auto from GPU_DEN count) - SERVE_HOST HTTP bind address (default: 0.0.0.0) - SERVE_PORT HTTP port (default: 8090) -""" - -from __future__ import annotations - -import io -import logging -import multiprocessing as mp -import os -import sys -import time -from typing import Optional - -WORKERS_DIR = os.path.dirname(os.path.abspath(__file__)) -if WORKERS_DIR not in sys.path: - sys.path.insert(0, WORKERS_DIR) - -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", -) -logger = logging.getLogger("serve") - -# ── Configuration ──────────────────────────────────────────────────────── -MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") -GPU_ENC = os.environ.get("GPU_ENC", "0") -GPU_DEN = os.environ.get("GPU_DEN", "1,2") -GPU_VAE = os.environ.get("GPU_VAE", "3") -TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(","))))) -SERVE_HOST = os.environ.get("SERVE_HOST", "0.0.0.0") -SERVE_PORT = int(os.environ.get("SERVE_PORT", "8090")) - -# ── Reuse core components ──────────────────────────────────────────────── -from run_e2e_sglang import ( - _patch_hunyuan_config_task_type, - _detect_encoder_modules, - _launch_stage, - StageClient, - terminate_processes, -) - -# Global state -enc_client: Optional[StageClient] = None -den_client: Optional[StageClient] = None -vae_client: Optional[StageClient] = None -_all_procs = [] -_request_counter = 0 - - -# ── Pipeline execution ─────────────────────────────────────────────────── - -async def run_pipeline( - prompt: str, - negative_prompt: str = "", - height: int = 544, - width: int = 960, - num_frames: int = 61, - num_steps: int = 50, - guidance_scale: float = 1.0, - seed: int = 42, -) -> bytes: - """Run full E2E pipeline, return mp4 bytes.""" - global _request_counter - import torch - from sglang_utils import build_req, inject_tensors_to_req - - req_id = _request_counter - _request_counter += 1 - - req_kwargs = dict( - prompt=prompt, negative_prompt=negative_prompt, - height=height, width=width, num_frames=num_frames, - num_inference_steps=num_steps, guidance_scale=guidance_scale, seed=seed, - ) - - t_total = time.monotonic() - - # ── Encoder ── - t0 = time.monotonic() - enc_output = await enc_client.forward([build_req(**req_kwargs)]) - t_enc = time.monotonic() - t0 - if enc_output.error: - raise RuntimeError(f"Encoder error: {enc_output.error}") - enc_result = enc_output.output - nixl_enc = "_nixl_transfer_meta" in enc_result - - # ── Denoiser ── - t0 = time.monotonic() - den_req = build_req(**req_kwargs) - den_req.do_classifier_free_guidance = (guidance_scale > 1.0) - if nixl_enc: - den_req._nixl_transfer_meta = enc_result["_nixl_transfer_meta"] - else: - inject_tensors_to_req(den_req, enc_result) - den_output = await den_client.forward([den_req]) - t_den = time.monotonic() - t0 - if den_output.error: - raise RuntimeError(f"Denoiser error: {den_output.error}") - den_result = den_output.output - nixl_den = "_nixl_transfer_meta" in den_result - - # ── VAE ── - t0 = time.monotonic() - vae_req = build_req(prompt="", height=height, width=width, - num_frames=num_frames, num_inference_steps=num_steps, - guidance_scale=0.0, seed=seed) - if nixl_den: - vae_req._nixl_transfer_meta = den_result["_nixl_transfer_meta"] - else: - vae_req.latents = den_result["latents"].cpu() - vae_output = await vae_client.forward([vae_req]) - t_vae = time.monotonic() - t0 - if vae_output.error: - raise RuntimeError(f"VAE error: {vae_output.error}") - - t_elapsed = time.monotonic() - t_total - - # ── Encode as mp4 ── - import numpy as np - import imageio - - frames_tensor = vae_output.output - if hasattr(frames_tensor, "cpu"): - frames_tensor = frames_tensor.cpu().float().numpy() - frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) - - buf = io.BytesIO() - imageio.mimwrite(buf, frames, format="mp4", fps=24, codec="libx264") - mp4_bytes = buf.getvalue() - - transfer = "NIXL" if (nixl_enc and nixl_den) else "ZMQ" - logger.info( - "req %d | %d frames %dx%d | enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs | %s | %.1f KB", - req_id, frames.shape[0], width, height, - t_enc, t_den, t_vae, t_elapsed, transfer, len(mp4_bytes) / 1024, - ) - return mp4_bytes - - -# ── FastAPI app ────────────────────────────────────────────────────────── - -from fastapi import FastAPI, HTTPException -from fastapi.responses import Response -from pydantic import BaseModel - -app = FastAPI(title="Disaggregated Diffusion Server") - - -class GenerateRequest(BaseModel): - prompt: str = "A cat walking on green grass" - negative_prompt: str = "" - height: int = 544 - width: int = 960 - num_frames: int = 61 - num_steps: int = 50 - guidance_scale: float = 1.0 - seed: int = 42 - - -@app.get("/health") -async def health(): - n_enc = len(GPU_ENC.split(",")) - n_den = len(GPU_DEN.split(",")) - n_vae = len(GPU_VAE.split(",")) - return { - "status": "ok", - "model": MODEL_PATH, - "gpus": { - "encoder": f"GPU {GPU_ENC} ({n_enc} process{'es' if n_enc > 1 else ''})", - "denoiser": f"GPU {GPU_DEN} ({n_den} processes, TP={TP_SIZE})", - "vae": f"GPU {GPU_VAE} ({n_vae} process{'es' if n_vae > 1 else ''})", - }, - "total_gpu_workers": n_enc + n_den + n_vae, - "transfer": "NIXL RDMA", - "requests_served": _request_counter, - } - - -@app.post("/generate") -async def generate(req: GenerateRequest): - try: - mp4_bytes = await run_pipeline( - prompt=req.prompt, negative_prompt=req.negative_prompt, - height=req.height, width=req.width, num_frames=req.num_frames, - num_steps=req.num_steps, guidance_scale=req.guidance_scale, - seed=req.seed, - ) - return Response(content=mp4_bytes, media_type="video/mp4") - except Exception as e: - logger.exception("Generation failed") - raise HTTPException(status_code=500, detail=str(e)) - - -# ── Stage lifecycle ────────────────────────────────────────────────────── - -def launch_all_stages(): - global enc_client, den_client, vae_client, _all_procs - - from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages - - _patch_hunyuan_config_task_type() - - t0 = time.monotonic() - logger.info("=" * 60) - logger.info(" Launching disaggregated diffusion stages") - logger.info(" Model: %s", MODEL_PATH) - logger.info(" Encoder: GPU %s (1 process)", GPU_ENC) - logger.info(" Denoiser: GPU %s (%d processes, TP=%d)", GPU_DEN, len(GPU_DEN.split(",")), TP_SIZE) - logger.info(" VAE: GPU %s (1 process)", GPU_VAE) - logger.info(" Total: %d GPU worker processes", len(GPU_ENC.split(",")) + len(GPU_DEN.split(",")) + len(GPU_VAE.split(","))) - logger.info("=" * 60) - - enc_procs, enc_args = _launch_stage( - "Encoder", GPU_ENC, - required_modules=_detect_encoder_modules(MODEL_PATH), - custom_stages_fn=build_encoder_stages, - tp_size=1, scheduler_port=15600, - ) - den_procs, den_args = _launch_stage( - "Denoiser", GPU_DEN, - required_modules=["transformer", "scheduler"], - custom_stages_fn=build_denoiser_stages, - tp_size=TP_SIZE, scheduler_port=15700, - ) - vae_procs, vae_args = _launch_stage( - "VAE", GPU_VAE, - required_modules=["vae", "scheduler"], - custom_stages_fn=build_vae_stages, - tp_size=1, scheduler_port=15800, - ) - - _all_procs.extend(enc_procs + den_procs + vae_procs) - - enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") - den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") - vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") - - logger.info("=" * 60) - logger.info(" All %d workers ready in %.1fs", len(_all_procs), time.monotonic() - t0) - logger.info(" HTTP server: http://%s:%d", SERVE_HOST, SERVE_PORT) - logger.info(" Transfer: NIXL RDMA (GPU-direct)") - logger.info("") - logger.info(" Try:") - logger.info(" curl http://localhost:%d/health", SERVE_PORT) - logger.info(' curl -X POST http://localhost:%d/generate \\', SERVE_PORT) - logger.info(' -H "Content-Type: application/json" \\') - logger.info(' -d \'{"prompt": "A cat on grass", "num_frames": 9, "num_steps": 3}\' \\') - logger.info(" --output test.mp4") - logger.info("=" * 60) - - -def shutdown_stages(): - for client in [enc_client, den_client, vae_client]: - if client: - try: - client.close() - except Exception: - pass - for p in _all_procs: - p.terminate() - for p in _all_procs: - p.join(timeout=10) - logger.info("All workers terminated.") - - -if __name__ == "__main__": - mp.set_start_method("spawn", force=True) - launch_all_stages() - - import uvicorn - try: - uvicorn.run(app, host=SERVE_HOST, port=SERVE_PORT, log_level="info") - finally: - shutdown_stages() diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py index 099710172b94..bfa2071a8966 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -81,7 +81,7 @@ def build_partial_pipeline( Auto-detects pipeline class from model_index.json, suppresses automatic stage creation, and syncs all component configs (even unloaded ones). """ - from sglang.multimodal_gen.runtime.pipelines import get_model_info + from sglang.multimodal_gen.runtime.pipelines_core import get_model_info model_info = get_model_info(server_args.model_path) base_pipeline_cls = model_info.pipeline_cls @@ -95,7 +95,7 @@ def _noop_create_stages(self, server_args): def _safe_init(self, **kwargs): # Call ComposedPipelineBase.__init__ directly, skipping LoRAPipeline # which tries to access self.modules['transformer'] - from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) ComposedPipelineBase.__init__(self, **kwargs) @@ -248,21 +248,19 @@ def build_config(server_args): def build_req( prompt: str, negative_prompt: Optional[str] = "", - height: int = 480, - width: int = 832, - num_frames: int = 17, - num_inference_steps: int = 20, - guidance_scale: float = 5.0, + height: int = 544, + width: int = 960, + num_frames: int = 61, + num_inference_steps: int = 50, + guidance_scale: float = 1.0, seed: int = 42, device: str = "cuda", **extra_fields, ) -> "Req": """Construct a minimal SGLang ``Req`` for running pipeline stages.""" - from sglang.multimodal_gen.configs.sample.base import DataType - from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req req = Req( - data_type=DataType.VIDEO, prompt=prompt, negative_prompt=negative_prompt, height=height, diff --git a/examples/disagg_diffusion/phase1_workers/vae_worker.py b/examples/disagg_diffusion/phase1_workers/vae_worker.py index c1c8bc5c8c58..e6843cd8a037 100755 --- a/examples/disagg_diffusion/phase1_workers/vae_worker.py +++ b/examples/disagg_diffusion/phase1_workers/vae_worker.py @@ -1,116 +1,179 @@ #!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — VAE Worker - -Loads only the VAE decoder. Accepts denoised latents and returns -the final image as base64-encoded PNG. - -Usage: - python vae_worker.py --model black-forest-labs/FLUX.1-schnell +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — VAE Worker (Dynamo RPC) + +Wraps an SGLang Scheduler subprocess running NixlReceiveStage + DecodingStage. +Receives latents via NIXL RDMA, decodes to video frames, saves as mp4. + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess (spawned by launch_partial_server) + +-- PartialGPUWorker + |-- NixlReceiveStage <-- RDMA-pull latents from denoiser + +-- DecodingStage <-- VAE decode to video frames """ import asyncio -import base64 -import io +import json import logging +import multiprocessing as mp import os import sys +import uuid import numpy as np -import torch import uvloop -from PIL import Image sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from protocol import VAEDecodeRequest, VAEDecodeResponse, b64_to_tensors # noqa: E402 - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 logger = logging.getLogger(__name__) -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") - - -class VAEStage: - """VAE decode stage: latents → image pixels.""" - - def __init__(self): - self.vae = None +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15800")) +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_videos") - def load_model(self): - from diffusers import AutoencoderKL - logger.info("Loading VAE from %s …", MODEL_PATH) - self.vae = AutoencoderKL.from_pretrained( - MODEL_PATH, subfolder="vae", torch_dtype=torch.bfloat16 - ) - self.vae.to(DEVICE) - - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("VAE ready — VRAM: %.0f MB", vram) +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from run_e2e_sglang import ( + _patch_hunyuan_config_task_type, + StageClient, + ) + from partial_gpu_worker import build_vae_stages, launch_partial_server + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from sglang_utils import build_req - @dynamo_endpoint(VAEDecodeRequest, VAEDecodeResponse) - async def generate(self, request: VAEDecodeRequest): - logger.info("Decoding latents …") + _patch_hunyuan_config_task_type() + os.makedirs(OUTPUT_DIR, exist_ok=True) - data = b64_to_tensors(request.latents_b64, DEVICE) - latents = data["latents"] - scaling_factor = data["scaling_factor"].item() - shift_factor = data.get("shift_factor") - if shift_factor is not None: - shift_factor = shift_factor.item() + server_args = ServerArgs.from_kwargs( + model_path=MODEL_PATH, + num_gpus=1, + tp_size=1, + scheduler_port=SCHEDULER_PORT, + ) + set_global_server_args(server_args) - # Undo pipeline's latent scaling - if shift_factor is not None: - latents = latents / scaling_factor + shift_factor - else: - latents = latents / scaling_factor + logger.info("Launching VAE Scheduler: port=%d", SCHEDULER_PORT) + processes = launch_partial_server( + server_args, + required_modules=["vae", "scheduler"], + custom_stages_fn=build_vae_stages, + ) - loop = asyncio.get_event_loop() + # Connect ZMQ client to local Scheduler + client = StageClient(server_args.scheduler_endpoint, "vae") + + # ── Dynamo RPC handlers ────────────────────────────────────────── + + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) + + req = build_req( + prompt="", + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=1, + guidance_scale=0.0, + seed=request.get("seed", 42), + ) + + # Pass NIXL metadata for NixlReceiveStage to RDMA-pull latents + transfer_meta = request.get("transfer_meta", {}) + if transfer_meta: + req._nixl_transfer_meta = transfer_meta + + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "video_path": "", "num_frames": 0} + return + + # Extract frames and save video + frames_tensor = output.output + request_id = request.get("request_id") or str(uuid.uuid4())[:8] + + loop = asyncio.get_event_loop() + filename, n_frames = await loop.run_in_executor( + None, _save_video_frames, frames_tensor, request_id, + ) + logger.info("Decoded — %d frames -> %s", n_frames, filename) + yield {"video_path": filename, "num_frames": n_frames} + + except Exception as e: + logger.error("VAE generate failed: %s", e, exc_info=True) + yield {"error": str(e), "video_path": "", "num_frames": 0} + + async def handle_health(request, context): + yield {"status": "ok", "stage": "vae", "model": MODEL_PATH} + + # ── Serve Dynamo endpoints ─────────────────────────────────────── + + ns = runtime.namespace("disagg_diffusion") + gen_ep = ns.component("vae").endpoint("generate") + health_ep = ns.component("vae").endpoint("health") + + logger.info("Serving: disagg_diffusion.vae.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) - def _decode(): - with torch.no_grad(): - decoded = self.vae.decode(latents, return_dict=False)[0] - decoded = (decoded / 2 + 0.5).clamp(0, 1) - return decoded.cpu().permute(0, 2, 3, 1).float().numpy() - pixels = await loop.run_in_executor(None, _decode) +def _save_video_frames(frames_tensor, request_id: str) -> tuple: + """Save decoded frames as mp4. Returns (filename, num_frames).""" + import torch - img = Image.fromarray((pixels[0] * 255).round().astype(np.uint8)) - buf = io.BytesIO() - img.save(buf, format="PNG") - image_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + if isinstance(frames_tensor, dict): + # OutputBatch may return dict — extract the video tensor + for v in frames_tensor.values(): + if hasattr(v, "shape"): + frames_tensor = v + break - response = VAEDecodeResponse(image_b64=image_b64) - logger.info("Decoded — image %dx%d", img.width, img.height) - yield response.model_dump() + if hasattr(frames_tensor, "cpu"): + frames_tensor = frames_tensor.cpu().float().numpy() + # [B, C, T, H, W] -> [T, H, W, C] + frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.vae.generate") + filename = f"{request_id}.mp4" + filepath = os.path.join(OUTPUT_DIR, filename) - stage = VAEStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) + try: + import imageio + imageio.mimwrite(filepath, frames, fps=24, codec="libx264") + except Exception as e: + logger.warning("mp4 export failed (%s), saving first frame as PNG", e) + from PIL import Image + img = Image.fromarray(frames[0]) + filepath = filepath.replace(".mp4", ".png") + filename = filename.replace(".mp4", ".png") + img.save(filepath) - logger.info("Serving VAE endpoint: disagg_diffusion.vae.generate") - await endpoint.serve_endpoint(stage.generate) + return filename, len(frames) if __name__ == "__main__": @@ -118,5 +181,6 @@ async def worker(runtime: DistributedRuntime): level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s", ) + mp.set_start_method("spawn", force=True) uvloop.install() asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py b/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py index 615ccff5f401..b2a316a5e3a3 100755 --- a/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py +++ b/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py @@ -1,164 +1,339 @@ #!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Phase 2: Disaggregated Diffusion Orchestrator - -Connects to the three stage workers (Encoder, Denoiser, VAE) via Dynamo -RPC and chains them into an end-to-end image generation pipeline. - -This plays the same role as the Frontend/Global Router in the EPD -architecture, but implemented as a lightweight client for the POC. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion Orchestrator — HTTP Server with Pipeline Parallelism + +Persistent server that accepts video generation requests and chains +Encoder → Denoiser → VAE via Dynamo RPC + NIXL RDMA. + +Pipeline parallelism: multiple requests can be in different stages +simultaneously. Per-stage semaphores control backpressure so each +GPU processes one request at a time, while the pipeline stays full. + +Each stage worker wraps an SGLang Scheduler subprocess via +launch_partial_server(), supporting TP for the denoiser and NIXL RDMA +for GPU-direct tensor transfer between stages. Usage: - # Ensure the three workers are already running (see launch/run_all.sh) - python run_disagg.py \\ - --prompt "A photo of a cat sitting on a windowsill" \\ - --output /tmp/disagg_output.png + python run_disagg.py [--port 8080] + +API: + POST /v1/videos/generations + GET /health + GET /health/stages + GET /pipeline/status + GET /videos/ """ import asyncio -import base64 import json import logging import os import sys import time +import uuid +from collections import defaultdict +from typing import Dict import uvloop sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "phase1_workers")) from protocol import ( # noqa: E402 - DenoiserRequest, - EncoderRequest, - VAEDecodeRequest, + DenoiserRequest, EncoderRequest, VAEDecodeRequest, + HealthRequest, ) - from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 logger = logging.getLogger(__name__) -PROMPT = os.environ.get("PROMPT", "A photo of a cat sitting on a windowsill") -MODEL = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -OUTPUT = os.environ.get("OUTPUT", "/tmp/disagg_output.png") -HEIGHT = int(os.environ.get("HEIGHT", "512")) -WIDTH = int(os.environ.get("WIDTH", "512")) -NUM_STEPS = int(os.environ.get("NUM_STEPS", "4")) -SEED = int(os.environ.get("SEED", "42")) +PORT = int(os.environ.get("PORT", "8080")) +HOST = os.environ.get("HOST", "0.0.0.0") +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_videos") +MAX_PIPELINE_DEPTH = int(os.environ.get("MAX_PIPELINE_DEPTH", "4")) async def call_stage(client, request_json: str) -> dict: - """Call a Dynamo endpoint, collect the streamed response. - - Each stage worker yields exactly one response dict. - """ result = None stream = await client.generate(request_json) async for chunk in stream: - # Dynamo streams return objects with a .data() method or raw dicts data = chunk.data() if hasattr(chunk, "data") else chunk if isinstance(data, str): data = json.loads(data) result = data - if result is None: raise RuntimeError("Empty response from stage") return result -@dynamo_worker() +async def query_stage_health(health_client, stage_name: str) -> dict: + """Query a stage's StageEngine health via its Dynamo health endpoint.""" + try: + health_req = HealthRequest() + stream = await health_client.generate(health_req.model_dump_json()) + result = None + async for chunk in stream: + data = chunk.data() if hasattr(chunk, "data") else chunk + if isinstance(data, str): + data = json.loads(data) + result = data + return {"stage": stage_name, **result} if result else {"stage": stage_name, "error": "empty response"} + except Exception as e: + return {"stage": stage_name, "error": str(e)} + + +class PipelineTracker: + """Tracks requests flowing through the 3-stage pipeline.""" + + STAGES = ("encoder", "denoiser", "vae") + + def __init__(self): + self._active: Dict[str, str] = {} + self._completed = 0 + self._failed = 0 + self._stage_times: Dict[str, list] = defaultdict(list) + self._lock = asyncio.Lock() + + async def enter(self, request_id: str, stage: str): + async with self._lock: + self._active[request_id] = stage + + async def leave(self, request_id: str, stage: str, elapsed: float): + async with self._lock: + self._stage_times[stage].append(elapsed) + if request_id in self._active and self._active[request_id] == stage: + if stage == "vae": + del self._active[request_id] + + async def mark_done(self, request_id: str): + async with self._lock: + self._active.pop(request_id, None) + self._completed += 1 + + async def mark_failed(self, request_id: str): + async with self._lock: + self._active.pop(request_id, None) + self._failed += 1 + + async def status(self) -> dict: + async with self._lock: + per_stage = defaultdict(list) + for rid, stage in self._active.items(): + per_stage[stage].append(rid) + avg_times = {} + for stage in self.STAGES: + times = self._stage_times[stage] + avg_times[stage] = round(sum(times) / len(times), 3) if times else 0 + return { + "active_requests": dict(per_stage), + "active_count": len(self._active), + "completed": self._completed, + "failed": self._failed, + "avg_stage_seconds": avg_times, + } + + +@dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - """Orchestrator: chain Encoder → Denoiser → VAE.""" - - # Create clients to the three stage endpoints - encoder_client = await runtime.endpoint( - "disagg_diffusion.encoder.generate" - ).client() - denoiser_client = await runtime.endpoint( - "disagg_diffusion.denoiser.generate" - ).client() - vae_client = await runtime.endpoint( - "disagg_diffusion.vae.generate" - ).client() - - logger.info("Connected to all three stage endpoints") - timings = {} - - # ── Stage 1: Encoder ───────────────────────────────────────────── - logger.info("[1/3] Encoding prompt …") - t0 = time.monotonic() - - encoder_req = EncoderRequest(prompt=PROMPT, model=MODEL) - encoder_resp = await call_stage(encoder_client, encoder_req.model_dump_json()) - - timings["encoder_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs — shapes: %s", timings["encoder_s"], encoder_resp.get("shapes")) - - # ── Stage 2: Denoiser ──────────────────────────────────────────── - logger.info("[2/3] Denoising %dx%d, %d steps …", WIDTH, HEIGHT, NUM_STEPS) - t0 = time.monotonic() - - denoiser_req = DenoiserRequest( - embeddings_b64=encoder_resp["embeddings_b64"], - model=MODEL, - height=HEIGHT, - width=WIDTH, - num_inference_steps=NUM_STEPS, - guidance_scale=0.0, - seed=SEED, - ) - denoiser_resp = await call_stage(denoiser_client, denoiser_req.model_dump_json()) - - timings["denoiser_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs — latent shape: %s", timings["denoiser_s"], denoiser_resp.get("shape")) - - # ── Stage 3: VAE Decode ────────────────────────────────────────── - logger.info("[3/3] VAE decoding …") - t0 = time.monotonic() - - vae_req = VAEDecodeRequest( - latents_b64=denoiser_resp["latents_b64"], - model=MODEL, - ) - vae_resp = await call_stage(vae_client, vae_req.model_dump_json()) - - timings["vae_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs", timings["vae_s"]) - - # ── Save output ────────────────────────────────────────────────── - image_bytes = base64.b64decode(vae_resp["image_b64"]) - with open(OUTPUT, "wb") as f: - f.write(image_bytes) - - timings["total_s"] = sum(timings.values()) - - logger.info("") - logger.info("=" * 50) - logger.info("Pipeline complete!") - logger.info(" Encoder: %.2fs", timings["encoder_s"]) - logger.info(" Denoiser: %.2fs", timings["denoiser_s"]) - logger.info(" VAE: %.2fs", timings["vae_s"]) - logger.info(" Total: %.2fs", timings["total_s"]) - logger.info(" Output: %s", OUTPUT) - logger.info("=" * 50) + ns = runtime.namespace("disagg_diffusion") + encoder_client = await ns.component("encoder").endpoint("generate").client() + denoiser_client = await ns.component("denoiser").endpoint("generate").client() + vae_client = await ns.component("vae").endpoint("generate").client() + + encoder_health_client = await ns.component("encoder").endpoint("health").client() + denoiser_health_client = await ns.component("denoiser").endpoint("health").client() + vae_health_client = await ns.component("vae").endpoint("health").client() + + health_clients = { + "encoder": encoder_health_client, + "denoiser": denoiser_health_client, + "vae": vae_health_client, + } + + logger.info("Waiting for stage workers …") + await encoder_client.wait_for_instances() + await denoiser_client.wait_for_instances() + await vae_client.wait_for_instances() + logger.info("All 3 stage workers connected") + + os.makedirs(OUTPUT_DIR, exist_ok=True) + + stage_sems = { + "encoder": asyncio.Semaphore(1), + "denoiser": asyncio.Semaphore(1), + "vae": asyncio.Semaphore(1), + } + admission = asyncio.Semaphore(MAX_PIPELINE_DEPTH) + tracker = PipelineTracker() + + async def run_stage(name: str, request_id: str, client, request_json: str) -> dict: + """Run a single stage with semaphore gating and tracking.""" + async with stage_sems[name]: + await tracker.enter(request_id, name) + t0 = time.monotonic() + result = await call_stage(client, request_json) + elapsed = time.monotonic() - t0 + await tracker.leave(request_id, name, elapsed) + logger.info("[%s] %s: %.2fs", request_id, name.capitalize(), elapsed) + return result, elapsed + + async def handle_generate(request: dict) -> dict: + request_id = str(uuid.uuid4())[:8] + seed = request.get("seed") or int(time.time()) % 1000000 + timings: Dict[str, float] = {} + + async with admission: + enc_req = EncoderRequest( + prompt=request["prompt"], + negative_prompt=request.get("negative_prompt", ""), + guidance_scale=request.get("guidance_scale", 1.0), + ) + enc_resp, timings["encoder_s"] = await run_stage( + "encoder", request_id, encoder_client, enc_req.model_dump_json(), + ) + + den_req = DenoiserRequest( + transfer_meta=enc_resp["transfer_meta"], + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=request.get("num_inference_steps", 50), + guidance_scale=request.get("guidance_scale", 1.0), + seed=seed, + ) + den_resp, timings["denoiser_s"] = await run_stage( + "denoiser", request_id, denoiser_client, den_req.model_dump_json(), + ) + + vae_req = VAEDecodeRequest( + transfer_meta=den_resp["transfer_meta"], + request_id=request_id, + ) + vae_resp, timings["vae_s"] = await run_stage( + "vae", request_id, vae_client, vae_req.model_dump_json(), + ) + + timings["total_s"] = round(sum(timings.values()), 3) + await tracker.mark_done(request_id) + logger.info("[%s] Total: %.2fs", request_id, timings["total_s"]) + + filename = vae_resp["video_path"] + resp_format = request.get("response_format", "url") + + if resp_format == "url": + data = [{"url": f"/videos/{filename}"}] + else: + import base64 + filepath = os.path.join(OUTPUT_DIR, filename) + with open(filepath, "rb") as f: + data = [{"b64_json": base64.b64encode(f.read()).decode("ascii")}] + + return { + "id": f"video-{request_id}", + "created": int(time.time()), + "data": data, + "timings": timings, + } + + from aiohttp import web + + async def handle_post(http_request: web.Request) -> web.Response: + try: + body = await http_request.json() + if "prompt" not in body: + return web.json_response({"error": "missing 'prompt' field"}, status=400) + result = await handle_generate(body) + return web.json_response(result) + except Exception as e: + logger.error("Request failed: %s", e, exc_info=True) + return web.json_response({"error": str(e)}, status=500) + + async def handle_health(http_request: web.Request) -> web.Response: + return web.json_response({"status": "ok"}) + + async def handle_stages_health(http_request: web.Request) -> web.Response: + results = await asyncio.gather( + *[ + query_stage_health(client, name) + for name, client in health_clients.items() + ] + ) + return web.json_response({"stages": list(results)}) + + async def handle_pipeline_status(http_request: web.Request) -> web.Response: + status = await tracker.status() + return web.json_response(status) + + async def handle_video(http_request: web.Request) -> web.Response: + filename = http_request.match_info["filename"] + if "/" in filename or "\\" in filename or ".." in filename: + return web.json_response({"error": "invalid filename"}, status=400) + filepath = os.path.join(OUTPUT_DIR, filename) + if not os.path.exists(filepath): + return web.json_response({"error": "not found"}, status=404) + return web.FileResponse(filepath, headers={"Content-Type": "video/mp4"}) + + async def handle_index(http_request: web.Request) -> web.Response: + files = sorted( + [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".mp4")], + key=lambda f: os.path.getmtime(os.path.join(OUTPUT_DIR, f)), + reverse=True, + ) if os.path.isdir(OUTPUT_DIR) else [] + latest = files[0] if files else None + video_tag = ( + f'' + if latest else "

No videos generated yet.

" + ) + history = "".join( + f'
  • {f}
    • ' for f in files[:20] + ) + html = ( + "Disagg Diffusion" + "" + f"

    Latest Video

    {video_tag}" + f"

    History

      {history}
    " + ) + return web.Response(text=html, content_type="text/html") + + http_app = web.Application() + http_app.router.add_get("/", handle_index) + http_app.router.add_post("/v1/videos/generations", handle_post) + http_app.router.add_get("/health", handle_health) + http_app.router.add_get("/health/stages", handle_stages_health) + http_app.router.add_get("/pipeline/status", handle_pipeline_status) + http_app.router.add_get("/videos/{filename}", handle_video) + + runner = web.AppRunner(http_app) + await runner.setup() + + bound_port = PORT + for attempt in range(10): + try: + site = web.TCPSite(runner, HOST, bound_port, reuse_address=True) + await site.start() + break + except OSError as e: + if e.errno == 98 and attempt < 9: + logger.warning("Port %d in use, trying %d", bound_port, bound_port + 1) + bound_port += 1 + else: + raise + + logger.info("Server listening on http://%s:%d (pipeline depth=%d)", HOST, bound_port, MAX_PIPELINE_DEPTH) + logger.info(" GET / <- latest video preview") + logger.info(" POST /v1/videos/generations") + logger.info(" GET /health") + logger.info(" GET /health/stages") + logger.info(" GET /pipeline/status") + logger.info(" GET /videos/.mp4") + + await asyncio.Event().wait() if __name__ == "__main__": - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", - ) + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s") uvloop.install() asyncio.run(worker()) From d4f148a6a119736361735a805790785814fa5cf0 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 22:21:28 +0800 Subject: [PATCH 06/22] docs: update README with high-quality default examples Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 9d5a1b75e0c1..33d5e966bb7f 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -51,7 +51,14 @@ CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py python phase2_orchestrator/run_disagg.py ``` -Test: +Test (high quality, 61 frames / ~2.5s video, 50 denoising steps): +```bash +curl -X POST http://localhost:8080/v1/videos/generations \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing in the background"}' +``` + +Quick test (lower quality, faster): ```bash curl -X POST http://localhost:8080/v1/videos/generations \ -H "Content-Type: application/json" \ @@ -60,12 +67,22 @@ curl -X POST http://localhost:8080/v1/videos/generations \ ### Standalone E2E (no Dynamo, no etcd) -Single-process launcher that starts all 3 stages and runs the pipeline: +Single-process launcher that starts all 3 stages and runs the pipeline. +Default: 61 frames, 50 steps, 544x960 resolution (~8 min on 4x H20 GPUs): ```bash +# High quality (default: 61 frames, 50 steps, 544x960) python phase1_workers/run_e2e_sglang.py + +# Custom prompt +PROMPT="A golden retriever running on a sunny beach" python phase1_workers/run_e2e_sglang.py + +# Quick test (faster, lower quality) +NUM_FRAMES=9 NUM_STEPS=3 python phase1_workers/run_e2e_sglang.py ``` +Output videos are saved to `/tmp/disagg_e2e/output_0.mp4`. + ## Phases ### Phase 0: Offline Validation (no Dynamo) @@ -73,6 +90,14 @@ python phase1_workers/run_e2e_sglang.py Single-GPU script proving diffusers supports split execution. ```bash +# High quality +python phase0_validate/validate_split.py \ + --model hunyuanvideo-community/HunyuanVideo \ + --prompt "A golden retriever running on a sunny beach" \ + --num-steps 30 --num-frames 61 \ + --output-dir /tmp/disagg_validate + +# Quick test python phase0_validate/validate_split.py \ --model hunyuanvideo-community/HunyuanVideo \ --prompt "A cat walking on grass" \ From 780598ba4d6d71cb7f1ccd782e1a8051a02df6a4 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 22:27:33 +0800 Subject: [PATCH 07/22] docs: rewrite README with architecture diagrams, flow, and roadmap Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 211 ++++++++++++++-------------- 1 file changed, 109 insertions(+), 102 deletions(-) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 33d5e966bb7f..8de31322683b 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,158 +1,165 @@ # Disaggregated Diffusion Inference (HunyuanVideo) -Split a monolithic video diffusion pipeline (Text Encoder -> Transformer -> VAE) into -independent stages on separate GPUs. Tensor data transfers between stages use -**NIXL RDMA** (GPU-direct); only small metadata travels over Dynamo RPC. +Split a monolithic video diffusion pipeline into independent stages on separate GPUs. +Tensor data transfers between stages use **NIXL RDMA** (GPU-direct); only small metadata +travels over the control plane. Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. ## Architecture ``` - Dynamo RPC (metadata only) - ┌────────────┬────────────────────────┐ - │ │ │ - ▼ ▼ ▼ -GPU 0: Encoder Worker GPU 1,2: Denoiser Worker GPU 3: VAE Worker - │ (Llama + CLIP) │ (DiT, TP=2) │ (3D VAE) - │ │ │ - └── NIXL RDMA ──────────┘── NIXL RDMA ──────────────┘ - (embeddings) (latents) - ▲ - │ - Orchestrator (HTTP, no GPU) +┌─────────────────────────────────────────────────────────────────────┐ +│ Orchestrator (HTTP API) │ +│ run_disagg.py / run_e2e_sglang.py │ +└──────┬──────────────────────┬───────────────────────────┬───────────┘ + │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC + │ (metadata) │ (metadata) │ (metadata) + ▼ ▼ ▼ +┌──────────────┐ ┌─────────────────────┐ ┌──────────────────────┐ +│ Encoder │ │ Denoiser │ │ VAE Decoder │ +│ Worker │ │ Worker │ │ Worker │ +│ │ │ │ │ │ +│ GPU 0 │ │ GPU 1,2 (TP=2) │ │ GPU 3 │ +│ ~18 GB VRAM │ │ ~24 GB VRAM/GPU │ │ ~6 GB VRAM │ +│ │ │ │ │ │ +│ Llama 8B │ │ HunyuanVideo DiT │ │ 3D VAE │ +│ + CLIP │ │ 9.5B params │ │ │ +└──────┬───────┘ └──────┬──────────────┘ └──────┬───────────────┘ + │ │ │ + └──── NIXL RDMA ──►└──── NIXL RDMA ─────────►│ + (embeddings, (latents, + GPU-direct) GPU-direct) ``` -Each worker is a Dynamo `@dynamo_worker` that: -1. Spawns SGLang Scheduler subprocess(es) via `launch_partial_server()` -2. Connects a ZMQ `StageClient` to its local Scheduler -3. Exposes `serve_endpoint("generate")` for Dynamo RPC -4. Bridges RPC requests to the Scheduler, which runs the model + NIXL transfer +### Request Flow -## Quick Start +``` +1. User ──POST /v1/videos/generations──► Orchestrator +2. Orchestrator ──EncoderRequest──► Encoder Worker +3. Encoder: TextEncoding → NixlSendStage (register embeddings on GPU) +4. Encoder ──{nixl_metadata}──► Orchestrator +5. Orchestrator ──DenoiserRequest + nixl_meta──► Denoiser Worker +6. Denoiser: NixlReceive (RDMA pull embeddings) → LatentPrep → Denoise (50 steps) → NixlSend +7. Denoiser ──{nixl_metadata}──► Orchestrator +8. Orchestrator ──VAERequest + nixl_meta──► VAE Worker +9. VAE: NixlReceive (RDMA pull latents) → Decode → Save MP4 +10. VAE ──{video_path}──► Orchestrator ──► User +``` -```bash -conda activate omni -export HF_HUB_CACHE=/path/to/huggingface/hub +### Worker Internal Architecture -# Terminal 0: etcd (service discovery) -etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 +Each worker wraps an SGLang Scheduler subprocess: -# Terminal 1: Encoder Worker (Llama 8B + CLIP, ~18 GB VRAM) -CUDA_VISIBLE_DEVICES=0 python phase1_workers/encoder_worker.py +``` +Dynamo Worker Process (e.g. encoder_worker.py) +├── @dynamo_worker +│ └── serve_endpoint("generate") ← Dynamo RPC from orchestrator +│ └── StageClient.forward() ← ZMQ to local Scheduler +│ +└── SGLang Scheduler subprocess ← spawned by launch_partial_server() + └── PartialGPUWorker + ├── TextEncodingStage ← model inference + └── NixlSendStage ← register tensors for RDMA +``` -# Terminal 2: Denoiser Worker (DiT TP=2, ~24 GB VRAM per GPU) -CUDA_VISIBLE_DEVICES=1,2 python phase1_workers/denoiser_worker.py +### Pipeline Parallelism -# Terminal 3: VAE Worker (~6 GB VRAM) -CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py +Multiple requests overlap across stages: -# Terminal 4: Orchestrator (no GPU) -python phase2_orchestrator/run_disagg.py ``` - -Test (high quality, 61 frames / ~2.5s video, 50 denoising steps): -```bash -curl -X POST http://localhost:8080/v1/videos/generations \ - -H "Content-Type: application/json" \ - -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing in the background"}' +Request 1: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ VAE ] +Request 2: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ VAE ] +Request 3: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ``` -Quick test (lower quality, faster): -```bash -curl -X POST http://localhost:8080/v1/videos/generations \ - -H "Content-Type: application/json" \ - -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_inference_steps": 3}' -``` +## Quick Start — Single Script (Recommended) -### Standalone E2E (no Dynamo, no etcd) - -Single-process launcher that starts all 3 stages and runs the pipeline. -Default: 61 frames, 50 steps, 544x960 resolution (~8 min on 4x H20 GPUs): +Launches all 3 stages + runs the full pipeline in one command. No etcd needed. ```bash -# High quality (default: 61 frames, 50 steps, 544x960) +conda activate omni +export HF_HUB_CACHE=/path/to/huggingface/hub + +# Default: 61 frames, 50 steps, 544x960, ~8 min on 4x H20 GPUs python phase1_workers/run_e2e_sglang.py # Custom prompt PROMPT="A golden retriever running on a sunny beach" python phase1_workers/run_e2e_sglang.py -# Quick test (faster, lower quality) +# Quick smoke test (~30s) NUM_FRAMES=9 NUM_STEPS=3 python phase1_workers/run_e2e_sglang.py ``` -Output videos are saved to `/tmp/disagg_e2e/output_0.mp4`. - -## Phases +Output: `/tmp/disagg_e2e/output_0.mp4` -### Phase 0: Offline Validation (no Dynamo) +## Multi-Process Deployment (Dynamo RPC) -Single-GPU script proving diffusers supports split execution. +For production use with independent workers and HTTP API: ```bash -# High quality -python phase0_validate/validate_split.py \ - --model hunyuanvideo-community/HunyuanVideo \ - --prompt "A golden retriever running on a sunny beach" \ - --num-steps 30 --num-frames 61 \ - --output-dir /tmp/disagg_validate - -# Quick test -python phase0_validate/validate_split.py \ - --model hunyuanvideo-community/HunyuanVideo \ - --prompt "A cat walking on grass" \ - --num-steps 3 --num-frames 9 \ - --output-dir /tmp/disagg_validate +conda activate omni +export HF_HUB_CACHE=/path/to/huggingface/hub + +# Terminal 0: etcd +etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 + +# Terminal 1-3: Workers +CUDA_VISIBLE_DEVICES=0 python phase1_workers/encoder_worker.py +CUDA_VISIBLE_DEVICES=1,2 python phase1_workers/denoiser_worker.py +CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py + +# Terminal 4: HTTP Orchestrator +python phase2_orchestrator/run_disagg.py ``` -### Phase 1: Dynamo Stage Workers (NIXL + SGLang Scheduler) +```bash +# Generate video (61 frames, 50 steps by default) +curl -X POST http://localhost:8080/v1/videos/generations \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing"}' +``` -Three Dynamo workers, each wrapping an SGLang Scheduler subprocess: +## Workers -| Worker | Model Component | VRAM | Endpoint | -|--------|----------------|------|----------| +| Worker | Model Component | VRAM | Dynamo Endpoint | +|--------|----------------|------|-----------------| | `encoder_worker.py` | Llama 8B + CLIP text encoders | ~18 GB | `disagg_diffusion.encoder.generate` | | `denoiser_worker.py` | HunyuanVideo DiT (TP=2) | ~24 GB/GPU | `disagg_diffusion.denoiser.generate` | | `vae_worker.py` | 3D VAE decoder | ~6 GB | `disagg_diffusion.vae.generate` | -### Phase 2: Orchestrator (Pipeline Parallel) - -Chains three stage endpoints with pipeline parallelism. Multiple concurrent -requests overlap across stages: - -``` -Request 1: [Encoder] → [Denoiser] → [ VAE ] -Request 2: [Encoder] → [Denoiser] → [ VAE ] -Request 3: [Encoder] → [Denoiser] → ... -``` - ## Environment Variables -### Worker Configuration - | Variable | Default | Description | |----------|---------|-------------| -| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model path | -| `SCHEDULER_PORT` | `15600/15700/15800` | ZMQ port for local Scheduler | -| `TP_SIZE` | auto from `CUDA_VISIBLE_DEVICES` | Tensor parallelism (denoiser) | -| `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory (VAE) | +| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | +| `NUM_FRAMES` | `61` | Number of video frames (~2.5s at 24fps) | +| `NUM_STEPS` | `50` | Denoising steps (more = higher quality) | +| `HEIGHT` / `WIDTH` | `544` / `960` | Output resolution | +| `GUIDANCE` | `1.0` | Guidance scale (1.0 = embedded guidance) | +| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment | +| `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism for denoiser | +| `OUTPUT_DIR` | `/tmp/disagg_e2e` | Video output directory | -### Orchestrator Configuration +## Supported Models -| Variable | Default | Description | -|----------|---------|-------------| -| `PORT` | `8080` | HTTP server port | -| `MAX_PIPELINE_DEPTH` | `4` | Max concurrent requests in pipeline | -| `OUTPUT_DIR` | `/tmp/disagg_videos` | Shared directory for video output | +- **`hunyuanvideo-community/HunyuanVideo`** — 13B, dual encoder (Llama 8B + CLIP), recommended +- `Wan-AI/Wan2.2-TI2V-5B-Diffusers` — 5B, single encoder -## Supported Models +## Roadmap -- `hunyuanvideo-community/HunyuanVideo` (13B, dual encoder, recommended) -- `Wan-AI/Wan2.2-TI2V-5B-Diffusers` (5B, single encoder) -- Any SGLang-supported diffusion model with encoder/denoiser/VAE stages +- [ ] **Dynamo RPC data plane** — replace NIXL with Dynamo's native tensor transport +- [ ] **Multi-node** — distribute stages across machines (currently single-node only) +- [ ] **Dynamic batching** — batch multiple prompts per denoiser pass +- [ ] **LoRA hot-swap** — switch LoRA adapters without restarting workers +- [ ] **Speculative decoding** — use smaller DiT for early steps, full DiT for final steps +- [ ] **Streaming output** — stream decoded frames as they're produced +- [ ] **HunyuanVideo 1.5** — upgrade to latest HunyuanVideo with improved quality +- [ ] **Wan2.2 14B** — support larger Wan model with TP +- [ ] **Profiling dashboard** — per-stage latency, GPU utilization, NIXL throughput metrics ## Dependencies ```bash -pip install ai-dynamo-runtime sglang imageio imageio-ffmpeg pyzmq setproctitle etcd-distro +pip install ai-dynamo-runtime sglang imageio imageio-ffmpeg pyzmq setproctitle ``` From 9afd1dcfeb100fea0271b5f33b7ccacf1598cbb7 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 22:32:27 +0800 Subject: [PATCH 08/22] feat: add run_all.sh launcher script, rewrite README - run_all.sh: one-command launch of etcd + 3 workers + orchestrator - README: remove phase terminology, add architecture/flow diagrams, roadmap Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 81 +++++++-------- examples/disagg_diffusion/run_all.sh | 144 +++++++++++++++++++++++++++ 2 files changed, 178 insertions(+), 47 deletions(-) create mode 100755 examples/disagg_diffusion/run_all.sh diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 8de31322683b..7332b9865624 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,8 +1,8 @@ # Disaggregated Diffusion Inference (HunyuanVideo) -Split a monolithic video diffusion pipeline into independent stages on separate GPUs. +Split a monolithic video diffusion pipeline into independent Dynamo workers on separate GPUs. Tensor data transfers between stages use **NIXL RDMA** (GPU-direct); only small metadata -travels over the control plane. +travels over Dynamo RPC. Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. @@ -10,8 +10,8 @@ Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. ``` ┌─────────────────────────────────────────────────────────────────────┐ -│ Orchestrator (HTTP API) │ -│ run_disagg.py / run_e2e_sglang.py │ +│ Orchestrator (HTTP API) │ +│ run_disagg.py │ └──────┬──────────────────────┬───────────────────────────┬───────────┘ │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC │ (metadata) │ (metadata) │ (metadata) @@ -35,16 +35,16 @@ Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. ### Request Flow ``` -1. User ──POST /v1/videos/generations──► Orchestrator -2. Orchestrator ──EncoderRequest──► Encoder Worker -3. Encoder: TextEncoding → NixlSendStage (register embeddings on GPU) -4. Encoder ──{nixl_metadata}──► Orchestrator -5. Orchestrator ──DenoiserRequest + nixl_meta──► Denoiser Worker -6. Denoiser: NixlReceive (RDMA pull embeddings) → LatentPrep → Denoise (50 steps) → NixlSend -7. Denoiser ──{nixl_metadata}──► Orchestrator -8. Orchestrator ──VAERequest + nixl_meta──► VAE Worker -9. VAE: NixlReceive (RDMA pull latents) → Decode → Save MP4 -10. VAE ──{video_path}──► Orchestrator ──► User + 1. User ──POST /v1/videos/generations──► Orchestrator + 2. Orchestrator ──EncoderRequest──► Encoder Worker + 3. Encoder: TextEncoding → NixlSendStage (register embeddings on GPU) + 4. Encoder ──{nixl_metadata}──► Orchestrator + 5. Orchestrator ──DenoiserRequest + nixl_meta──► Denoiser Worker + 6. Denoiser: NixlReceive (RDMA pull embeddings) → LatentPrep → Denoise (N steps) → NixlSend + 7. Denoiser ──{nixl_metadata}──► Orchestrator + 8. Orchestrator ──VAERequest + nixl_meta──► VAE Worker + 9. VAE: NixlReceive (RDMA pull latents) → Decode → Save MP4 +10. VAE ──{video_path}──► Orchestrator ──► User ``` ### Worker Internal Architecture @@ -73,34 +73,27 @@ Request 2: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ Request 3: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ``` -## Quick Start — Single Script (Recommended) +## Quick Start -Launches all 3 stages + runs the full pipeline in one command. No etcd needed. +One script launches everything (etcd + 3 workers + orchestrator): ```bash conda activate omni export HF_HUB_CACHE=/path/to/huggingface/hub -# Default: 61 frames, 50 steps, 544x960, ~8 min on 4x H20 GPUs -python phase1_workers/run_e2e_sglang.py +# Launch all services + send a test request +./run_all.sh --test -# Custom prompt -PROMPT="A golden retriever running on a sunny beach" python phase1_workers/run_e2e_sglang.py +# Quick smoke test (9 frames, 3 steps, ~30s) +./run_all.sh --test --quick -# Quick smoke test (~30s) -NUM_FRAMES=9 NUM_STEPS=3 python phase1_workers/run_e2e_sglang.py +# Just launch services (no test request) +./run_all.sh ``` -Output: `/tmp/disagg_e2e/output_0.mp4` - -## Multi-Process Deployment (Dynamo RPC) - -For production use with independent workers and HTTP API: +Or launch each service manually: ```bash -conda activate omni -export HF_HUB_CACHE=/path/to/huggingface/hub - # Terminal 0: etcd etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 @@ -109,12 +102,13 @@ CUDA_VISIBLE_DEVICES=0 python phase1_workers/encoder_worker.py CUDA_VISIBLE_DEVICES=1,2 python phase1_workers/denoiser_worker.py CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py -# Terminal 4: HTTP Orchestrator +# Terminal 4: Orchestrator python phase2_orchestrator/run_disagg.py ``` +Generate a video (61 frames, 50 steps, 544x960 by default): + ```bash -# Generate video (61 frames, 50 steps by default) curl -X POST http://localhost:8080/v1/videos/generations \ -H "Content-Type: application/json" \ -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing"}' @@ -133,13 +127,10 @@ curl -X POST http://localhost:8080/v1/videos/generations \ | Variable | Default | Description | |----------|---------|-------------| | `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | -| `NUM_FRAMES` | `61` | Number of video frames (~2.5s at 24fps) | -| `NUM_STEPS` | `50` | Denoising steps (more = higher quality) | -| `HEIGHT` / `WIDTH` | `544` / `960` | Output resolution | -| `GUIDANCE` | `1.0` | Guidance scale (1.0 = embedded guidance) | -| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment | +| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage | | `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism for denoiser | -| `OUTPUT_DIR` | `/tmp/disagg_e2e` | Video output directory | +| `PORT` | `8080` | Orchestrator HTTP port | +| `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory | ## Supported Models @@ -148,15 +139,11 @@ curl -X POST http://localhost:8080/v1/videos/generations \ ## Roadmap -- [ ] **Dynamo RPC data plane** — replace NIXL with Dynamo's native tensor transport -- [ ] **Multi-node** — distribute stages across machines (currently single-node only) -- [ ] **Dynamic batching** — batch multiple prompts per denoiser pass -- [ ] **LoRA hot-swap** — switch LoRA adapters without restarting workers -- [ ] **Speculative decoding** — use smaller DiT for early steps, full DiT for final steps -- [ ] **Streaming output** — stream decoded frames as they're produced -- [ ] **HunyuanVideo 1.5** — upgrade to latest HunyuanVideo with improved quality -- [ ] **Wan2.2 14B** — support larger Wan model with TP -- [ ] **Profiling dashboard** — per-stage latency, GPU utilization, NIXL throughput metrics +- [ ] **Dynamic scaling** — auto-scale workers based on queue depth, add/remove denoiser replicas +- [ ] **Streaming output** — stream decoded frames to client as they are produced +- [ ] **Orchestrator improvements** — smarter scheduling, request priority, load balancing across replicas +- [ ] **Metrics & observability** — per-stage latency, GPU utilization, NIXL throughput, Prometheus export +- [ ] **Request cancellation** — cancel in-flight requests, free GPU resources immediately ## Dependencies diff --git a/examples/disagg_diffusion/run_all.sh b/examples/disagg_diffusion/run_all.sh new file mode 100755 index 000000000000..7ee7a1c5f32e --- /dev/null +++ b/examples/disagg_diffusion/run_all.sh @@ -0,0 +1,144 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Launch all disaggregated diffusion services (etcd + 3 workers + orchestrator) +# and optionally send a test request. +# +# Usage: +# ./run_all.sh # launch all services +# ./run_all.sh --test # launch + send a test request +# ./run_all.sh --test --quick # launch + quick smoke test (9 frames, 3 steps) +# +# Environment variables: +# MODEL_PATH HuggingFace model (default: hunyuanvideo-community/HunyuanVideo) +# GPU_ENC GPU for encoder (default: 0) +# GPU_DEN GPUs for denoiser (default: 1,2) +# GPU_VAE GPU for VAE (default: 3) +# PORT HTTP port (default: 8080) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +WORKERS_DIR="$SCRIPT_DIR/phase1_workers" +ORCH_DIR="$SCRIPT_DIR/phase2_orchestrator" +LOG_DIR="/tmp/disagg_logs" + +GPU_ENC="${GPU_ENC:-0}" +GPU_DEN="${GPU_DEN:-1,2}" +GPU_VAE="${GPU_VAE:-3}" +PORT="${PORT:-8080}" + +DO_TEST=false +QUICK=false +for arg in "$@"; do + case "$arg" in + --test) DO_TEST=true ;; + --quick) QUICK=true ;; + esac +done + +mkdir -p "$LOG_DIR" +PIDS=() + +cleanup() { + echo "" + echo "Shutting down all services..." + for pid in "${PIDS[@]}"; do + kill "$pid" 2>/dev/null || true + done + # Kill child processes (sglang schedulers) + for pid in "${PIDS[@]}"; do + pkill -P "$pid" 2>/dev/null || true + done + wait 2>/dev/null + echo "All services stopped." +} +trap cleanup EXIT INT TERM + +echo "==========================================" +echo " Disaggregated Diffusion — Launch All" +echo "==========================================" +echo " Encoder: GPU $GPU_ENC" +echo " Denoiser: GPU $GPU_DEN" +echo " VAE: GPU $GPU_VAE" +echo " HTTP port: $PORT" +echo " Logs: $LOG_DIR/" +echo "==========================================" + +# 1. etcd +if ! pgrep -x etcd > /dev/null 2>&1; then + echo "[1/5] Starting etcd..." + etcd --data-dir /tmp/etcd_disagg \ + --listen-client-urls http://0.0.0.0:2379 \ + --advertise-client-urls http://127.0.0.1:2379 \ + > "$LOG_DIR/etcd.log" 2>&1 & + PIDS+=($!) + sleep 2 +else + echo "[1/5] etcd already running, skipping." +fi + +# 2. Encoder Worker +echo "[2/5] Starting Encoder Worker (GPU $GPU_ENC)..." +CUDA_VISIBLE_DEVICES="$GPU_ENC" python "$WORKERS_DIR/encoder_worker.py" \ + > "$LOG_DIR/encoder.log" 2>&1 & +PIDS+=($!) + +# 3. Denoiser Worker +echo "[3/5] Starting Denoiser Worker (GPU $GPU_DEN)..." +CUDA_VISIBLE_DEVICES="$GPU_DEN" python "$WORKERS_DIR/denoiser_worker.py" \ + > "$LOG_DIR/denoiser.log" 2>&1 & +PIDS+=($!) + +# 4. VAE Worker +echo "[4/5] Starting VAE Worker (GPU $GPU_VAE)..." +CUDA_VISIBLE_DEVICES="$GPU_VAE" python "$WORKERS_DIR/vae_worker.py" \ + > "$LOG_DIR/vae.log" 2>&1 & +PIDS+=($!) + +# 5. Orchestrator +echo "[5/5] Starting Orchestrator (port $PORT)..." +PORT="$PORT" python "$ORCH_DIR/run_disagg.py" \ + > "$LOG_DIR/orchestrator.log" 2>&1 & +PIDS+=($!) + +echo "" +echo "All services launching. Waiting for workers to be ready..." +echo " tail -f $LOG_DIR/encoder.log # monitor encoder" +echo " tail -f $LOG_DIR/denoiser.log # monitor denoiser" +echo " tail -f $LOG_DIR/vae.log # monitor vae" +echo " tail -f $LOG_DIR/orchestrator.log" +echo "" + +# Wait for orchestrator HTTP to be ready +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Orchestrator ready at http://localhost:$PORT" + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Orchestrator not ready after 120s. Check logs in $LOG_DIR/" + exit 1 + fi + sleep 1 +done + +# Test request +if [ "$DO_TEST" = true ]; then + echo "" + echo "Sending test request..." + if [ "$QUICK" = true ]; then + curl -s -X POST "http://localhost:$PORT/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_inference_steps": 3}' | python -m json.tool + else + curl -s -X POST "http://localhost:$PORT/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing in the background"}' | python -m json.tool + fi +fi + +echo "" +echo "Services running. Press Ctrl+C to stop all." +wait From dfb608cafea4679bfe51b8aa1041bc077219b3be Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Mon, 16 Mar 2026 22:47:10 +0800 Subject: [PATCH 09/22] refactor: move shared utilities to sglang_utils, remove dead code - Move StageClient, patch_hunyuan_config, detect_encoder_modules, save_video to sglang_utils.py (was duplicated in run_e2e_sglang.py and vae_worker.py) - Workers now import from sglang_utils instead of run_e2e_sglang.py, avoiding module-level side effects on import - Remove unused build_config() function - Remove duplicate _save_video_frames from vae_worker.py Co-Authored-By: Claude Opus 4.6 --- .../phase1_workers/denoiser_worker.py | 8 +- .../phase1_workers/encoder_worker.py | 11 +- .../phase1_workers/run_e2e_sglang.py | 97 ++-------- .../phase1_workers/sglang_utils.py | 169 ++++++++++++------ .../phase1_workers/vae_worker.py | 49 +---- 5 files changed, 140 insertions(+), 194 deletions(-) diff --git a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py index b214830ebeb3..b451b1f98947 100755 --- a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py +++ b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py @@ -46,17 +46,13 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from run_e2e_sglang import ( - _patch_hunyuan_config_task_type, - StageClient, - ) + from sglang_utils import StageClient, patch_hunyuan_config, build_req from partial_gpu_worker import build_denoiser_stages, launch_partial_server from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, set_global_server_args, ) - from sglang_utils import build_req - _patch_hunyuan_config_task_type() + patch_hunyuan_config() # Auto-detect GPU count from CUDA_VISIBLE_DEVICES num_gpus = len(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")) diff --git a/examples/disagg_diffusion/phase1_workers/encoder_worker.py b/examples/disagg_diffusion/phase1_workers/encoder_worker.py index 7f678f5a6320..3b7d5e3113fe 100755 --- a/examples/disagg_diffusion/phase1_workers/encoder_worker.py +++ b/examples/disagg_diffusion/phase1_workers/encoder_worker.py @@ -44,21 +44,18 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from run_e2e_sglang import ( - _patch_hunyuan_config_task_type, - _detect_encoder_modules, - StageClient, + from sglang_utils import ( + StageClient, patch_hunyuan_config, detect_encoder_modules, build_req, ) from partial_gpu_worker import build_encoder_stages, launch_partial_server from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, set_global_server_args, ) - from sglang_utils import build_req - _patch_hunyuan_config_task_type() + patch_hunyuan_config() # Launch SGLang Scheduler subprocess with text encoder stages - enc_modules = _detect_encoder_modules(MODEL_PATH) + enc_modules = detect_encoder_modules(MODEL_PATH) server_args = ServerArgs.from_kwargs( model_path=MODEL_PATH, num_gpus=1, diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index 6fb17b605bbc..58f12bb5d3b5 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -78,79 +78,16 @@ OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_e2e") -# ── SGLang compatibility patches ───────────────────────────────────────── - -def _patch_hunyuan_config_task_type(): - """HunyuanConfig inherits ``task_type`` from PipelineConfig without a - default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply - ``task_type=T2V`` when omitted. Idempotent. - """ - from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType - try: - from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( - HunyuanConfig, FastHunyuanConfig, - ) - except ImportError: - return - - for cls in (HunyuanConfig, FastHunyuanConfig): - if getattr(cls, "_task_type_patched", False): - continue - orig = cls.__init__ - - def _patched(self, *a, task_type=ModelTaskType.T2V, _orig=orig, **kw): - _orig(self, *a, task_type=task_type, **kw) - - cls.__init__ = _patched - cls._task_type_patched = True - +from sglang_utils import ( # noqa: E402 + StageClient, + patch_hunyuan_config, + detect_encoder_modules, + save_video, +) -def _detect_encoder_modules(model_path: str) -> List[str]: - """Return the required_modules list for the encoder stage.""" - try: - from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - maybe_download_model_index, verify_model_config_and_directory, - ) - config = (verify_model_config_and_directory(model_path) - if os.path.exists(model_path) - else maybe_download_model_index(model_path)) - modules = ["text_encoder", "tokenizer"] - if "text_encoder_2" in config: - modules += ["text_encoder_2", "tokenizer_2"] - modules.append("scheduler") - return modules - except Exception: - pass - # Fallback: include dual encoders for known models - if "hunyuan" in model_path.lower(): - return ["text_encoder", "text_encoder_2", - "tokenizer", "tokenizer_2", "scheduler"] - return ["text_encoder", "tokenizer", "scheduler"] - - -# ── Non-singleton ZMQ client ──────────────────────────────────────────── - -class StageClient: - """Async ZMQ REQ client that talks to a SGLang Scheduler.""" - - def __init__(self, endpoint: str, name: str = ""): - import zmq.asyncio - self._name = name - self._ctx = zmq.asyncio.Context() - self._sock = self._ctx.socket(zmq.REQ) - self._sock.connect(endpoint) - self._lock = asyncio.Lock() - logger.info("StageClient(%s) connected to %s", name, endpoint) - - async def forward(self, reqs): - """Send request(s) and receive response.""" - async with self._lock: - await self._sock.send_pyobj(reqs) - return await self._sock.recv_pyobj() - - def close(self): - self._sock.close() - self._ctx.term() +# Backward-compat aliases for any external code that imports old names +_patch_hunyuan_config_task_type = patch_hunyuan_config +_detect_encoder_modules = detect_encoder_modules # ── Stage launchers ───────────────────────────────────────────────────── @@ -297,17 +234,9 @@ async def run_single_pipeline( def _save_video(frames_tensor, req_id: int): """Save decoded video tensor [B,C,T,H,W] as mp4.""" try: - import numpy as np - import imageio - - os.makedirs(OUTPUT_DIR, exist_ok=True) - if hasattr(frames_tensor, "cpu"): - frames_tensor = frames_tensor.cpu().float().numpy() - # [B, C, T, H, W] -> [T, H, W, C] - frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) out_path = os.path.join(OUTPUT_DIR, f"output_{req_id}.mp4") - imageio.mimwrite(out_path, frames, fps=24, codec="libx264") - logger.info("req %d | Saved %d frames to %s", req_id, frames.shape[0], out_path) + filepath, n_frames = save_video(frames_tensor, out_path) + logger.info("req %d | Saved %d frames to %s", req_id, n_frames, filepath) except Exception as e: logger.warning("req %d | Could not save video: %s", req_id, e) @@ -353,7 +282,7 @@ async def main(): from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages # Apply patches once before any SGLang config is created - _patch_hunyuan_config_task_type() + patch_hunyuan_config() logger.info("=" * 72) logger.info(" Disaggregated Diffusion E2E — SGLang Backend") @@ -376,7 +305,7 @@ async def main(): enc_procs, enc_args = _launch_stage( "Encoder", GPU_ENC, - required_modules=_detect_encoder_modules(MODEL_PATH), + required_modules=detect_encoder_modules(MODEL_PATH), custom_stages_fn=build_encoder_stages, tp_size=1, scheduler_port=15600, diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py index bfa2071a8966..887c71114cd4 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -6,10 +6,14 @@ Provides helpers to construct ServerArgs, load partial pipelines (only the modules each worker needs), and convert between Dynamo protocol types and SGLang's Req dataclass. + +Also contains shared utilities (StageClient, model detection, compatibility +patches) used by both Dynamo workers and the standalone E2E script. """ from __future__ import annotations +import asyncio import logging import os from typing import Any, Dict, List, Optional @@ -191,60 +195,6 @@ def get_component_backend(module) -> str: return f"unknown ({mod}.{cls})" -def build_config(server_args): - """Construct a minimal Config from a diffusion ServerArgs.""" - import types - - try: - from dynamo.sglang.args import Config, DynamoConfig - dynamo_args = DynamoConfig.__new__(DynamoConfig) - except ImportError: - dynamo_args = types.SimpleNamespace() - Config = None - - dynamo_args.component = "disagg_diffusion" - dynamo_args.namespace = "disagg_diffusion" - dynamo_args.diffusion_worker = True - dynamo_args.use_kv_events = False - dynamo_args.media_output_fs_url = "file:///tmp/disagg_videos" - dynamo_args.media_output_http_url = None - dynamo_args.use_sglang_tokenizer = False - dynamo_args.multimodal_processor = False - dynamo_args.multimodal_encode_worker = False - dynamo_args.multimodal_worker = False - dynamo_args.embedding_worker = False - dynamo_args.image_diffusion_worker = False - dynamo_args.video_generation_worker = False - dynamo_args.disagg_config = None - dynamo_args.disagg_config_key = None - dynamo_args.endpoint = "generate" - dynamo_args.discovery_backend = "etcd" - dynamo_args.request_plane = "tcp" - dynamo_args.event_plane = "tcp" - dynamo_args.connector = [] - dynamo_args.enable_local_indexer = False - dynamo_args.durable_kv_events = False - dynamo_args.endpoint_types = "generate" - dynamo_args.dump_config_to = None - dynamo_args.multimodal_embedding_cache_capacity_gb = 0.0 - dynamo_args.output_modalities = ["video"] - dynamo_args.dyn_tool_call_parser = None - dynamo_args.dyn_reasoning_parser = None - dynamo_args.custom_jinja_template = None - - if not hasattr(server_args, "disaggregation_mode"): - server_args.disaggregation_mode = "null" - - if Config is not None: - return Config(server_args, dynamo_args) - - cfg = types.SimpleNamespace() - cfg.server_args = server_args - cfg.dynamo_args = dynamo_args - cfg.serving_mode = getattr(server_args, "disaggregation_mode", "null") - return cfg - - def build_req( prompt: str, negative_prompt: Optional[str] = "", @@ -337,3 +287,114 @@ def inject_tensors_to_req( else: setattr(req, key, value) return req + + +# ═══════════════════════════════════════════════════════════════════════ +# Shared utilities — used by Dynamo workers and the standalone E2E script +# ═══════════════════════════════════════════════════════════════════════ + + +class StageClient: + """Async ZMQ REQ client that talks to a SGLang Scheduler subprocess.""" + + def __init__(self, endpoint: str, name: str = ""): + import zmq.asyncio + self._name = name + self._ctx = zmq.asyncio.Context() + self._sock = self._ctx.socket(zmq.REQ) + self._sock.connect(endpoint) + self._lock = asyncio.Lock() + logger.info("StageClient(%s) connected to %s", name, endpoint) + + async def forward(self, reqs): + """Send request(s) and receive response.""" + async with self._lock: + await self._sock.send_pyobj(reqs) + return await self._sock.recv_pyobj() + + def close(self): + self._sock.close() + self._ctx.term() + + +def patch_hunyuan_config(): + """HunyuanConfig inherits ``task_type`` from PipelineConfig without a + default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply + ``task_type=T2V`` when omitted. Idempotent. + """ + from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType + try: + from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( + HunyuanConfig, FastHunyuanConfig, + ) + except ImportError: + return + + for cls in (HunyuanConfig, FastHunyuanConfig): + if getattr(cls, "_task_type_patched", False): + continue + orig = cls.__init__ + + def _patched(self, *a, task_type=ModelTaskType.T2V, _orig=orig, **kw): + _orig(self, *a, task_type=task_type, **kw) + + cls.__init__ = _patched + cls._task_type_patched = True + + +def detect_encoder_modules(model_path: str) -> List[str]: + """Return the required_modules list for the encoder stage. + + Auto-detects dual-encoder models (e.g. HunyuanVideo with Llama + CLIP). + """ + try: + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + maybe_download_model_index, verify_model_config_and_directory, + ) + config = (verify_model_config_and_directory(model_path) + if os.path.exists(model_path) + else maybe_download_model_index(model_path)) + modules = ["text_encoder", "tokenizer"] + if "text_encoder_2" in config: + modules += ["text_encoder_2", "tokenizer_2"] + modules.append("scheduler") + return modules + except Exception: + pass + # Fallback: include dual encoders for known models + if "hunyuan" in model_path.lower(): + return ["text_encoder", "text_encoder_2", + "tokenizer", "tokenizer_2", "scheduler"] + return ["text_encoder", "tokenizer", "scheduler"] + + +def save_video(frames_tensor, output_path: str, fps: int = 24): + """Save decoded video tensor [B,C,T,H,W] as mp4. + + Returns (filepath, num_frames). + """ + import numpy as np + + if isinstance(frames_tensor, dict): + for v in frames_tensor.values(): + if hasattr(v, "shape"): + frames_tensor = v + break + + if hasattr(frames_tensor, "cpu"): + frames_tensor = frames_tensor.cpu().float().numpy() + + # [B, C, T, H, W] -> [T, H, W, C] + frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) + + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + try: + import imageio + imageio.mimwrite(output_path, frames, fps=fps, codec="libx264") + except Exception as e: + logger.warning("mp4 export failed (%s), saving first frame as PNG", e) + from PIL import Image + output_path = output_path.rsplit(".", 1)[0] + ".png" + Image.fromarray(frames[0]).save(output_path) + + return output_path, len(frames) diff --git a/examples/disagg_diffusion/phase1_workers/vae_worker.py b/examples/disagg_diffusion/phase1_workers/vae_worker.py index e6843cd8a037..c48e12002b8f 100755 --- a/examples/disagg_diffusion/phase1_workers/vae_worker.py +++ b/examples/disagg_diffusion/phase1_workers/vae_worker.py @@ -30,7 +30,6 @@ import sys import uuid -import numpy as np import uvloop sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) @@ -46,17 +45,13 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from run_e2e_sglang import ( - _patch_hunyuan_config_task_type, - StageClient, - ) + from sglang_utils import StageClient, patch_hunyuan_config, build_req, save_video from partial_gpu_worker import build_vae_stages, launch_partial_server from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, set_global_server_args, ) - from sglang_utils import build_req - _patch_hunyuan_config_task_type() + patch_hunyuan_config() os.makedirs(OUTPUT_DIR, exist_ok=True) server_args = ServerArgs.from_kwargs( @@ -107,11 +102,13 @@ async def handle_generate(request, context): # Extract frames and save video frames_tensor = output.output request_id = request.get("request_id") or str(uuid.uuid4())[:8] + out_path = os.path.join(OUTPUT_DIR, f"{request_id}.mp4") loop = asyncio.get_event_loop() - filename, n_frames = await loop.run_in_executor( - None, _save_video_frames, frames_tensor, request_id, + filepath, n_frames = await loop.run_in_executor( + None, save_video, frames_tensor, out_path, ) + filename = os.path.basename(filepath) logger.info("Decoded — %d frames -> %s", n_frames, filename) yield {"video_path": filename, "num_frames": n_frames} @@ -142,40 +139,6 @@ async def handle_health(request, context): p.join(timeout=10) -def _save_video_frames(frames_tensor, request_id: str) -> tuple: - """Save decoded frames as mp4. Returns (filename, num_frames).""" - import torch - - if isinstance(frames_tensor, dict): - # OutputBatch may return dict — extract the video tensor - for v in frames_tensor.values(): - if hasattr(v, "shape"): - frames_tensor = v - break - - if hasattr(frames_tensor, "cpu"): - frames_tensor = frames_tensor.cpu().float().numpy() - - # [B, C, T, H, W] -> [T, H, W, C] - frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) - - filename = f"{request_id}.mp4" - filepath = os.path.join(OUTPUT_DIR, filename) - - try: - import imageio - imageio.mimwrite(filepath, frames, fps=24, codec="libx264") - except Exception as e: - logger.warning("mp4 export failed (%s), saving first frame as PNG", e) - from PIL import Image - img = Image.fromarray(frames[0]) - filepath = filepath.replace(".mp4", ".png") - filename = filename.replace(".mp4", ".png") - img.save(filepath) - - return filename, len(frames) - - if __name__ == "__main__": logging.basicConfig( level=logging.INFO, From 3badde7e0f553fc8b7981a36b8c4384080506756 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 14:36:06 +0800 Subject: [PATCH 10/22] perf: eliminate asyncio.run() overhead, remove dead code, consolidate worker boilerplate Performance: - Replace asyncio.run() with loop.run_until_complete() in NIXL transfer to reuse existing event loop and keep _keep_alive tasks alive - Remove unnecessary __count metadata tensors and CPU/GPU tensor filtering in NixlSendStage - Use batch.logging_info directly instead of creating RequestTimings per call Code quality: - Remove dead build_server_args(), _ensure_distributed_init(), and extract_tensors_from_req() from sglang_utils - Remove backward-compat aliases in run_e2e_sglang.py - Add launch_stage_server() helper consolidating worker boilerplate - Refactor encoder/denoiser/vae workers to use launch_stage_server() - Fix sglang module paths (pipelines_core -> pipelines, etc.) for compatibility with sglang 0.5.5 - Clean up unused imports across all files Verified: E2E test passes with 3 high-quality requests (61 frames, 50 steps, TP=2 denoiser, NIXL RDMA transfers). Co-Authored-By: Claude Opus 4.6 --- .../phase1_workers/denoiser_worker.py | 33 +---- .../phase1_workers/encoder_worker.py | 35 +---- .../phase1_workers/nixl_transfer.py | 6 +- .../phase1_workers/partial_gpu_worker.py | 60 +++----- .../phase1_workers/run_e2e_sglang.py | 10 +- .../phase1_workers/sglang_utils.py | 137 ++++++------------ .../phase1_workers/vae_worker.py | 26 +--- 7 files changed, 90 insertions(+), 217 deletions(-) diff --git a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py index b451b1f98947..67fb0d6b8667 100755 --- a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py +++ b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py @@ -46,39 +46,20 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from sglang_utils import StageClient, patch_hunyuan_config, build_req - from partial_gpu_worker import build_denoiser_stages, launch_partial_server - from sglang.multimodal_gen.runtime.server_args import ( - ServerArgs, set_global_server_args, - ) - - patch_hunyuan_config() + from sglang_utils import launch_stage_server, build_req + from partial_gpu_worker import build_denoiser_stages # Auto-detect GPU count from CUDA_VISIBLE_DEVICES num_gpus = len(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")) tp_size = int(os.environ.get("TP_SIZE", str(num_gpus))) - server_args = ServerArgs.from_kwargs( - model_path=MODEL_PATH, - num_gpus=num_gpus, - tp_size=tp_size, - scheduler_port=SCHEDULER_PORT, - ) - set_global_server_args(server_args) - - logger.info( - "Launching denoiser Scheduler: num_gpus=%d, tp=%d, port=%d", - num_gpus, tp_size, SCHEDULER_PORT, - ) - processes = launch_partial_server( - server_args, - required_modules=["transformer", "scheduler"], - custom_stages_fn=build_denoiser_stages, + logger.info("Launching denoiser Scheduler: num_gpus=%d, tp=%d, port=%d", num_gpus, tp_size, SCHEDULER_PORT) + processes, client, server_args = launch_stage_server( + MODEL_PATH, ["transformer", "scheduler"], build_denoiser_stages, + SCHEDULER_PORT, tp_size=tp_size, num_gpus=num_gpus, + client_name="denoiser", ) - # Connect ZMQ client to local Scheduler - client = StageClient(server_args.scheduler_endpoint, "denoiser") - # ── Dynamo RPC handlers ────────────────────────────────────────── async def handle_generate(request, context): diff --git a/examples/disagg_diffusion/phase1_workers/encoder_worker.py b/examples/disagg_diffusion/phase1_workers/encoder_worker.py index 3b7d5e3113fe..644d35e5e6a3 100755 --- a/examples/disagg_diffusion/phase1_workers/encoder_worker.py +++ b/examples/disagg_diffusion/phase1_workers/encoder_worker.py @@ -44,38 +44,15 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from sglang_utils import ( - StageClient, patch_hunyuan_config, detect_encoder_modules, build_req, - ) - from partial_gpu_worker import build_encoder_stages, launch_partial_server - from sglang.multimodal_gen.runtime.server_args import ( - ServerArgs, set_global_server_args, - ) + from sglang_utils import launch_stage_server, detect_encoder_modules, build_req + from partial_gpu_worker import build_encoder_stages - patch_hunyuan_config() - - # Launch SGLang Scheduler subprocess with text encoder stages enc_modules = detect_encoder_modules(MODEL_PATH) - server_args = ServerArgs.from_kwargs( - model_path=MODEL_PATH, - num_gpus=1, - tp_size=1, - scheduler_port=SCHEDULER_PORT, + logger.info("Launching encoder Scheduler: modules=%s, port=%d", enc_modules, SCHEDULER_PORT) + processes, client, server_args = launch_stage_server( + MODEL_PATH, enc_modules, build_encoder_stages, + SCHEDULER_PORT, client_name="encoder", ) - set_global_server_args(server_args) - - logger.info( - "Launching encoder Scheduler: modules=%s, port=%d", - enc_modules, SCHEDULER_PORT, - ) - processes = launch_partial_server( - server_args, - required_modules=enc_modules, - custom_stages_fn=build_encoder_stages, - ) - - # Connect ZMQ client to local Scheduler - client = StageClient(server_args.scheduler_endpoint, "encoder") # ── Dynamo RPC handlers ────────────────────────────────────────── diff --git a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py index 52bfe9fe404c..3ee77d6e88b7 100644 --- a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py +++ b/examples/disagg_diffusion/phase1_workers/nixl_transfer.py @@ -22,7 +22,7 @@ import asyncio import logging -from typing import Any, Dict, List, Optional +from typing import Dict import torch @@ -61,7 +61,7 @@ def __init__(self): def send(self, tensors: Dict[str, torch.Tensor]) -> dict: """Register tensors and return metadata dict (synchronous wrapper).""" - return asyncio.run(self._async_send(tensors)) + return asyncio.get_event_loop().run_until_complete(self._async_send(tensors)) async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: # Clean completed tasks @@ -99,7 +99,7 @@ class NixlTensorReceiver: def recv(self, meta: dict, device: str = "cuda") -> Dict[str, torch.Tensor]: """Pull tensors described by metadata. Returns {name: tensor}.""" - return asyncio.run(self._async_recv(meta, device)) + return asyncio.get_event_loop().run_until_complete(self._async_recv(meta, device)) async def _async_recv(self, meta: dict, device: str) -> Dict[str, torch.Tensor]: connector = await _PersistentConnector.get() diff --git a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py index 72dac1566830..b677c64a9b8c 100644 --- a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py @@ -57,8 +57,8 @@ get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch -from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req, OutputBatch +from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs # layerwise_offload may not exist in all sglang versions — guard import @@ -102,12 +102,10 @@ def _nixl_pull(self, batch: Req, meta: dict) -> Req: self._receiver = NixlTensorReceiver() tensors = self._receiver.recv(meta, device="cuda") # Reconstruct indexed fields (e.g. prompt_embeds_0, prompt_embeds_1 - # + __prompt_embeds_count → prompt_embeds list) + # → prompt_embeds list) reconstructed = {} indexed = {} # base_name → {idx: tensor} for k, v in tensors.items(): - if k.startswith("__") and k.endswith("_count"): - continue parts = k.rsplit("_", 1) if len(parts) == 2 and parts[1].isdigit(): indexed.setdefault(parts[0], {})[int(parts[1])] = v @@ -153,24 +151,16 @@ def __init__(self, output_fields: List[str]): self._output_fields = output_fields self._sender = None - @staticmethod - def _make_timings(): - """Create a RequestTimings so gpu_worker.execute_forward doesn't crash.""" - try: - from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings - return RequestTimings(request_id="nixl") - except Exception: - return None - def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: tensors = self._extract_tensors(batch) + logging_info = batch.logging_info if not tensors: - return OutputBatch(output={}, timings=self._make_timings()) + return OutputBatch(output={}, logging_info=logging_info) from nixl_transfer import NIXL_AVAILABLE if NIXL_AVAILABLE: - return self._nixl_send(tensors) - return self._fallback_send(tensors) + return self._nixl_send(tensors, logging_info) + return self._fallback_send(tensors, logging_info) def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: """Flatten list-valued fields into individual tensors.""" @@ -188,35 +178,23 @@ def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: # Dual-encoder: store each element separately for NIXL for i, t in enumerate(val): result[f"{field}_{i}"] = t - result[f"__{field}_count"] = torch.tensor(len(val)) elif isinstance(val, torch.Tensor): result[field] = val return result - def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + def _nixl_send(self, tensors: Dict[str, torch.Tensor], logging_info) -> OutputBatch: from nixl_transfer import NixlTensorSender if self._sender is None: self._sender = NixlTensorSender() - # Filter out non-GPU metadata tensors for NIXL - gpu_tensors = {k: v for k, v in tensors.items() - if isinstance(v, torch.Tensor) and v.is_cuda} - cpu_tensors = {k: v for k, v in tensors.items() - if isinstance(v, torch.Tensor) and not v.is_cuda} - meta = self._sender.send(gpu_tensors) - # Include CPU metadata tensors directly (e.g. __count fields) - meta["cpu_tensors"] = cpu_tensors - return OutputBatch(output={"_nixl_transfer_meta": meta}, timings=self._make_timings()) - - def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + meta = self._sender.send(tensors) + return OutputBatch(output={"_nixl_transfer_meta": meta}, logging_info=logging_info) + + def _fallback_send(self, tensors: Dict[str, torch.Tensor], logging_info) -> OutputBatch: """Fallback: send raw tensors via ZMQ pickle.""" # Reconstruct list-valued fields for backward compat output: Dict[str, object] = {} - counts = {} for k, v in tensors.items(): - if k.startswith("__") and k.endswith("_count"): - base = k[2:-6] - counts[base] = int(v.item()) - elif "_" in k and k.rsplit("_", 1)[1].isdigit(): + if "_" in k and k.rsplit("_", 1)[1].isdigit(): base, idx = k.rsplit("_", 1) output.setdefault(f"_list_{base}", {})[int(idx)] = v else: @@ -227,7 +205,7 @@ def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: real_key = base[6:] output[real_key] = [idx_map[i] for i in sorted(idx_map)] del output[base] - return OutputBatch(output=output, timings=self._make_timings()) + return OutputBatch(output=output, logging_info=logging_info) # ═══════════════════════════════════════════════════════════════════════ @@ -241,7 +219,7 @@ def build_encoder_stages(pipeline, server_args): Automatically detects all loaded text encoders/tokenizers so that both single-encoder (Wan) and dual-encoder (HunyuanVideo) models work. """ - from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( + from sglang.multimodal_gen.runtime.pipelines.stages.text_encoding import ( TextEncodingStage, ) from sglang_utils import get_component_backend @@ -274,13 +252,13 @@ def build_encoder_stages(pipeline, server_args): def build_denoiser_stages(pipeline, server_args): """NixlReceive → LatentPrep → TimestepPrep → Denoising → NixlSend.""" - from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( + from sglang.multimodal_gen.runtime.pipelines.stages.latent_preparation import ( LatentPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( + from sglang.multimodal_gen.runtime.pipelines.stages.timestep_preparation import ( TimestepPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + from sglang.multimodal_gen.runtime.pipelines.stages.denoising import ( DenoisingStage, ) from sglang_utils import get_component_backend @@ -300,7 +278,7 @@ def build_denoiser_stages(pipeline, server_args): def build_vae_stages(pipeline, server_args): """NixlReceive → DecodingStage.""" - from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( + from sglang.multimodal_gen.runtime.pipelines.stages.decoding import ( DecodingStage, ) from sglang_utils import get_component_backend diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index 58f12bb5d3b5..ab660ec1f00a 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -85,10 +85,6 @@ save_video, ) -# Backward-compat aliases for any external code that imports old names -_patch_hunyuan_config_task_type = patch_hunyuan_config -_detect_encoder_modules = detect_encoder_modules - # ── Stage launchers ───────────────────────────────────────────────────── @@ -330,9 +326,9 @@ async def main(): logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) # ── Connect clients ────────────────────────────────────────── - enc_client = StageClient(enc_args.scheduler_endpoint, "encoder") - den_client = StageClient(den_args.scheduler_endpoint, "denoiser") - vae_client = StageClient(vae_args.scheduler_endpoint, "vae") + enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") + den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") + vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") # ── Warmup ─────────────────────────────────────────────────── logger.info("Warmup request …") diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py index 887c71114cd4..81818ca1a777 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -3,9 +3,9 @@ """SGLang PipelineStage utilities for disaggregated diffusion workers. -Provides helpers to construct ServerArgs, load partial pipelines (only the -modules each worker needs), and convert between Dynamo protocol types and -SGLang's Req dataclass. +Provides helpers to load partial pipelines (only the modules each worker +needs), launch stage servers, and convert between Dynamo protocol types +and SGLang's Req dataclass. Also contains shared utilities (StageClient, model detection, compatibility patches) used by both Dynamo workers and the standalone E2E script. @@ -16,66 +16,13 @@ import asyncio import logging import os -from typing import Any, Dict, List, Optional +from typing import Dict, List, Optional import torch logger = logging.getLogger(__name__) -def build_server_args(model_path: str, **overrides): - """Create a ServerArgs, initialize torch.distributed, and set the global singleton.""" - from sglang.multimodal_gen.runtime.server_args import ( - ServerArgs, - set_global_server_args, - ) - - defaults = dict( - model_path=model_path, - num_gpus=1, - ) - defaults.update(overrides) - server_args = ServerArgs.from_kwargs(**defaults) - set_global_server_args(server_args) - - _ensure_distributed_init(server_args) - - return server_args - - -def _ensure_distributed_init(server_args): - """Initialize torch.distributed and model-parallel groups via SGLang.""" - import inspect - from sglang.multimodal_gen.runtime.distributed import ( - model_parallel_is_initialized, - maybe_init_distributed_environment_and_model_parallel, - ) - - if model_parallel_is_initialized(): - return - - os.environ.setdefault("MASTER_ADDR", "localhost") - os.environ.setdefault("MASTER_PORT", str(server_args.master_port)) - os.environ.setdefault("LOCAL_RANK", "0") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - - kwargs = dict( - tp_size=server_args.tp_size, - enable_cfg_parallel=server_args.enable_cfg_parallel, - ulysses_degree=server_args.ulysses_degree, - ring_degree=server_args.ring_degree, - sp_size=server_args.sp_degree, - dp_size=server_args.dp_size, - distributed_init_method=f"tcp://127.0.0.1:{server_args.master_port}", - ) - sig = inspect.signature(maybe_init_distributed_environment_and_model_parallel) - if "dist_timeout" in sig.parameters: - kwargs["dist_timeout"] = server_args.dist_timeout - - maybe_init_distributed_environment_and_model_parallel(**kwargs) - - def build_partial_pipeline( server_args, required_modules: List[str], @@ -85,7 +32,7 @@ def build_partial_pipeline( Auto-detects pipeline class from model_index.json, suppresses automatic stage creation, and syncs all component configs (even unloaded ones). """ - from sglang.multimodal_gen.runtime.pipelines_core import get_model_info + from sglang.multimodal_gen.runtime.pipelines import get_model_info model_info = get_model_info(server_args.model_path) base_pipeline_cls = model_info.pipeline_cls @@ -99,7 +46,7 @@ def _noop_create_stages(self, server_args): def _safe_init(self, **kwargs): # Call ComposedPipelineBase.__init__ directly, skipping LoRAPipeline # which tries to access self.modules['transformer'] - from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( ComposedPipelineBase, ) ComposedPipelineBase.__init__(self, **kwargs) @@ -208,9 +155,11 @@ def build_req( **extra_fields, ) -> "Req": """Construct a minimal SGLang ``Req`` for running pipeline stages.""" - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req + from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req + from sglang.multimodal_gen.configs.sample.base import DataType req = Req( + data_type=DataType.VIDEO, prompt=prompt, negative_prompt=negative_prompt, height=height, @@ -229,34 +178,6 @@ def build_req( return req -def extract_tensors_from_req( - req, - keys: List[str], -) -> Dict[str, object]: - """Pull named tensor fields out of a ``Req`` for NIXL transfer. - - Single-element lists are unwrapped to bare tensors. - Multi-element lists (dual-encoder outputs) are preserved as lists. - """ - result: Dict[str, object] = {} - for key in keys: - val = getattr(req, key, None) - if val is None: - continue - if isinstance(val, list): - if len(val) == 0: - continue - if len(val) == 1: - result[key] = val[0] - else: - # Keep multi-element lists intact (e.g. dual-encoder outputs - # with incompatible shapes) - result[key] = val - elif isinstance(val, torch.Tensor): - result[key] = val - return result - - def inject_tensors_to_req( req, tensors: Dict[str, object], @@ -322,9 +243,9 @@ def patch_hunyuan_config(): default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply ``task_type=T2V`` when omitted. Idempotent. """ - from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType + from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType try: - from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( + from sglang.multimodal_gen.configs.pipelines.hunyuan import ( HunyuanConfig, FastHunyuanConfig, ) except ImportError: @@ -398,3 +319,39 @@ def save_video(frames_tensor, output_path: str, fps: int = 24): Image.fromarray(frames[0]).save(output_path) return output_path, len(frames) + + +def launch_stage_server(model_path, required_modules, custom_stages_fn, + scheduler_port, tp_size=1, num_gpus=None, + client_name=""): + """Patch configs, create ServerArgs, launch Scheduler, return (processes, client, server_args). + + Consolidates the boilerplate shared by encoder, denoiser, and VAE workers: + patch_hunyuan_config → ServerArgs → set_global → launch_partial_server → StageClient. + """ + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from partial_gpu_worker import launch_partial_server + + patch_hunyuan_config() + + if num_gpus is None: + num_gpus = tp_size + + server_args = ServerArgs.from_kwargs( + model_path=model_path, + num_gpus=num_gpus, + tp_size=tp_size, + scheduler_port=scheduler_port, + ) + set_global_server_args(server_args) + + processes = launch_partial_server( + server_args, + required_modules=required_modules, + custom_stages_fn=custom_stages_fn, + ) + + client = StageClient(server_args.scheduler_endpoint(), client_name) + return processes, client, server_args diff --git a/examples/disagg_diffusion/phase1_workers/vae_worker.py b/examples/disagg_diffusion/phase1_workers/vae_worker.py index c48e12002b8f..6a65cdb383d7 100755 --- a/examples/disagg_diffusion/phase1_workers/vae_worker.py +++ b/examples/disagg_diffusion/phase1_workers/vae_worker.py @@ -45,33 +45,17 @@ @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - from sglang_utils import StageClient, patch_hunyuan_config, build_req, save_video - from partial_gpu_worker import build_vae_stages, launch_partial_server - from sglang.multimodal_gen.runtime.server_args import ( - ServerArgs, set_global_server_args, - ) + from sglang_utils import launch_stage_server, build_req, save_video + from partial_gpu_worker import build_vae_stages - patch_hunyuan_config() os.makedirs(OUTPUT_DIR, exist_ok=True) - server_args = ServerArgs.from_kwargs( - model_path=MODEL_PATH, - num_gpus=1, - tp_size=1, - scheduler_port=SCHEDULER_PORT, - ) - set_global_server_args(server_args) - logger.info("Launching VAE Scheduler: port=%d", SCHEDULER_PORT) - processes = launch_partial_server( - server_args, - required_modules=["vae", "scheduler"], - custom_stages_fn=build_vae_stages, + processes, client, server_args = launch_stage_server( + MODEL_PATH, ["vae", "scheduler"], build_vae_stages, + SCHEDULER_PORT, client_name="vae", ) - # Connect ZMQ client to local Scheduler - client = StageClient(server_args.scheduler_endpoint, "vae") - # ── Dynamo RPC handlers ────────────────────────────────────────── async def handle_generate(request, context): From a2d7e5813a681419770dddf91fafc2222433af17 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 15:58:31 +0800 Subject: [PATCH 11/22] feat: add multi-worker stage pools with round-robin dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each pipeline stage (Encoder, Denoiser, VAE) can now run N workers. The orchestrator round-robins requests across workers per stage independently, enabling full GPU utilization (e.g. 8 GPU with 2 workers per stage). GPU spec uses `;` to separate workers: GPU_DEN="1,2;5,6" launches 2 TP=2 denoiser workers. Backward compatible — no `;` means 1 worker. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 31 +++- .../phase1_workers/run_e2e_sglang.py | 142 +++++++++++++----- .../phase1_workers/sglang_utils.py | 23 +++ 3 files changed, 155 insertions(+), 41 deletions(-) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 7332b9865624..352fc72e816c 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -73,6 +73,34 @@ Request 2: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ Request 3: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ``` +### Multi-Worker Stage Pools (Standalone E2E) + +The standalone E2E script (`phase1_workers/run_e2e_sglang.py`) supports +**multiple workers per stage** to utilise all available GPUs. Each stage +runs an independent pool; the orchestrator round-robins requests across +workers in each pool. + +``` + ┌─ Encoder_0 (GPU 0) ─┐ ┌─ Denoiser_0 TP=2 (GPU 1,2) ─┐ ┌─ VAE_0 (GPU 3) ─┐ + Request ─┤ ├──►─┤ ├──►─┤ ├─► Video + └─ Encoder_1 (GPU 4) ─┘ └─ Denoiser_1 TP=2 (GPU 5,6) ─┘ └─ VAE_1 (GPU 7) ─┘ + round-robin round-robin round-robin +``` + +Use `;` to separate workers, `,` for TP GPUs within a worker: + +```bash +# 8 GPU — 2 workers per stage +GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python phase1_workers/run_e2e_sglang.py + +# Asymmetric pools (1 encoder, 3 denoisers, 1 VAE) +GPU_ENC="0" GPU_DEN="1,2;3,4;5,6" GPU_VAE="7" python phase1_workers/run_e2e_sglang.py +``` + +No pre-pairing is needed — each `send()` allocates a new NIXL buffer and +the `_keep_alive` coroutine holds it until the receiver's RDMA pull completes, +so any sender→receiver combination is safe. + ## Quick Start One script launches everything (etcd + 3 workers + orchestrator): @@ -127,7 +155,7 @@ curl -X POST http://localhost:8080/v1/videos/generations \ | Variable | Default | Description | |----------|---------|-------------| | `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | -| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage | +| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage (use `;` for multi-worker pools, e.g. `"0;4"`) | | `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism for denoiser | | `PORT` | `8080` | Orchestrator HTTP port | | `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory | @@ -139,6 +167,7 @@ curl -X POST http://localhost:8080/v1/videos/generations \ ## Roadmap +- [x] **Multi-worker stage pools** — configurable N workers per stage with round-robin dispatch (standalone E2E) - [ ] **Dynamic scaling** — auto-scale workers based on queue depth, add/remove denoiser replicas - [ ] **Streaming output** — stream decoded frames to client as they are produced - [ ] **Orchestrator improvements** — smarter scheduling, request priority, load balancing across replicas diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index ab660ec1f00a..289721847ccc 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -8,26 +8,38 @@ processes on different GPUs, then runs the full pipeline: Encoder -> Denoiser -> VAE +Each stage can have multiple workers (a "pool"). The orchestrator +round-robins requests across workers in each pool independently. +Use ``;`` in GPU_ENC / GPU_DEN / GPU_VAE to separate workers. + Measures per-stage timing. Supports concurrent requests for benchmarking. -GPU assignment: +GPU assignment (single-worker, 4 GPU): Encoder : GPU 0 (1 GPU) Denoiser: GPU 1,2 (TP=2) VAE : GPU 3 (1 GPU) +GPU assignment (multi-worker, 8 GPU): + Encoder : GPU 0, 4 (2 workers × 1 GPU) + Denoiser: GPU 1,2 | 5,6 (2 workers × TP=2) + VAE : GPU 3, 7 (2 workers × 1 GPU) + Usage: - # Single request (correctness check) + # Single request, single worker per stage (4 GPU) python run_e2e_sglang.py + # Multi-worker pools (8 GPU) + GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python run_e2e_sglang.py + # Benchmark: 4 requests, 2 concurrent NUM_REQUESTS=4 CONCURRENCY=2 python run_e2e_sglang.py Environment variables: MODEL_PATH Model to use (default: hunyuanvideo-community/HunyuanVideo) PROMPT Text prompt (default: A cat walking on green grass) - GPU_ENC GPU(s) for encoder (default: 0) - GPU_DEN GPU(s) for denoiser (default: 1,2) - GPU_VAE GPU(s) for VAE (default: 3) + GPU_ENC GPU(s) for encoder (default: 0; use "0;4" for 2 workers) + GPU_DEN GPU(s) for denoiser (default: 1,2; use "1,2;5,6" for 2 TP=2 workers) + GPU_VAE GPU(s) for VAE (default: 3; use "3;7" for 2 workers) TP_SIZE Tensor parallelism for denoiser (default: auto from GPU_DEN) NUM_REQUESTS Number of pipeline runs (default: 1) CONCURRENCY Max concurrent pipelines (default: 1) @@ -66,7 +78,7 @@ GPU_ENC = os.environ.get("GPU_ENC", "0") GPU_DEN = os.environ.get("GPU_DEN", "1,2") GPU_VAE = os.environ.get("GPU_VAE", "3") -TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(","))))) +TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(";")[0].split(","))))) NUM_REQUESTS = int(os.environ.get("NUM_REQUESTS", "1")) CONCURRENCY = int(os.environ.get("CONCURRENCY", "1")) NUM_FRAMES = int(os.environ.get("NUM_FRAMES", "61")) @@ -80,6 +92,7 @@ from sglang_utils import ( # noqa: E402 StageClient, + StageWorkerPool, patch_hunyuan_config, detect_encoder_modules, save_video, @@ -139,17 +152,68 @@ def terminate_processes(processes, name=""): logger.info("Terminated %s processes", name) +def _launch_stage_pool( + stage_name: str, + gpu_spec: str, + required_modules: List[str], + custom_stages_fn, + tp_size: int = 1, + base_port: int = 15600, +) -> tuple: + """Launch N workers for one stage, return (all_processes, pool). + + *gpu_spec* uses ``;`` to separate workers and ``,`` for TP GPUs within + a worker. E.g. ``"1,2;5,6"`` means 2 workers each with TP=2. + Single-worker specs (no ``;``) are backward compatible. + """ + worker_gpu_lists = gpu_spec.split(";") + all_procs = [] + clients = [] + for i, cuda_devices in enumerate(worker_gpu_lists): + port = base_port + i * 10 + procs, server_args = _launch_stage( + f"{stage_name}[{i}]", cuda_devices.strip(), + required_modules, custom_stages_fn, + tp_size=tp_size, scheduler_port=port, + ) + all_procs.append(procs) + clients.append(StageClient(server_args.scheduler_endpoint(), + f"{stage_name.lower()}_{i}")) + pool = StageWorkerPool(clients, stage_name) + logger.info("%s pool: %d worker(s)", stage_name, pool.num_workers) + return all_procs, pool + + +def _format_gpu_spec(gpu_spec: str, tp_size: int = 1) -> str: + """Format a GPU spec for the logging header. + + Returns e.g. ``"GPU 0, 4 (2 workers)"`` or ``"GPU 1,2 | 5,6 (2 workers, TP=2)"``. + """ + workers = gpu_spec.split(";") + n = len(workers) + sep = " | " if tp_size > 1 else ", " + gpus = sep.join(w.strip() for w in workers) + if n > 1: + tp_info = f", TP={tp_size}" if tp_size > 1 else "" + return f"GPU {gpus} ({n} workers{tp_info})" + tp_info = f" (TP={tp_size})" if tp_size > 1 else "" + return f"GPU {gpus}{tp_info}" + + # ── Pipeline execution ────────────────────────────────────────────────── async def run_single_pipeline( req_id: int, - encoder_client: StageClient, - denoiser_client: StageClient, - vae_client: StageClient, + encoder_pool: StageWorkerPool, + denoiser_pool: StageWorkerPool, + vae_pool: StageWorkerPool, seed: int, save_output: bool = False, ) -> dict: - """Run one Encoder -> Denoiser -> VAE pipeline, return timing dict.""" + """Run one Encoder -> Denoiser -> VAE pipeline, return timing dict. + + Each pool.forward() round-robins across workers in the pool. + """ import torch from sglang_utils import build_req, inject_tensors_to_req @@ -164,7 +228,7 @@ async def run_single_pipeline( # ── Encoder ────────────────────────────────────────────────────── t0 = time.monotonic() - enc_output = await encoder_client.forward([build_req(**req_kwargs)]) + enc_output = await encoder_pool.forward([build_req(**req_kwargs)]) timings["encoder_s"] = time.monotonic() - t0 if enc_output.error: raise RuntimeError(f"Encoder error: {enc_output.error}") @@ -188,7 +252,7 @@ async def run_single_pipeline( else: # ZMQ fallback — tensors already in enc_result inject_tensors_to_req(den_req, enc_result) - den_output = await denoiser_client.forward([den_req]) + den_output = await denoiser_pool.forward([den_req]) timings["denoiser_s"] = time.monotonic() - t0 if den_output.error: raise RuntimeError(f"Denoiser error: {den_output.error}") @@ -209,7 +273,7 @@ async def run_single_pipeline( vae_req._nixl_transfer_meta = den_result["_nixl_transfer_meta"] else: vae_req.latents = den_result["latents"].cpu() - vae_output = await vae_client.forward([vae_req]) + vae_output = await vae_pool.forward([vae_req]) timings["vae_s"] = time.monotonic() - t0 if vae_output.error: raise RuntimeError(f"VAE error: {vae_output.error}") @@ -284,56 +348,51 @@ async def main(): logger.info(" Disaggregated Diffusion E2E — SGLang Backend") logger.info(" Model: %s", MODEL_PATH) logger.info(" Prompt: %s", PROMPT) - logger.info(" Encoder: GPU %s", GPU_ENC) - logger.info(" Denoiser: GPU %s (TP=%d)", GPU_DEN, TP_SIZE) - logger.info(" VAE: GPU %s", GPU_VAE) + logger.info(" Encoder: %s", _format_gpu_spec(GPU_ENC)) + logger.info(" Denoiser: %s", _format_gpu_spec(GPU_DEN, TP_SIZE)) + logger.info(" VAE: %s", _format_gpu_spec(GPU_VAE)) logger.info(" Requests: %d Concurrency: %d", NUM_REQUESTS, CONCURRENCY) logger.info(" Frames: %d Steps: %d Size: %dx%d Guidance: %.1f", NUM_FRAMES, NUM_STEPS, WIDTH, HEIGHT, GUIDANCE) logger.info("=" * 72) - enc_procs = den_procs = vae_procs = None - enc_client = den_client = vae_client = None + enc_all_procs = den_all_procs = vae_all_procs = None + enc_pool = den_pool = vae_pool = None try: - # ── Launch all 3 stages ────────────────────────────────────── + # ── Launch all 3 stage pools ───────────────────────────────── t_launch = time.monotonic() - enc_procs, enc_args = _launch_stage( + enc_all_procs, enc_pool = _launch_stage_pool( "Encoder", GPU_ENC, required_modules=detect_encoder_modules(MODEL_PATH), custom_stages_fn=build_encoder_stages, tp_size=1, - scheduler_port=15600, + base_port=15600, ) - den_procs, den_args = _launch_stage( + den_all_procs, den_pool = _launch_stage_pool( "Denoiser", GPU_DEN, required_modules=["transformer", "scheduler"], custom_stages_fn=build_denoiser_stages, tp_size=TP_SIZE, - scheduler_port=15700, + base_port=15700, ) - vae_procs, vae_args = _launch_stage( + vae_all_procs, vae_pool = _launch_stage_pool( "VAE", GPU_VAE, required_modules=["vae", "scheduler"], custom_stages_fn=build_vae_stages, tp_size=1, - scheduler_port=15800, + base_port=15800, ) logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) - # ── Connect clients ────────────────────────────────────────── - enc_client = StageClient(enc_args.scheduler_endpoint(), "encoder") - den_client = StageClient(den_args.scheduler_endpoint(), "denoiser") - vae_client = StageClient(vae_args.scheduler_endpoint(), "vae") - # ── Warmup ─────────────────────────────────────────────────── logger.info("Warmup request …") warmup = await run_single_pipeline( - -1, enc_client, den_client, vae_client, SEED, save_output=False, + -1, enc_pool, den_pool, vae_pool, SEED, save_output=False, ) logger.info( "Warmup done — enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", @@ -345,7 +404,7 @@ async def main(): if NUM_REQUESTS <= 1: t_wall = time.monotonic() timings = await run_single_pipeline( - 0, enc_client, den_client, vae_client, SEED, save_output=True, + 0, enc_pool, den_pool, vae_pool, SEED, save_output=True, ) wall_elapsed = time.monotonic() - t_wall print_timing_report([timings], wall_elapsed) @@ -356,7 +415,7 @@ async def main(): async def _run_one(i): async with sem: return await run_single_pipeline( - i, enc_client, den_client, vae_client, + i, enc_pool, den_pool, vae_pool, SEED + i, save_output=(i == 0), ) @@ -367,17 +426,20 @@ async def _run_one(i): print_timing_report(all_timings, wall_elapsed) finally: - for client in [enc_client, den_client, vae_client]: - if client is not None: + for pool in [enc_pool, den_pool, vae_pool]: + if pool is not None: try: - client.close() + pool.close() except Exception: pass - for procs, name in [ - (enc_procs, "encoder"), (den_procs, "denoiser"), (vae_procs, "vae"), + for all_procs, name in [ + (enc_all_procs, "encoder"), + (den_all_procs, "denoiser"), + (vae_all_procs, "vae"), ]: - if procs is not None: - terminate_processes(procs, name) + if all_procs is not None: + for i, procs in enumerate(all_procs): + terminate_processes(procs, f"{name}[{i}]") if __name__ == "__main__": diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py index 81818ca1a777..d7a5e06cc6ac 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -238,6 +238,29 @@ def close(self): self._ctx.term() +class StageWorkerPool: + """Pool of StageClients for one stage with round-robin dispatch.""" + + def __init__(self, clients: List[StageClient], name: str = ""): + self._clients = clients + self._name = name + self._counter = 0 + + @property + def num_workers(self) -> int: + return len(self._clients) + + async def forward(self, reqs): + """Round-robin dispatch to next available worker.""" + idx = self._counter % len(self._clients) + self._counter += 1 + return await self._clients[idx].forward(reqs) + + def close(self): + for c in self._clients: + c.close() + + def patch_hunyuan_config(): """HunyuanConfig inherits ``task_type`` from PipelineConfig without a default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply From 62afdaae2aafae0336e840a0a02c51d64020ec70 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 17:08:09 +0800 Subject: [PATCH 12/22] docs: add multi-worker benchmark usage to README MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add concurrent benchmark example to the multi-worker section. Verified with 50-request run on 8×H20 (2 workers/stage): median=10.79s/req, std=0.55s, 0 failures. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 352fc72e816c..edb54f073a34 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -90,9 +90,14 @@ workers in each pool. Use `;` to separate workers, `,` for TP GPUs within a worker: ```bash -# 8 GPU — 2 workers per stage +# 8 GPU — 2 workers per stage, single request GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python phase1_workers/run_e2e_sglang.py +# 8 GPU — benchmark with 50 requests, 2 concurrent +GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" \ + NUM_REQUESTS=50 CONCURRENCY=2 NUM_FRAMES=9 NUM_STEPS=3 \ + python phase1_workers/run_e2e_sglang.py + # Asymmetric pools (1 encoder, 3 denoisers, 1 VAE) GPU_ENC="0" GPU_DEN="1,2;3,4;5,6" GPU_VAE="7" python phase1_workers/run_e2e_sglang.py ``` From c03f0ca8ffc96ff37c998833a4f688586fae6fd0 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 17:48:46 +0800 Subject: [PATCH 13/22] fix: serialize per-pool access to avoid NIXL concurrent RDMA race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NIXL/UCX crashes when multiple workers in the same stage do concurrent RDMA transfers (NIXL_ERR_REMOTE_DISCONNECT). Add a per-pool asyncio.Lock so only one request per stage is in-flight at a time. Pipeline parallelism across stages is preserved — Encoder, Denoiser, and VAE each run simultaneously on different requests. Use CONCURRENCY>=3 to keep the 3-stage pipeline fully saturated. Warmup now sends N requests (one per worker) to initialise all NIXL connectors before concurrent traffic begins. Verified: 50 requests, 8×H20, concurrency=4, 0 failures. Wall time: 272s (vs 545s sequential) — 2x throughput. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 4 +-- .../phase1_workers/run_e2e_sglang.py | 26 +++++++++++++------ .../phase1_workers/sglang_utils.py | 22 ++++++++++++---- 3 files changed, 37 insertions(+), 15 deletions(-) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index edb54f073a34..b45ba866837d 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -93,9 +93,9 @@ Use `;` to separate workers, `,` for TP GPUs within a worker: # 8 GPU — 2 workers per stage, single request GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python phase1_workers/run_e2e_sglang.py -# 8 GPU — benchmark with 50 requests, 2 concurrent +# 8 GPU — benchmark with 50 requests, pipeline-parallel (CONCURRENCY≥3) GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" \ - NUM_REQUESTS=50 CONCURRENCY=2 NUM_FRAMES=9 NUM_STEPS=3 \ + NUM_REQUESTS=50 CONCURRENCY=4 NUM_FRAMES=9 NUM_STEPS=3 \ python phase1_workers/run_e2e_sglang.py # Asymmetric pools (1 encoder, 3 denoisers, 1 VAE) diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py index 289721847ccc..76c9c1a5203d 100644 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py @@ -390,15 +390,25 @@ async def main(): logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) # ── Warmup ─────────────────────────────────────────────────── - logger.info("Warmup request …") - warmup = await run_single_pipeline( - -1, enc_pool, den_pool, vae_pool, SEED, save_output=False, - ) - logger.info( - "Warmup done — enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", - warmup["encoder_s"], warmup["denoiser_s"], - warmup["vae_s"], warmup["total_s"], + # Send one warmup request per worker (round-robin) so every + # worker's NIXL connector + UCX transport is fully initialised + # before concurrent requests hit them. + num_warmup = max( + enc_pool.num_workers, den_pool.num_workers, vae_pool.num_workers, ) + logger.info("Warmup: %d sequential request(s) …", num_warmup) + for wi in range(num_warmup): + warmup = await run_single_pipeline( + -(wi + 1), enc_pool, den_pool, vae_pool, SEED, + save_output=False, + ) + logger.info( + " warmup %d/%d — enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", + wi + 1, num_warmup, + warmup["encoder_s"], warmup["denoiser_s"], + warmup["vae_s"], warmup["total_s"], + ) + logger.info("Warmup done") # ── Run pipeline(s) ────────────────────────────────────────── if NUM_REQUESTS <= 1: diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/phase1_workers/sglang_utils.py index d7a5e06cc6ac..a0e1a94b6080 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/phase1_workers/sglang_utils.py @@ -239,22 +239,34 @@ def close(self): class StageWorkerPool: - """Pool of StageClients for one stage with round-robin dispatch.""" + """Pool of StageClients for one stage with round-robin dispatch. + + A per-pool asyncio.Lock serialises access so that at most one request + is processed per stage at a time. This avoids concurrent NIXL RDMA + operations within the same stage (which trigger UCX race conditions) + while still allowing **pipeline parallelism** across stages — i.e. + Encoder, Denoiser, and VAE can each be active simultaneously on + different requests. + + Set ``CONCURRENCY >= 3`` to keep the 3-stage pipeline fully saturated. + """ def __init__(self, clients: List[StageClient], name: str = ""): self._clients = clients self._name = name self._counter = 0 + self._lock = asyncio.Lock() @property def num_workers(self) -> int: return len(self._clients) async def forward(self, reqs): - """Round-robin dispatch to next available worker.""" - idx = self._counter % len(self._clients) - self._counter += 1 - return await self._clients[idx].forward(reqs) + """Round-robin dispatch to next available worker (serialised).""" + async with self._lock: + idx = self._counter % len(self._clients) + self._counter += 1 + return await self._clients[idx].forward(reqs) def close(self): for c in self._clients: From 856cd6ba1c5ea781acdabc4bcadf871ad6e2d296 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 20:53:26 +0800 Subject: [PATCH 14/22] refactor: flatten phase prefixes, fix sglang 0.5.8 import paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove development-era directory prefixes (phase0_validate → validate, phase1_workers → workers, phase2_orchestrator → orchestrator) and delete the stale launch/ directory superseded by root run_all.sh. Also update sglang imports for 0.5.8 API changes: - pipelines.{schedule_batch,stages} → pipelines_core.* - configs.pipelines → configs.pipeline_configs - configs.sample.base → configs.sample.sampling_params - scheduler_endpoint() method → property Verified: 8-GPU multi-worker (2 enc, 2 den TP=2, 2 vae) with 3 concurrent requests, all completed successfully. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/README.md | 38 +- examples/disagg_diffusion/launch/run_all.sh | 86 ---- .../run_disagg.py | 0 .../phase1_workers/run_e2e_sglang.py | 457 ------------------ examples/disagg_diffusion/run_all.sh | 79 ++- .../validate_split.py | 0 .../{phase1_workers => workers}/__init__.py | 0 .../denoiser_worker.py | 0 .../encoder_worker.py | 0 .../nixl_transfer.py | 0 .../partial_gpu_worker.py | 14 +- .../{phase1_workers => workers}/protocol.py | 0 .../sglang_utils.py | 54 +-- .../{phase1_workers => workers}/vae_worker.py | 0 14 files changed, 87 insertions(+), 641 deletions(-) delete mode 100755 examples/disagg_diffusion/launch/run_all.sh rename examples/disagg_diffusion/{phase2_orchestrator => orchestrator}/run_disagg.py (100%) delete mode 100644 examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py rename examples/disagg_diffusion/{phase0_validate => validate}/validate_split.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/__init__.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/denoiser_worker.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/encoder_worker.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/nixl_transfer.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/partial_gpu_worker.py (97%) rename examples/disagg_diffusion/{phase1_workers => workers}/protocol.py (100%) rename examples/disagg_diffusion/{phase1_workers => workers}/sglang_utils.py (86%) rename examples/disagg_diffusion/{phase1_workers => workers}/vae_worker.py (100%) diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index b45ba866837d..f756ea006cdb 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -73,12 +73,11 @@ Request 2: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ Request 3: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ``` -### Multi-Worker Stage Pools (Standalone E2E) +### Multi-Worker Scaling -The standalone E2E script (`phase1_workers/run_e2e_sglang.py`) supports -**multiple workers per stage** to utilise all available GPUs. Each stage -runs an independent pool; the orchestrator round-robins requests across -workers in each pool. +Each stage supports **multiple workers** — just launch more processes with the +same Dynamo component name. They register via etcd and the orchestrator +round-robins requests automatically. Use `;` in GPU specs to separate workers: ``` ┌─ Encoder_0 (GPU 0) ─┐ ┌─ Denoiser_0 TP=2 (GPU 1,2) ─┐ ┌─ VAE_0 (GPU 3) ─┐ @@ -87,21 +86,16 @@ workers in each pool. round-robin round-robin round-robin ``` -Use `;` to separate workers, `,` for TP GPUs within a worker: - ```bash -# 8 GPU — 2 workers per stage, single request -GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python phase1_workers/run_e2e_sglang.py - -# 8 GPU — benchmark with 50 requests, pipeline-parallel (CONCURRENCY≥3) -GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" \ - NUM_REQUESTS=50 CONCURRENCY=4 NUM_FRAMES=9 NUM_STEPS=3 \ - python phase1_workers/run_e2e_sglang.py +# 8 GPU — 2 workers per stage +GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" ./run_all.sh --test --quick -# Asymmetric pools (1 encoder, 3 denoisers, 1 VAE) -GPU_ENC="0" GPU_DEN="1,2;3,4;5,6" GPU_VAE="7" python phase1_workers/run_e2e_sglang.py +# Asymmetric (1 encoder, 3 denoisers, 1 VAE) +GPU_ENC="0" GPU_DEN="1,2;3,4;5,6" GPU_VAE="7" ./run_all.sh ``` +Single-worker specs (no `;`) are fully backward compatible. + No pre-pairing is needed — each `send()` allocates a new NIXL buffer and the `_keep_alive` coroutine holds it until the receiver's RDMA pull completes, so any sender→receiver combination is safe. @@ -131,12 +125,12 @@ Or launch each service manually: etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 # Terminal 1-3: Workers -CUDA_VISIBLE_DEVICES=0 python phase1_workers/encoder_worker.py -CUDA_VISIBLE_DEVICES=1,2 python phase1_workers/denoiser_worker.py -CUDA_VISIBLE_DEVICES=3 python phase1_workers/vae_worker.py +CUDA_VISIBLE_DEVICES=0 python workers/encoder_worker.py +CUDA_VISIBLE_DEVICES=1,2 python workers/denoiser_worker.py +CUDA_VISIBLE_DEVICES=3 python workers/vae_worker.py # Terminal 4: Orchestrator -python phase2_orchestrator/run_disagg.py +python orchestrator/run_disagg.py ``` Generate a video (61 frames, 50 steps, 544x960 by default): @@ -160,7 +154,7 @@ curl -X POST http://localhost:8080/v1/videos/generations \ | Variable | Default | Description | |----------|---------|-------------| | `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | -| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage (use `;` for multi-worker pools, e.g. `"0;4"`) | +| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage (use `;` to launch multiple workers, e.g. `"0;4"`) | | `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism for denoiser | | `PORT` | `8080` | Orchestrator HTTP port | | `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory | @@ -172,7 +166,7 @@ curl -X POST http://localhost:8080/v1/videos/generations \ ## Roadmap -- [x] **Multi-worker stage pools** — configurable N workers per stage with round-robin dispatch (standalone E2E) +- [x] **Multi-worker scaling** — launch N workers per stage via `run_all.sh`; Dynamo auto-discovers and round-robins - [ ] **Dynamic scaling** — auto-scale workers based on queue depth, add/remove denoiser replicas - [ ] **Streaming output** — stream decoded frames to client as they are produced - [ ] **Orchestrator improvements** — smarter scheduling, request priority, load balancing across replicas diff --git a/examples/disagg_diffusion/launch/run_all.sh b/examples/disagg_diffusion/launch/run_all.sh deleted file mode 100755 index aa3f131b3285..000000000000 --- a/examples/disagg_diffusion/launch/run_all.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Disaggregated Diffusion POC — all-in-one launcher -# -# Starts three stage workers (Encoder, Denoiser, VAE) as background -# processes, waits for them to be ready, then runs the orchestrator. -# -# Requirements: -# - 3 GPUs (or 1 large GPU >48 GB — set SINGLE_GPU=1) -# - diffusers, transformers, torch, dynamo runtime, uvloop -# -# Usage: -# bash launch/run_all.sh [MODEL] [PROMPT] -# -# Examples: -# bash launch/run_all.sh -# bash launch/run_all.sh black-forest-labs/FLUX.1-schnell "A sunset over mountains" -# SINGLE_GPU=1 bash launch/run_all.sh # all stages on GPU 0 - -set -euo pipefail - -export MODEL_PATH="${1:-black-forest-labs/FLUX.1-schnell}" -export PROMPT="${2:-A photo of a cat sitting on a windowsill}" -SCRIPT_DIR="$(cd "$(dirname "$0")/.." && pwd)" -WORKERS_DIR="${SCRIPT_DIR}/phase1_workers" -ORCH_DIR="${SCRIPT_DIR}/phase2_orchestrator" - -echo "============================================" -echo " Disaggregated Diffusion POC" -echo " Model: ${MODEL_PATH}" -echo " Prompt: ${PROMPT}" -echo "============================================" -echo "" - -PIDS=() - -cleanup() { - echo "" - echo "Shutting down workers …" - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - wait 2>/dev/null - echo "Done." -} -trap cleanup EXIT - -if [[ "${SINGLE_GPU:-0}" == "1" ]]; then - GPU_ENC=0; GPU_DEN=0; GPU_VAE=0 -else - GPU_ENC=0; GPU_DEN=1; GPU_VAE=2 -fi - -# --- Start workers -------------------------------------------------------- - -echo "[1/3] Starting Encoder Worker (GPU ${GPU_ENC}) …" -CUDA_VISIBLE_DEVICES=${GPU_ENC} python "${WORKERS_DIR}/encoder_worker.py" \ - 2>&1 | sed 's/^/ [encoder] /' & -PIDS+=($!) - -echo "[2/3] Starting Denoiser Worker (GPU ${GPU_DEN}) …" -CUDA_VISIBLE_DEVICES=${GPU_DEN} python "${WORKERS_DIR}/denoiser_worker.py" \ - 2>&1 | sed 's/^/ [denoiser] /' & -PIDS+=($!) - -echo "[3/3] Starting VAE Worker (GPU ${GPU_VAE}) …" -CUDA_VISIBLE_DEVICES=${GPU_VAE} python "${WORKERS_DIR}/vae_worker.py" \ - 2>&1 | sed 's/^/ [vae] /' & -PIDS+=($!) - -# Wait for model loading. In production use Dynamo health checks. -WAIT_SECS="${WAIT_SECS:-90}" -echo "" -echo "Waiting ${WAIT_SECS}s for workers to load models …" -sleep "${WAIT_SECS}" - -# --- Run orchestrator ----------------------------------------------------- - -echo "" -echo "Running orchestrator …" -echo "" - -export OUTPUT="/tmp/disagg_output.png" -python "${ORCH_DIR}/run_disagg.py" diff --git a/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py similarity index 100% rename from examples/disagg_diffusion/phase2_orchestrator/run_disagg.py rename to examples/disagg_diffusion/orchestrator/run_disagg.py diff --git a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py b/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py deleted file mode 100644 index 76c9c1a5203d..000000000000 --- a/examples/disagg_diffusion/phase1_workers/run_e2e_sglang.py +++ /dev/null @@ -1,457 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""End-to-end disaggregated diffusion pipeline with SGLang backend. - -Launches Encoder, Denoiser (TP=2), and VAE as separate SGLang scheduler -processes on different GPUs, then runs the full pipeline: - Encoder -> Denoiser -> VAE - -Each stage can have multiple workers (a "pool"). The orchestrator -round-robins requests across workers in each pool independently. -Use ``;`` in GPU_ENC / GPU_DEN / GPU_VAE to separate workers. - -Measures per-stage timing. Supports concurrent requests for benchmarking. - -GPU assignment (single-worker, 4 GPU): - Encoder : GPU 0 (1 GPU) - Denoiser: GPU 1,2 (TP=2) - VAE : GPU 3 (1 GPU) - -GPU assignment (multi-worker, 8 GPU): - Encoder : GPU 0, 4 (2 workers × 1 GPU) - Denoiser: GPU 1,2 | 5,6 (2 workers × TP=2) - VAE : GPU 3, 7 (2 workers × 1 GPU) - -Usage: - # Single request, single worker per stage (4 GPU) - python run_e2e_sglang.py - - # Multi-worker pools (8 GPU) - GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" python run_e2e_sglang.py - - # Benchmark: 4 requests, 2 concurrent - NUM_REQUESTS=4 CONCURRENCY=2 python run_e2e_sglang.py - -Environment variables: - MODEL_PATH Model to use (default: hunyuanvideo-community/HunyuanVideo) - PROMPT Text prompt (default: A cat walking on green grass) - GPU_ENC GPU(s) for encoder (default: 0; use "0;4" for 2 workers) - GPU_DEN GPU(s) for denoiser (default: 1,2; use "1,2;5,6" for 2 TP=2 workers) - GPU_VAE GPU(s) for VAE (default: 3; use "3;7" for 2 workers) - TP_SIZE Tensor parallelism for denoiser (default: auto from GPU_DEN) - NUM_REQUESTS Number of pipeline runs (default: 1) - CONCURRENCY Max concurrent pipelines (default: 1) - NUM_FRAMES Number of video frames (default: 61) - NUM_STEPS Denoising steps (default: 50) - HEIGHT Frame height (default: 544) - WIDTH Frame width (default: 960) - GUIDANCE Guidance scale (default: 1.0; >1.0 enables CFG) -""" - -from __future__ import annotations - -import asyncio -import logging -import multiprocessing as mp -import os -import statistics -import sys -import time -from typing import List - -# --- ensure workers dir is on sys.path so subprocesses find sglang_utils --- -WORKERS_DIR = os.path.dirname(os.path.abspath(__file__)) -if WORKERS_DIR not in sys.path: - sys.path.insert(0, WORKERS_DIR) - -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", -) -logger = logging.getLogger("e2e") - -# ── Configuration ──────────────────────────────────────────────────────── -MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") -PROMPT = os.environ.get("PROMPT", "A cat walking on green grass") -GPU_ENC = os.environ.get("GPU_ENC", "0") -GPU_DEN = os.environ.get("GPU_DEN", "1,2") -GPU_VAE = os.environ.get("GPU_VAE", "3") -TP_SIZE = int(os.environ.get("TP_SIZE", str(len(GPU_DEN.split(";")[0].split(","))))) -NUM_REQUESTS = int(os.environ.get("NUM_REQUESTS", "1")) -CONCURRENCY = int(os.environ.get("CONCURRENCY", "1")) -NUM_FRAMES = int(os.environ.get("NUM_FRAMES", "61")) -NUM_STEPS = int(os.environ.get("NUM_STEPS", "50")) -HEIGHT = int(os.environ.get("HEIGHT", "544")) -WIDTH = int(os.environ.get("WIDTH", "960")) -GUIDANCE = float(os.environ.get("GUIDANCE", "1.0")) -SEED = int(os.environ.get("SEED", "42")) -OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_e2e") - - -from sglang_utils import ( # noqa: E402 - StageClient, - StageWorkerPool, - patch_hunyuan_config, - detect_encoder_modules, - save_video, -) - - -# ── Stage launchers ───────────────────────────────────────────────────── - -def _launch_stage( - stage_name: str, - cuda_devices: str, - required_modules: List[str], - custom_stages_fn, - tp_size: int = 1, - scheduler_port: int = 15600, -): - """Launch a partial scheduler for one pipeline stage. - - Returns (processes, server_args). - """ - os.environ["CUDA_VISIBLE_DEVICES"] = cuda_devices - num_gpus = len(cuda_devices.split(",")) - - from sglang.multimodal_gen.runtime.server_args import ( - ServerArgs, set_global_server_args, - ) - from partial_gpu_worker import launch_partial_server - - server_args = ServerArgs.from_kwargs( - model_path=MODEL_PATH, - num_gpus=num_gpus, - tp_size=tp_size, - scheduler_port=scheduler_port, - ) - set_global_server_args(server_args) - - logger.info( - "Launching %s: CUDA_VISIBLE_DEVICES=%s num_gpus=%d tp=%d port=%d", - stage_name, cuda_devices, num_gpus, tp_size, server_args.scheduler_port, - ) - - processes = launch_partial_server( - server_args, - required_modules=required_modules, - custom_stages_fn=custom_stages_fn, - ) - - logger.info("%s ready (%d processes)", stage_name, len(processes)) - return processes, server_args - - -def terminate_processes(processes, name=""): - for p in processes: - p.terminate() - for p in processes: - p.join(timeout=10) - logger.info("Terminated %s processes", name) - - -def _launch_stage_pool( - stage_name: str, - gpu_spec: str, - required_modules: List[str], - custom_stages_fn, - tp_size: int = 1, - base_port: int = 15600, -) -> tuple: - """Launch N workers for one stage, return (all_processes, pool). - - *gpu_spec* uses ``;`` to separate workers and ``,`` for TP GPUs within - a worker. E.g. ``"1,2;5,6"`` means 2 workers each with TP=2. - Single-worker specs (no ``;``) are backward compatible. - """ - worker_gpu_lists = gpu_spec.split(";") - all_procs = [] - clients = [] - for i, cuda_devices in enumerate(worker_gpu_lists): - port = base_port + i * 10 - procs, server_args = _launch_stage( - f"{stage_name}[{i}]", cuda_devices.strip(), - required_modules, custom_stages_fn, - tp_size=tp_size, scheduler_port=port, - ) - all_procs.append(procs) - clients.append(StageClient(server_args.scheduler_endpoint(), - f"{stage_name.lower()}_{i}")) - pool = StageWorkerPool(clients, stage_name) - logger.info("%s pool: %d worker(s)", stage_name, pool.num_workers) - return all_procs, pool - - -def _format_gpu_spec(gpu_spec: str, tp_size: int = 1) -> str: - """Format a GPU spec for the logging header. - - Returns e.g. ``"GPU 0, 4 (2 workers)"`` or ``"GPU 1,2 | 5,6 (2 workers, TP=2)"``. - """ - workers = gpu_spec.split(";") - n = len(workers) - sep = " | " if tp_size > 1 else ", " - gpus = sep.join(w.strip() for w in workers) - if n > 1: - tp_info = f", TP={tp_size}" if tp_size > 1 else "" - return f"GPU {gpus} ({n} workers{tp_info})" - tp_info = f" (TP={tp_size})" if tp_size > 1 else "" - return f"GPU {gpus}{tp_info}" - - -# ── Pipeline execution ────────────────────────────────────────────────── - -async def run_single_pipeline( - req_id: int, - encoder_pool: StageWorkerPool, - denoiser_pool: StageWorkerPool, - vae_pool: StageWorkerPool, - seed: int, - save_output: bool = False, -) -> dict: - """Run one Encoder -> Denoiser -> VAE pipeline, return timing dict. - - Each pool.forward() round-robins across workers in the pool. - """ - import torch - from sglang_utils import build_req, inject_tensors_to_req - - timings = {"req_id": req_id} - negative_prompt = "bad quality" if GUIDANCE > 1.0 else "" - req_kwargs = dict( - prompt=PROMPT, negative_prompt=negative_prompt, - height=HEIGHT, width=WIDTH, num_frames=NUM_FRAMES, - num_inference_steps=NUM_STEPS, guidance_scale=GUIDANCE, seed=seed, - ) - t_pipeline = time.monotonic() - - # ── Encoder ────────────────────────────────────────────────────── - t0 = time.monotonic() - enc_output = await encoder_pool.forward([build_req(**req_kwargs)]) - timings["encoder_s"] = time.monotonic() - t0 - if enc_output.error: - raise RuntimeError(f"Encoder error: {enc_output.error}") - enc_result = enc_output.output - # enc_result is either {"_nixl_transfer_meta": {...}} (NIXL) or - # {"prompt_embeds": tensor, ...} (ZMQ fallback) - nixl_mode = "_nixl_transfer_meta" in enc_result - logger.info( - "req %d | Encoder done %.2fs — transfer: %s", - req_id, timings["encoder_s"], - "NIXL RDMA" if nixl_mode else f"ZMQ (keys: {list(enc_result.keys())})", - ) - - # ── Denoiser ───────────────────────────────────────────────────── - t0 = time.monotonic() - den_req = build_req(**req_kwargs) - den_req.do_classifier_free_guidance = (GUIDANCE > 1.0) - if nixl_mode: - # Pass NIXL metadata — NixlReceiveStage will RDMA-pull the tensors - den_req._nixl_transfer_meta = enc_result["_nixl_transfer_meta"] - else: - # ZMQ fallback — tensors already in enc_result - inject_tensors_to_req(den_req, enc_result) - den_output = await denoiser_pool.forward([den_req]) - timings["denoiser_s"] = time.monotonic() - t0 - if den_output.error: - raise RuntimeError(f"Denoiser error: {den_output.error}") - den_result = den_output.output - nixl_mode_den = "_nixl_transfer_meta" in den_result - logger.info( - "req %d | Denoiser done %.2fs — transfer: %s", - req_id, timings["denoiser_s"], - "NIXL RDMA" if nixl_mode_den else "ZMQ", - ) - - # ── VAE ────────────────────────────────────────────────────────── - t0 = time.monotonic() - vae_req = build_req(prompt="", height=HEIGHT, width=WIDTH, - num_frames=NUM_FRAMES, num_inference_steps=NUM_STEPS, - guidance_scale=0.0, seed=seed) - if nixl_mode_den: - vae_req._nixl_transfer_meta = den_result["_nixl_transfer_meta"] - else: - vae_req.latents = den_result["latents"].cpu() - vae_output = await vae_pool.forward([vae_req]) - timings["vae_s"] = time.monotonic() - t0 - if vae_output.error: - raise RuntimeError(f"VAE error: {vae_output.error}") - - timings["total_s"] = time.monotonic() - t_pipeline - logger.info( - "req %d | VAE done %.2fs — total pipeline: %.2fs", - req_id, timings["vae_s"], timings["total_s"], - ) - - # Save output as mp4 - if save_output and vae_output.output is not None: - _save_video(vae_output.output, req_id) - - return timings - - -def _save_video(frames_tensor, req_id: int): - """Save decoded video tensor [B,C,T,H,W] as mp4.""" - try: - out_path = os.path.join(OUTPUT_DIR, f"output_{req_id}.mp4") - filepath, n_frames = save_video(frames_tensor, out_path) - logger.info("req %d | Saved %d frames to %s", req_id, n_frames, filepath) - except Exception as e: - logger.warning("req %d | Could not save video: %s", req_id, e) - - -def print_timing_report(all_timings: list, wall_elapsed: float): - """Print per-stage timing statistics.""" - stages = ["encoder_s", "denoiser_s", "vae_s", "total_s"] - n = len(all_timings) - - logger.info("") - logger.info("=" * 72) - logger.info(" Timing Report (%d requests, concurrency=%d)", n, CONCURRENCY) - logger.info("=" * 72) - - for t in all_timings: - logger.info( - " req %2d | enc=%6.2fs den=%6.2fs vae=%6.2fs total=%6.2fs", - t["req_id"], t["encoder_s"], t["denoiser_s"], t["vae_s"], t["total_s"], - ) - - if n > 1: - logger.info("-" * 72) - for stage in stages: - vals = [t[stage] for t in all_timings] - mean = statistics.mean(vals) - med = statistics.median(vals) - mn, mx = min(vals), max(vals) - std = statistics.stdev(vals) if n >= 2 else 0.0 - logger.info( - " %-10s mean=%6.2fs median=%6.2fs min=%6.2fs max=%6.2fs std=%5.2fs", - stage, mean, med, mn, mx, std, - ) - - logger.info("-" * 72) - throughput = n / wall_elapsed if wall_elapsed > 0 else 0 - logger.info(" Wall time: %.2fs | Throughput: %.2f req/s", wall_elapsed, throughput) - logger.info("=" * 72) - - -# ── Main ──────────────────────────────────────────────────────────────── - -async def main(): - from partial_gpu_worker import build_encoder_stages, build_denoiser_stages, build_vae_stages - - # Apply patches once before any SGLang config is created - patch_hunyuan_config() - - logger.info("=" * 72) - logger.info(" Disaggregated Diffusion E2E — SGLang Backend") - logger.info(" Model: %s", MODEL_PATH) - logger.info(" Prompt: %s", PROMPT) - logger.info(" Encoder: %s", _format_gpu_spec(GPU_ENC)) - logger.info(" Denoiser: %s", _format_gpu_spec(GPU_DEN, TP_SIZE)) - logger.info(" VAE: %s", _format_gpu_spec(GPU_VAE)) - logger.info(" Requests: %d Concurrency: %d", NUM_REQUESTS, CONCURRENCY) - logger.info(" Frames: %d Steps: %d Size: %dx%d Guidance: %.1f", - NUM_FRAMES, NUM_STEPS, WIDTH, HEIGHT, GUIDANCE) - logger.info("=" * 72) - - enc_all_procs = den_all_procs = vae_all_procs = None - enc_pool = den_pool = vae_pool = None - - try: - # ── Launch all 3 stage pools ───────────────────────────────── - t_launch = time.monotonic() - - enc_all_procs, enc_pool = _launch_stage_pool( - "Encoder", GPU_ENC, - required_modules=detect_encoder_modules(MODEL_PATH), - custom_stages_fn=build_encoder_stages, - tp_size=1, - base_port=15600, - ) - - den_all_procs, den_pool = _launch_stage_pool( - "Denoiser", GPU_DEN, - required_modules=["transformer", "scheduler"], - custom_stages_fn=build_denoiser_stages, - tp_size=TP_SIZE, - base_port=15700, - ) - - vae_all_procs, vae_pool = _launch_stage_pool( - "VAE", GPU_VAE, - required_modules=["vae", "scheduler"], - custom_stages_fn=build_vae_stages, - tp_size=1, - base_port=15800, - ) - - logger.info("All stages launched in %.1fs", time.monotonic() - t_launch) - - # ── Warmup ─────────────────────────────────────────────────── - # Send one warmup request per worker (round-robin) so every - # worker's NIXL connector + UCX transport is fully initialised - # before concurrent requests hit them. - num_warmup = max( - enc_pool.num_workers, den_pool.num_workers, vae_pool.num_workers, - ) - logger.info("Warmup: %d sequential request(s) …", num_warmup) - for wi in range(num_warmup): - warmup = await run_single_pipeline( - -(wi + 1), enc_pool, den_pool, vae_pool, SEED, - save_output=False, - ) - logger.info( - " warmup %d/%d — enc=%.2fs den=%.2fs vae=%.2fs total=%.2fs", - wi + 1, num_warmup, - warmup["encoder_s"], warmup["denoiser_s"], - warmup["vae_s"], warmup["total_s"], - ) - logger.info("Warmup done") - - # ── Run pipeline(s) ────────────────────────────────────────── - if NUM_REQUESTS <= 1: - t_wall = time.monotonic() - timings = await run_single_pipeline( - 0, enc_pool, den_pool, vae_pool, SEED, save_output=True, - ) - wall_elapsed = time.monotonic() - t_wall - print_timing_report([timings], wall_elapsed) - else: - logger.info("Firing %d requests (concurrency=%d) …", NUM_REQUESTS, CONCURRENCY) - sem = asyncio.Semaphore(CONCURRENCY) - - async def _run_one(i): - async with sem: - return await run_single_pipeline( - i, enc_pool, den_pool, vae_pool, - SEED + i, save_output=(i == 0), - ) - - t_wall = time.monotonic() - tasks = [asyncio.create_task(_run_one(i)) for i in range(NUM_REQUESTS)] - all_timings = list(await asyncio.gather(*tasks)) - wall_elapsed = time.monotonic() - t_wall - print_timing_report(all_timings, wall_elapsed) - - finally: - for pool in [enc_pool, den_pool, vae_pool]: - if pool is not None: - try: - pool.close() - except Exception: - pass - for all_procs, name in [ - (enc_all_procs, "encoder"), - (den_all_procs, "denoiser"), - (vae_all_procs, "vae"), - ]: - if all_procs is not None: - for i, procs in enumerate(all_procs): - terminate_processes(procs, f"{name}[{i}]") - - -if __name__ == "__main__": - mp.set_start_method("spawn", force=True) - asyncio.run(main()) diff --git a/examples/disagg_diffusion/run_all.sh b/examples/disagg_diffusion/run_all.sh index 7ee7a1c5f32e..b9fba7bd5015 100755 --- a/examples/disagg_diffusion/run_all.sh +++ b/examples/disagg_diffusion/run_all.sh @@ -2,26 +2,33 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# Launch all disaggregated diffusion services (etcd + 3 workers + orchestrator) +# Launch all disaggregated diffusion services (etcd + workers + orchestrator) # and optionally send a test request. # +# Each stage supports multiple workers: use ';' to separate workers in GPU +# specs. Each worker is an independent top-level process; Dynamo discovers +# them via etcd and the orchestrator round-robins requests automatically. +# # Usage: -# ./run_all.sh # launch all services +# ./run_all.sh # launch all services (1 worker/stage) # ./run_all.sh --test # launch + send a test request # ./run_all.sh --test --quick # launch + quick smoke test (9 frames, 3 steps) # +# # Multi-worker (8 GPU): +# GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" ./run_all.sh --test --quick +# # Environment variables: # MODEL_PATH HuggingFace model (default: hunyuanvideo-community/HunyuanVideo) -# GPU_ENC GPU for encoder (default: 0) -# GPU_DEN GPUs for denoiser (default: 1,2) -# GPU_VAE GPU for VAE (default: 3) +# GPU_ENC GPU(s) for encoder (default: 0; use "0;4" for 2 workers) +# GPU_DEN GPU(s) for denoiser (default: 1,2; use "1,2;5,6" for 2 TP=2 workers) +# GPU_VAE GPU(s) for VAE (default: 3; use "3;7" for 2 workers) # PORT HTTP port (default: 8080) set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -WORKERS_DIR="$SCRIPT_DIR/phase1_workers" -ORCH_DIR="$SCRIPT_DIR/phase2_orchestrator" +WORKERS_DIR="$SCRIPT_DIR/workers" +ORCH_DIR="$SCRIPT_DIR/orchestrator" LOG_DIR="/tmp/disagg_logs" GPU_ENC="${GPU_ENC:-0}" @@ -79,35 +86,55 @@ else echo "[1/5] etcd already running, skipping." fi -# 2. Encoder Worker -echo "[2/5] Starting Encoder Worker (GPU $GPU_ENC)..." -CUDA_VISIBLE_DEVICES="$GPU_ENC" python "$WORKERS_DIR/encoder_worker.py" \ - > "$LOG_DIR/encoder.log" 2>&1 & -PIDS+=($!) +# 2. Encoder Worker(s) +STEP=2 +IFS=';' read -ra ENC_GPUS <<< "$GPU_ENC" +for i in "${!ENC_GPUS[@]}"; do + port=$((15600 + i * 10)) + echo "[$STEP] Starting Encoder Worker $i (GPU ${ENC_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${ENC_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/encoder_worker.py" > "$LOG_DIR/encoder_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done -# 3. Denoiser Worker -echo "[3/5] Starting Denoiser Worker (GPU $GPU_DEN)..." -CUDA_VISIBLE_DEVICES="$GPU_DEN" python "$WORKERS_DIR/denoiser_worker.py" \ - > "$LOG_DIR/denoiser.log" 2>&1 & -PIDS+=($!) +# 3. Denoiser Worker(s) +IFS=';' read -ra DEN_GPUS <<< "$GPU_DEN" +for i in "${!DEN_GPUS[@]}"; do + port=$((15700 + i * 10)) + echo "[$STEP] Starting Denoiser Worker $i (GPU ${DEN_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${DEN_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/denoiser_worker.py" > "$LOG_DIR/denoiser_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done -# 4. VAE Worker -echo "[4/5] Starting VAE Worker (GPU $GPU_VAE)..." -CUDA_VISIBLE_DEVICES="$GPU_VAE" python "$WORKERS_DIR/vae_worker.py" \ - > "$LOG_DIR/vae.log" 2>&1 & -PIDS+=($!) +# 4. VAE Worker(s) +IFS=';' read -ra VAE_GPUS <<< "$GPU_VAE" +for i in "${!VAE_GPUS[@]}"; do + port=$((15800 + i * 10)) + echo "[$STEP] Starting VAE Worker $i (GPU ${VAE_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${VAE_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/vae_worker.py" > "$LOG_DIR/vae_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done + +N_WORKERS=$(( ${#ENC_GPUS[@]} + ${#DEN_GPUS[@]} + ${#VAE_GPUS[@]} )) +echo "" +echo "Launched $N_WORKERS worker(s): ${#ENC_GPUS[@]} encoder, ${#DEN_GPUS[@]} denoiser, ${#VAE_GPUS[@]} vae" # 5. Orchestrator -echo "[5/5] Starting Orchestrator (port $PORT)..." +echo "[$STEP] Starting Orchestrator (port $PORT)..." PORT="$PORT" python "$ORCH_DIR/run_disagg.py" \ > "$LOG_DIR/orchestrator.log" 2>&1 & PIDS+=($!) echo "" echo "All services launching. Waiting for workers to be ready..." -echo " tail -f $LOG_DIR/encoder.log # monitor encoder" -echo " tail -f $LOG_DIR/denoiser.log # monitor denoiser" -echo " tail -f $LOG_DIR/vae.log # monitor vae" +echo " tail -f $LOG_DIR/encoder_*.log # monitor encoder(s)" +echo " tail -f $LOG_DIR/denoiser_*.log # monitor denoiser(s)" +echo " tail -f $LOG_DIR/vae_*.log # monitor vae(s)" echo " tail -f $LOG_DIR/orchestrator.log" echo "" diff --git a/examples/disagg_diffusion/phase0_validate/validate_split.py b/examples/disagg_diffusion/validate/validate_split.py similarity index 100% rename from examples/disagg_diffusion/phase0_validate/validate_split.py rename to examples/disagg_diffusion/validate/validate_split.py diff --git a/examples/disagg_diffusion/phase1_workers/__init__.py b/examples/disagg_diffusion/workers/__init__.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/__init__.py rename to examples/disagg_diffusion/workers/__init__.py diff --git a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py b/examples/disagg_diffusion/workers/denoiser_worker.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/denoiser_worker.py rename to examples/disagg_diffusion/workers/denoiser_worker.py diff --git a/examples/disagg_diffusion/phase1_workers/encoder_worker.py b/examples/disagg_diffusion/workers/encoder_worker.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/encoder_worker.py rename to examples/disagg_diffusion/workers/encoder_worker.py diff --git a/examples/disagg_diffusion/phase1_workers/nixl_transfer.py b/examples/disagg_diffusion/workers/nixl_transfer.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/nixl_transfer.py rename to examples/disagg_diffusion/workers/nixl_transfer.py diff --git a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py similarity index 97% rename from examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py rename to examples/disagg_diffusion/workers/partial_gpu_worker.py index b677c64a9b8c..b3b5a221a96e 100644 --- a/examples/disagg_diffusion/phase1_workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -57,8 +57,8 @@ get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker -from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req, OutputBatch -from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs # layerwise_offload may not exist in all sglang versions — guard import @@ -219,7 +219,7 @@ def build_encoder_stages(pipeline, server_args): Automatically detects all loaded text encoders/tokenizers so that both single-encoder (Wan) and dual-encoder (HunyuanVideo) models work. """ - from sglang.multimodal_gen.runtime.pipelines.stages.text_encoding import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( TextEncodingStage, ) from sglang_utils import get_component_backend @@ -252,13 +252,13 @@ def build_encoder_stages(pipeline, server_args): def build_denoiser_stages(pipeline, server_args): """NixlReceive → LatentPrep → TimestepPrep → Denoising → NixlSend.""" - from sglang.multimodal_gen.runtime.pipelines.stages.latent_preparation import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines.stages.timestep_preparation import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( TimestepPreparationStage, ) - from sglang.multimodal_gen.runtime.pipelines.stages.denoising import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( DenoisingStage, ) from sglang_utils import get_component_backend @@ -278,7 +278,7 @@ def build_denoiser_stages(pipeline, server_args): def build_vae_stages(pipeline, server_args): """NixlReceive → DecodingStage.""" - from sglang.multimodal_gen.runtime.pipelines.stages.decoding import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( DecodingStage, ) from sglang_utils import get_component_backend diff --git a/examples/disagg_diffusion/phase1_workers/protocol.py b/examples/disagg_diffusion/workers/protocol.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/protocol.py rename to examples/disagg_diffusion/workers/protocol.py diff --git a/examples/disagg_diffusion/phase1_workers/sglang_utils.py b/examples/disagg_diffusion/workers/sglang_utils.py similarity index 86% rename from examples/disagg_diffusion/phase1_workers/sglang_utils.py rename to examples/disagg_diffusion/workers/sglang_utils.py index a0e1a94b6080..2468cb791153 100644 --- a/examples/disagg_diffusion/phase1_workers/sglang_utils.py +++ b/examples/disagg_diffusion/workers/sglang_utils.py @@ -8,7 +8,7 @@ and SGLang's Req dataclass. Also contains shared utilities (StageClient, model detection, compatibility -patches) used by both Dynamo workers and the standalone E2E script. +patches) used by Dynamo workers. """ from __future__ import annotations @@ -32,7 +32,7 @@ def build_partial_pipeline( Auto-detects pipeline class from model_index.json, suppresses automatic stage creation, and syncs all component configs (even unloaded ones). """ - from sglang.multimodal_gen.runtime.pipelines import get_model_info + from sglang.multimodal_gen.runtime.pipelines_core import get_model_info model_info = get_model_info(server_args.model_path) base_pipeline_cls = model_info.pipeline_cls @@ -46,7 +46,7 @@ def _noop_create_stages(self, server_args): def _safe_init(self, **kwargs): # Call ComposedPipelineBase.__init__ directly, skipping LoRAPipeline # which tries to access self.modules['transformer'] - from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) ComposedPipelineBase.__init__(self, **kwargs) @@ -155,8 +155,8 @@ def build_req( **extra_fields, ) -> "Req": """Construct a minimal SGLang ``Req`` for running pipeline stages.""" - from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req - from sglang.multimodal_gen.configs.sample.base import DataType + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req + from sglang.multimodal_gen.configs.sample.sampling_params import DataType req = Req( data_type=DataType.VIDEO, @@ -238,49 +238,14 @@ def close(self): self._ctx.term() -class StageWorkerPool: - """Pool of StageClients for one stage with round-robin dispatch. - - A per-pool asyncio.Lock serialises access so that at most one request - is processed per stage at a time. This avoids concurrent NIXL RDMA - operations within the same stage (which trigger UCX race conditions) - while still allowing **pipeline parallelism** across stages — i.e. - Encoder, Denoiser, and VAE can each be active simultaneously on - different requests. - - Set ``CONCURRENCY >= 3`` to keep the 3-stage pipeline fully saturated. - """ - - def __init__(self, clients: List[StageClient], name: str = ""): - self._clients = clients - self._name = name - self._counter = 0 - self._lock = asyncio.Lock() - - @property - def num_workers(self) -> int: - return len(self._clients) - - async def forward(self, reqs): - """Round-robin dispatch to next available worker (serialised).""" - async with self._lock: - idx = self._counter % len(self._clients) - self._counter += 1 - return await self._clients[idx].forward(reqs) - - def close(self): - for c in self._clients: - c.close() - - def patch_hunyuan_config(): """HunyuanConfig inherits ``task_type`` from PipelineConfig without a default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply ``task_type=T2V`` when omitted. Idempotent. """ - from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType + from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType try: - from sglang.multimodal_gen.configs.pipelines.hunyuan import ( + from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( HunyuanConfig, FastHunyuanConfig, ) except ImportError: @@ -388,5 +353,8 @@ def launch_stage_server(model_path, required_modules, custom_stages_fn, custom_stages_fn=custom_stages_fn, ) - client = StageClient(server_args.scheduler_endpoint(), client_name) + endpoint = server_args.scheduler_endpoint + if callable(endpoint): + endpoint = endpoint() + client = StageClient(endpoint, client_name) return processes, client, server_args diff --git a/examples/disagg_diffusion/phase1_workers/vae_worker.py b/examples/disagg_diffusion/workers/vae_worker.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/vae_worker.py rename to examples/disagg_diffusion/workers/vae_worker.py From de76961f7fd460cde8d2ffccfc477d72912c7d78 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Tue, 17 Mar 2026 21:15:45 +0800 Subject: [PATCH 15/22] fix: adapt NixlSendStage + orchestrator for sglang 0.5.8 API sglang 0.5.8 removed Req.logging_info and OutputBatch's logging_info parameter, and gpu_worker.execute_forward() now expects OutputBatch.timings to be set (accesses timings.total_duration_ms). Changes: - NixlSendStage: drop logging_info usage, provide RequestTimings in every OutputBatch so gpu_worker can record total_duration_ms - orchestrator: use round_robin() dispatch, dynamic per-stage semaphores based on discovered worker count Verified: 4/4 requests (1 single + 3 concurrent) produce valid mp4 videos via NIXL RDMA end-to-end. Co-Authored-By: Claude Opus 4.6 --- .../orchestrator/run_disagg.py | 23 +++++++++++------- .../workers/partial_gpu_worker.py | 24 ++++++++++++------- 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/examples/disagg_diffusion/orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py index b2a316a5e3a3..5c6dc627e809 100755 --- a/examples/disagg_diffusion/orchestrator/run_disagg.py +++ b/examples/disagg_diffusion/orchestrator/run_disagg.py @@ -38,7 +38,7 @@ import uvloop -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "phase1_workers")) +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "workers")) from protocol import ( # noqa: E402 DenoiserRequest, EncoderRequest, VAEDecodeRequest, @@ -56,7 +56,7 @@ async def call_stage(client, request_json: str) -> dict: result = None - stream = await client.generate(request_json) + stream = await client.round_robin(request_json) async for chunk in stream: data = chunk.data() if hasattr(chunk, "data") else chunk if isinstance(data, str): @@ -155,16 +155,23 @@ async def worker(runtime: DistributedRuntime): await encoder_client.wait_for_instances() await denoiser_client.wait_for_instances() await vae_client.wait_for_instances() - logger.info("All 3 stage workers connected") + + # Discover registered worker instances per stage + enc_ids = encoder_client.instance_ids() + den_ids = denoiser_client.instance_ids() + vae_ids = vae_client.instance_ids() + n_enc, n_den, n_vae = len(enc_ids) or 1, len(den_ids) or 1, len(vae_ids) or 1 + logger.info("Workers: encoder=%d, denoiser=%d, vae=%d", n_enc, n_den, n_vae) os.makedirs(OUTPUT_DIR, exist_ok=True) stage_sems = { - "encoder": asyncio.Semaphore(1), - "denoiser": asyncio.Semaphore(1), - "vae": asyncio.Semaphore(1), + "encoder": asyncio.Semaphore(n_enc), + "denoiser": asyncio.Semaphore(n_den), + "vae": asyncio.Semaphore(n_vae), } - admission = asyncio.Semaphore(MAX_PIPELINE_DEPTH) + pipeline_depth = MAX_PIPELINE_DEPTH if MAX_PIPELINE_DEPTH > 0 else (n_enc + n_den + n_vae) + admission = asyncio.Semaphore(pipeline_depth) tracker = PipelineTracker() async def run_stage(name: str, request_id: str, client, request_json: str) -> dict: @@ -322,7 +329,7 @@ async def handle_index(http_request: web.Request) -> web.Response: else: raise - logger.info("Server listening on http://%s:%d (pipeline depth=%d)", HOST, bound_port, MAX_PIPELINE_DEPTH) + logger.info("Server listening on http://%s:%d (pipeline depth=%d)", HOST, bound_port, pipeline_depth) logger.info(" GET / <- latest video preview") logger.info(" POST /v1/videos/generations") logger.info(" GET /health") diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py index b3b5a221a96e..3e2decf913e1 100644 --- a/examples/disagg_diffusion/workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -151,16 +151,24 @@ def __init__(self, output_fields: List[str]): self._output_fields = output_fields self._sender = None + @staticmethod + def _make_timings() -> object: + """Create a RequestTimings object compatible with sglang's executor.""" + try: + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import RequestTimings + return RequestTimings(request_id="nixl-send") + except Exception: + return None + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: tensors = self._extract_tensors(batch) - logging_info = batch.logging_info if not tensors: - return OutputBatch(output={}, logging_info=logging_info) + return OutputBatch(output={}, timings=self._make_timings()) from nixl_transfer import NIXL_AVAILABLE if NIXL_AVAILABLE: - return self._nixl_send(tensors, logging_info) - return self._fallback_send(tensors, logging_info) + return self._nixl_send(tensors) + return self._fallback_send(tensors) def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: """Flatten list-valued fields into individual tensors.""" @@ -182,14 +190,14 @@ def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: result[field] = val return result - def _nixl_send(self, tensors: Dict[str, torch.Tensor], logging_info) -> OutputBatch: + def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: from nixl_transfer import NixlTensorSender if self._sender is None: self._sender = NixlTensorSender() meta = self._sender.send(tensors) - return OutputBatch(output={"_nixl_transfer_meta": meta}, logging_info=logging_info) + return OutputBatch(output={"_nixl_transfer_meta": meta}, timings=self._make_timings()) - def _fallback_send(self, tensors: Dict[str, torch.Tensor], logging_info) -> OutputBatch: + def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: """Fallback: send raw tensors via ZMQ pickle.""" # Reconstruct list-valued fields for backward compat output: Dict[str, object] = {} @@ -205,7 +213,7 @@ def _fallback_send(self, tensors: Dict[str, torch.Tensor], logging_info) -> Outp real_key = base[6:] output[real_key] = [idx_map[i] for i in sorted(idx_map)] del output[base] - return OutputBatch(output=output, logging_info=logging_info) + return OutputBatch(output=output, timings=self._make_timings()) # ═══════════════════════════════════════════════════════════════════════ From 3b2bd4fd2aa60644e65f02f7484a38b2c5d15c18 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Wed, 18 Mar 2026 09:51:01 +0800 Subject: [PATCH 16/22] feat: add debug-level NIXL transfer tensor stats logging Log tensor shape/dtype/mean/std/min/max at NIXL send and receive points for diagnosing data transfer issues. Uses logger.debug() so no overhead at default INFO level. Co-Authored-By: Claude Opus 4.6 --- .../workers/partial_gpu_worker.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py index 3e2decf913e1..fabc0c406747 100644 --- a/examples/disagg_diffusion/workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -113,6 +113,18 @@ def _nixl_pull(self, batch: Req, meta: dict) -> Req: reconstructed[k] = v for base, idx_map in indexed.items(): reconstructed[base] = [idx_map[i] for i in sorted(idx_map)] + # Debug: log tensor stats after NIXL pull + for k, v in reconstructed.items(): + if isinstance(v, list): + for i, t in enumerate(v): + if hasattr(t, 'float'): + f = t.float() + logger.debug("NIXL_RECV %s[%d]: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", + k, i, t.shape, t.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) + elif hasattr(v, 'float'): + f = v.float() + logger.debug("NIXL_RECV %s: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", + k, v.shape, v.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) from sglang_utils import inject_tensors_to_req inject_tensors_to_req(batch, reconstructed) return batch @@ -165,6 +177,12 @@ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: if not tensors: return OutputBatch(output={}, timings=self._make_timings()) + # Debug: log tensor stats before NIXL send + for k, t in tensors.items(): + f = t.float() + logger.debug("NIXL_SEND %s: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", + k, t.shape, t.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) + from nixl_transfer import NIXL_AVAILABLE if NIXL_AVAILABLE: return self._nixl_send(tensors) From a8a74651d4d78ad11483d61997c4d6aec07df583 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 12:15:16 +0800 Subject: [PATCH 17/22] docs: add DESIGN.md for disaggregated diffusion pipeline Comprehensive design document covering the architecture, component design, and SGLang/HunyuanVideo implementation details. Sections 1-3 are model-agnostic and backend-agnostic; Section 4 isolates SGLang specifics. Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/DESIGN.md | 518 ++++++++++++++++++++++++++++ 1 file changed, 518 insertions(+) create mode 100644 examples/disagg_diffusion/DESIGN.md diff --git a/examples/disagg_diffusion/DESIGN.md b/examples/disagg_diffusion/DESIGN.md new file mode 100644 index 000000000000..4b3685ea4321 --- /dev/null +++ b/examples/disagg_diffusion/DESIGN.md @@ -0,0 +1,518 @@ +# Disaggregated Diffusion Pipeline — Design Document + +## 1. Overview + +### Motivation + +Modern diffusion pipelines (text-to-video, text-to-image, omni-modal) are +composed of heterogeneous stages — text encoding, iterative denoising, VAE +decoding — each with fundamentally different compute profiles: + +- **Different optimization strategies per stage.** Encoders are + memory-bound single-pass transforms; denoisers are compute-bound + multi-step loops that benefit from tensor parallelism; VAE decoders + are memory-intensive but run only once. Forcing all three into a + single process prevents stage-specific tuning (parallelism, batching, + memory management, quantization). + +- **Shifting compute balance.** As diffusion models mature, the DiT no + longer dominates the entire pipeline — faster denoisers (fewer steps, + distilled models) shift the bottleneck to encoding and decoding. + Multi-task pipelines are emerging (e.g. OneVideo: encode + denoise + + decode + audio in one model) where each task demands independent + scaling. + +- **Omni-modal future.** Models that jointly produce video, audio, + image, and text require stage-level separation so each modality's + compute can scale independently without wasting GPU resources. + +### Solution + +Decompose the pipeline into **N independent stages**, each running as a +Dynamo RPC worker on dedicated GPU(s). An orchestrator chains stages +together, using NIXL RDMA for GPU-direct tensor transfer between them. + +| Goal | Mechanism | +|---|---| +| Stage-level scaling | N workers per stage, auto-discovered via etcd | +| Pipeline parallelism | Semaphore admission, independent worker pools | +| GPU-direct transfer | NIXL RDMA — only ~1.5 KB metadata over RPC | +| Loose coupling | Workers are independent processes, no shared state | +| Dynamic scaling | Add/remove workers at runtime, no restart | +| Auto routing | Idle-queue dispatch, backpressure, retry on failure | + + +## 2. Architecture + +### 2.1 Architecture Diagram + +``` + ┌───────────┐ + │ etcd │ + │ registry │ + └─────┬─────┘ + register │ discover + ┌───────────────────────────────┼───────────────────────────────┐ + │ Orchestrator │ + │ │ + │ HTTP ──► handle_generate() ──► dispatch_with_retry() │ + │ PipelineTracker per-stage WorkerManager │ + │ Semaphore(depth) acquire → direct() → release │ + │ │ + └──────┬──────────────────────┬──────────────────────┬──────────┘ + │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC + │ (JSON ~1 KB) │ (NIXL meta ~1.5 KB) │ (NIXL meta ~1.5 KB) + ▼ ▼ ▼ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-0 │──────────│Denoiser-0│────────────│ VAE-0 │ + │ GPU 0 │ NIXL │ GPU 1,2 │ NIXL │ GPU 5 │ + └──────────┘ RDMA └──────────┘ RDMA └──────────┘ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-1 │──────────│Denoiser-1│────────────│ VAE-1 │ + │ GPU 3 │ NIXL │ GPU 4,5 │ NIXL │ GPU 7 │ + └──────────┘ RDMA └──────────┘ RDMA └──────────┘ + ...↕ ...↕ ...↕ + dynamic dynamic dynamic + add/remove add/remove add/remove +``` + +**Key points:** +- Workers self-register with etcd; the orchestrator discovers them + automatically. +- Dynamo RPC carries only small JSON payloads and NIXL metadata + (~1.5 KB). Actual tensors (embeddings, latents) transfer GPU-to-GPU + via NIXL RDMA, never touching the CPU or RPC channel. +- Each stage can have a different number of workers and GPU count. + The denoiser typically uses TP > 1 (multi-GPU), while encoder and VAE + each use a single GPU. + +### 2.2 Process Model + +Each worker is an independent OS process with no shared state: + +``` +Dynamo Worker Process (e.g. encoder_worker.py) +├── @dynamo_worker ← Dynamo runtime bootstrap +│ ├── serve_endpoint("generate") ← Dynamo RPC from orchestrator +│ │ └── handle_generate() +│ │ └── StageClient.forward() ← ZMQ to local backend subprocess +│ └── serve_endpoint("health") +│ +└── Backend subprocess (spawned at startup) + └── Inference engine (any backend) + ├── ReceiveStage ← RDMA-pull tensors from previous stage + ├── ComputeStage ← model-specific inference + └── SendStage ← register output tensors as RDMA-readable +``` + +- The Dynamo worker process handles RPC and control-plane logic. +- The backend subprocess runs the actual model inference. It can be any + inference engine — SGLang, vLLM, a custom PyTorch loop, etc. +- Communication between the two is via ZMQ REQ/REP (same-host IPC). + +### 2.3 Loose Coupling & Dynamic Scaling + +Workers are completely independent — they know nothing about each other +or the orchestrator: + +- **Startup:** Worker process starts → registers with etcd (via Dynamo + runtime) → orchestrator auto-discovers the new instance. +- **Shutdown:** Worker process exits → etcd lease expires → + orchestrator stops routing to it. +- **Add worker:** Start a new process on any available GPU → etcd → + orchestrator sees it within seconds. No restart, no reconfiguration. +- **Remove worker:** Kill the process → etcd lease expires → traffic + drains naturally. +- **Auto-scale:** An external controller can monitor queue depth per + stage (`GET /pipeline/status`) and spawn/kill workers as needed. + +### 2.4 Auto Routing + +Each stage has a `WorkerManager` that maintains an idle-pool queue: + +``` +WorkerManager("denoiser", client, [0, 1, 2]) +│ +├── _idle_queue: asyncio.Queue ← [0, 1, 2] initially +│ +├── acquire_worker() → int ← blocks until a worker is idle +├── dispatch(wid, rid, json) ← client.direct(json, wid) +│ └── on completion/failure → release wid back to _idle_queue +│ +└── status() → dict ← per-worker completed/latency/state +``` + +- `acquire_worker()` blocks the caller when all workers are busy + (backpressure). +- `dispatch()` sends to a specific worker via `client.direct()` and + tracks busy/idle state for observability. +- `dispatch_with_retry()` wraps this: on failure, acquires a different + worker and retries (configurable via `STAGE_DISPATCH_RETRIES`). + +### 2.5 Request Flow + +A single video generation request follows this path: + +``` +Client Orchestrator Encoder-k Denoiser-j VAE-i + │ │ │ │ │ + │─POST /v1/videos/──────►│ │ │ │ + │ generations │ │ │ │ + │ [acquire semaphore] │ │ │ + │ │ │ │ │ + │ │──EncoderRequest──────►│ │ │ + │ │ (prompt, cfg) [encode text] │ │ + │ │◄──NIXL metadata──────│ │ │ + │ │ (~1.5 KB) │ │ │ + │ │ │ │ │ + │ │──DenoiserRequest───────────────────►│ │ + │ │ (NIXL meta, params) [RDMA pull embeddings]│ + │ │ [denoise N steps] │ + │ │◄──NIXL metadata────────────────────│ │ + │ │ │ │ + │ │──VAEDecodeRequest──────────────────────────────►│ + │ │ (NIXL meta, req_id) [RDMA pull] │ + │ │ [decode] │ + │ │◄──{video_path}──────────────────────────────────│ + │ [release semaphore] │ │ │ + │◄──{url, timings}───────│ │ │ │ +``` + +Only ~1.5 KB of NIXL metadata travels over Dynamo RPC between stages. +The actual tensor data (embeddings: ~tens of MB, latents: ~hundreds of +MB) transfers GPU-to-GPU via NIXL RDMA without CPU involvement. + +### 2.6 Pipeline Parallelism + +An `asyncio.Semaphore(pipeline_depth)` gates admission. Each stage has +its own independent worker pool, so multiple requests overlap: + +``` +Time ──────────────────────────────────────────────────► + +Req A: [Enc-0][======Den-0======][VAE-0] +Req B: [Enc-1][======Den-1======][VAE-1] +Req C: [Enc-0][======Den-0======][VAE-0] +Req D: [Enc-1][======Den-1======][VAE-1] +``` + +- `pipeline_depth` defaults to `MAX_PIPELINE_DEPTH` (env, default 4) + or the total number of workers across all stages. +- Each request independently acquires workers from each stage's pool. +- The denoiser is typically the bottleneck (50 diffusion steps), so + encoder and VAE workers are freed quickly to serve other requests. + + +## 3. Component Design + +### 3.1 Orchestrator + +`orchestrator/run_disagg.py` — aiohttp HTTP server that chains stages. + +**Endpoints:** + +| Method | Path | Description | +|---|---|---| +| `POST` | `/v1/videos/generations` | Submit generation request | +| `GET` | `/health` | Orchestrator liveness | +| `GET` | `/health/stages` | Per-stage health (queries each worker) | +| `GET` | `/pipeline/status` | Active requests, queue depth, latencies | +| `GET` | `/videos/` | Serve generated video file | + +**Key components:** + +- `PipelineTracker` — Tracks active requests per stage, completed/failed + counts, and rolling average stage latencies. +- `dispatch_with_retry(mgr, request_id, json)` — Acquire worker → + dispatch → on failure retry on a different worker (up to + `STAGE_DISPATCH_RETRIES` attempts). +- `admission = asyncio.Semaphore(pipeline_depth)` — Limits concurrent + in-flight requests across the entire pipeline. + +### 3.2 WorkerManager + +`orchestrator/worker_manager.py` — Per-stage worker pool with +busy/idle tracking. + +```python +class WorkerManager: + def __init__(self, stage_name: str, client, worker_ids: List[int]): ... + async def acquire_worker(self) -> int: ... # blocks until idle + async def dispatch(self, worker_id, request_id, request_json) -> (dict, float): ... + def status(self) -> dict: ... # per-worker stats +``` + +- Backed by `asyncio.Queue` (idle pool) — `acquire_worker()` awaits + the queue, `dispatch()` returns the worker to it on completion. +- `_call_direct()` sends to a specific worker via + `client.direct(json, worker_id)`. +- `status()` returns per-worker `{id, status, request_id, completed, + avg_latency_s}` plus stage-level `{queue_depth, completed, failed}`. + +### 3.3 Worker Interface + +All three workers follow an identical pattern: + +```python +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + # 1. Launch backend subprocess + processes, client, server_args = launch_stage_server( + MODEL_PATH, required_modules, build_stage_fn, SCHEDULER_PORT, + ) + + # 2. Define Dynamo RPC handlers + async def handle_generate(request, context): + output = await client.forward([build_req(...)]) + yield result_dict # JSON response with NIXL metadata or output + + async def handle_health(request, context): + yield {"status": "ok", "stage": "..."} + + # 3. Serve endpoints + gen_ep = runtime.endpoint("disagg_diffusion..generate") + health_ep = runtime.endpoint("disagg_diffusion..health") + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) +``` + +Each worker: +- Spawns a backend subprocess (which loads model weights and runs + inference). +- Bridges Dynamo RPC ↔ backend via `StageClient` (async ZMQ). +- Handles NIXL metadata forwarding: output from one stage's send + becomes the next stage's receive metadata. +- Includes a ZMQ fallback path: when NIXL is unavailable, tensors are + serialized via `torch.save()` + base64 over the RPC channel. + +**Stage-specific behavior:** + +| Stage | Input | Compute | Output | +|---|---|---|---| +| Encoder | prompt text | Text encoding | NIXL metadata (embeddings) | +| Denoiser | NIXL metadata (embeddings) + params | RDMA pull → N denoise steps | NIXL metadata (latents) | +| VAE | NIXL metadata (latents) | RDMA pull → VAE decode | video file path | + +### 3.4 Tensor Transfer (NIXL) + +`workers/nixl_transfer.py` — GPU-direct RDMA transfer between stages. + +**Sender (`NixlTensorSender`):** + +```python +sender = NixlTensorSender() +meta = sender.send({"latents": tensor}) # → ~1.5 KB metadata dict +``` + +1. Flatten all tensors into a single contiguous GPU buffer + (`torch.cat`). +2. Create a NIXL `Descriptor` wrapping the flat buffer. +3. Register as `readable` via `connector.create_readable(descriptor)`. +4. Return metadata: tensor keys, shapes, dtypes, NIXL descriptor. +5. Hold `(readable, flat_buffer, timestamp)` in `_pending` list. +6. `_sweep()` polls `readable.status` on each subsequent `send()` — + releases buffers on `COMPLETE` or timeout (`NIXL_BUFFER_TIMEOUT_S`, + default 120s). + +**Receiver (`NixlTensorReceiver`):** + +```python +receiver = NixlTensorReceiver() +tensors = receiver.recv(meta, device="cuda") # → {"latents": tensor} +``` + +1. Parse metadata to compute total byte size and per-tensor specs. +2. Allocate a flat `torch.uint8` buffer directly on the target GPU. +3. `connector.begin_read(rdma_meta, descriptor)` → RDMA pull from + sender's GPU. +4. `read_op.wait_for_completion()` — blocks until transfer finishes. +5. Slice the flat buffer into individual tensors using stored shapes + and dtypes. + +### 3.5 Protocol Types + +`workers/protocol.py` — Pydantic models for Dynamo RPC serialization. + +```python +class EncoderRequest(BaseModel): + prompt: str + negative_prompt: str = "" + guidance_scale: float = 1.0 + +class DenoiserRequest(BaseModel): + transfer_meta: Dict[str, Any] # NIXL metadata from encoder + tensor_data: Dict[str, Any] = {} # ZMQ fallback + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + seed: int = 42 + +class VAEDecodeRequest(BaseModel): + transfer_meta: Dict[str, Any] # NIXL metadata from denoiser + tensor_data: Dict[str, Any] = {} # ZMQ fallback + request_id: str = "" +``` + +Responses carry either `transfer_meta` (NIXL path) or `tensor_data` +(ZMQ fallback) but never both. + + +## 4. Implementation: SGLang Backend + +This section describes the current backend implementation using SGLang's +multimodal generation runtime, with HunyuanVideo as the reference model. +The generic architecture (Sections 1-3) is backend-agnostic — any +inference engine that can run pipeline stages can replace SGLang. + +### 4.1 PartialGPUWorker + +`workers/partial_gpu_worker.py` — Extends SGLang's `GPUWorker` to load +only the modules each stage needs. + +```python +class PartialGPUWorker(GPUWorker): + def __init__(self, required_modules, custom_stages_fn, **kwargs): ... + def init_device_and_model(self): + # 1. Set up distributed environment (TP, SP, CFG parallel) + # 2. build_partial_pipeline() — load only required_modules + # 3. custom_stages_fn(pipeline, server_args) — build stage list + # 4. Register stages with pipeline +``` + +- Overrides only `init_device_and_model()`. All other `GPUWorker` + behavior (forward execution, memory analysis, LoRA) is inherited. +- `build_partial_pipeline()` (`sglang_utils.py`) dynamically creates a + subclass of the model's pipeline that suppresses automatic stage + creation and LoRA initialization, loading only the specified modules. + +**Custom pipeline stages (NIXL integration):** + +- `NixlReceiveStage(PipelineStage)` — Prepended at the start of + denoiser/VAE pipelines. Reads `_nixl_transfer_meta` from the `Req` + and RDMA-pulls tensors. Falls back to device-move for ZMQ path. + Includes retry logic for `REMOTE_DISCONNECT` errors. +- `NixlSendStage(PipelineStage)` — Appended as the last stage in + encoder/denoiser pipelines. Extracts tensors from `Req`, registers + with NIXL, returns `OutputBatch` containing only metadata. + +**Stage builders** (picklable functions passed to subprocess): + +| Function | Stages | +|---|---| +| `build_encoder_stages()` | `TextEncodingStage` → `NixlSendStage` | +| `build_denoiser_stages()` | `NixlReceiveStage` → `LatentPreparationStage` → `TimestepPreparationStage` → `DenoisingStage` → `NixlSendStage` | +| `build_vae_stages()` | `NixlReceiveStage` → `DecodingStage` | + +### 4.2 Subprocess Launcher + +`workers/partial_gpu_worker.py:launch_partial_server()` — Spawns +SGLang Scheduler subprocess(es) with `PartialGPUWorker` monkey-patched +in place of the default `GPUWorker`. + +``` +launch_partial_server(server_args, required_modules, custom_stages_fn) +│ +├── For each GPU (rank 0..N-1): +│ ├── Create readiness pipe +│ ├── mp.Process(target=_run_partial_scheduler_process) +│ │ ├── Monkey-patch: sched_mod.GPUWorker = _PatchedGPUWorker +│ │ └── run_scheduler_process(...) ← standard SGLang entry point +│ └── Start process +│ +├── Wire master/slave pipes (TP > 1: rank 0 is master, ranks 1..N are slaves) +├── Wait for all readiness signals +└── Return process list +``` + +`launch_stage_server()` (`sglang_utils.py`) wraps this with config +setup: `patch_hunyuan_config()` → `ServerArgs.from_kwargs()` → +`launch_partial_server()` → `StageClient(endpoint)`. + +### 4.3 StageClient + +`workers/sglang_utils.py:StageClient` — Async ZMQ REQ/REP client +connecting the Dynamo worker main process to the SGLang Scheduler +subprocess. + +```python +class StageClient: + def __init__(self, endpoint: str, name: str = ""): ... + async def forward(self, reqs): # send_pyobj → recv_pyobj with timeout + def close(self): ... +``` + +- `asyncio.Lock` serializes concurrent calls (ZMQ REQ socket is + single-flight). +- Configurable timeout: `STAGE_FORWARD_TIMEOUT_S` (default 120s). + +### 4.4 HunyuanVideo Specifics + +- **Dual encoder detection:** `detect_encoder_modules()` reads + `model_index.json` and auto-detects `text_encoder_2` / `tokenizer_2` + (HunyuanVideo uses Llama + CLIP). Falls back to heuristic for known + model names. +- **HunyuanConfig patching:** `patch_hunyuan_config()` wraps + `HunyuanConfig.__init__` to supply `task_type=T2V` when omitted + (the base class requires it but HunyuanConfig doesn't default it). +- **Triton norm contiguous workaround:** + `_patch_triton_norm_contiguous()` wraps SGLang's triton + `norm_infer` to call `.contiguous()` on non-contiguous tensors from + HunyuanVideo's attention reshapes, avoiding the triton kernel + assertion `x.stride(-1) == 1`. +- **Component config sync:** `_sync_all_component_configs()` reads + `config.json` for every component in `model_index.json` and updates + `server_args.pipeline_config`, ensuring correct parameters (e.g. + `z_dim`) even for components whose weights are not loaded by the + current stage. + + +## 5. Roadmap + +- [x] Multi-worker scaling — N workers per stage, etcd auto-discovery +- [x] Pipeline parallelism — overlapping requests across stages +- [x] NIXL GPU-direct transfer — GPU-to-GPU RDMA, only metadata over RPC +- [ ] Runtime scaling — add/remove workers without restart; + external auto-scaler integration based on queue depth +- [ ] Smart routing — load-aware dispatch (not just idle-queue), + affinity-based routing, request priority +- [ ] Fault tolerance — worker health-check + eviction, + dead worker detection, graceful degradation +- [ ] Streaming output — stream decoded frames as produced +- [ ] Metrics & observability — per-stage latency histograms, + GPU utilization, NIXL throughput, Prometheus export +- [ ] Request cancellation — cancel in-flight requests, free GPU + immediately +- [ ] Multi-model support — OneVideo, omni-modal pipelines, + heterogeneous stage graphs (not just linear 3-stage) + + +## 6. File Map + +``` +examples/disagg_diffusion/ +├── DESIGN.md ← this document +├── README.md ← usage guide and quick start +├── run_all.sh ← launch etcd + all workers + orchestrator +├── stress_test.sh ← concurrent load testing script +│ +├── orchestrator/ +│ ├── run_disagg.py ← HTTP server, PipelineTracker, dispatch_with_retry +│ └── worker_manager.py ← WorkerManager: idle-pool, acquire/dispatch/status +│ +├── workers/ +│ ├── __init__.py +│ ├── protocol.py ← Pydantic request/response models (Dynamo RPC) +│ ├── encoder_worker.py ← Encoder Dynamo worker (text encoding → NIXL send) +│ ├── denoiser_worker.py ← Denoiser Dynamo worker (NIXL recv → denoise → NIXL send) +│ ├── vae_worker.py ← VAE Dynamo worker (NIXL recv → decode → save video) +│ ├── nixl_transfer.py ← NixlTensorSender / NixlTensorReceiver (RDMA) +│ ├── partial_gpu_worker.py ← PartialGPUWorker, NixlSend/ReceiveStage, subprocess launcher +│ └── sglang_utils.py ← StageClient, build_partial_pipeline, launch_stage_server +│ +└── validate/ + └── validate_split.py ← validation script for split correctness +``` From 84feca9febd92a43f76e2facd643035914448edc Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 12:18:03 +0800 Subject: [PATCH 18/22] refactor: per-worker targeted dispatch, ZMQ fallback, sglang 0.5.8 compat Replace round-robin semaphore-based orchestrator with WorkerManager that tracks per-worker busy/idle state and dispatches via client.direct(). Add dispatch_with_retry for automatic failover to a different worker. - Add ZMQ tensor fallback path (base64-encoded torch.save) for when NIXL is unavailable; controlled via DISABLE_NIXL env var - Fix sglang 0.5.8 API: runtime.endpoint() replaces ns.component(), OutputBatch(metrics=) replaces timings=, fix add_stage arg order - Add NIXL pull retry with exponential backoff for REMOTE_DISCONNECT - Fix NIXL sender buffer leak: replace broken async _keep_alive with synchronous _sweep() that polls readable.status - Add StageClient.forward() timeout (STAGE_FORWARD_TIMEOUT_S) - Add stress_test.sh for concurrent load testing Co-Authored-By: Claude Opus 4.6 --- .../orchestrator/run_disagg.py | 167 +++++++----- .../orchestrator/worker_manager.py | 141 ++++++++++ examples/disagg_diffusion/run_all.sh | 2 +- examples/disagg_diffusion/stress_test.sh | 251 ++++++++++++++++++ .../workers/denoiser_worker.py | 36 ++- .../workers/encoder_worker.py | 27 +- .../disagg_diffusion/workers/nixl_transfer.py | 58 ++-- .../workers/partial_gpu_worker.py | 45 ++-- examples/disagg_diffusion/workers/protocol.py | 2 + .../disagg_diffusion/workers/sglang_utils.py | 11 +- .../disagg_diffusion/workers/vae_worker.py | 20 +- 11 files changed, 641 insertions(+), 119 deletions(-) create mode 100644 examples/disagg_diffusion/orchestrator/worker_manager.py create mode 100755 examples/disagg_diffusion/stress_test.sh diff --git a/examples/disagg_diffusion/orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py index 5c6dc627e809..660f7c39070e 100755 --- a/examples/disagg_diffusion/orchestrator/run_disagg.py +++ b/examples/disagg_diffusion/orchestrator/run_disagg.py @@ -8,8 +8,10 @@ Encoder → Denoiser → VAE via Dynamo RPC + NIXL RDMA. Pipeline parallelism: multiple requests can be in different stages -simultaneously. Per-stage semaphores control backpressure so each -GPU processes one request at a time, while the pipeline stays full. +simultaneously. Per-stage WorkerManagers track busy/idle state for +each GPU worker and dispatch requests to specific instances via +``client.direct()``, enabling backpressure-aware scheduling and +per-worker observability through ``/pipeline/status``. Each stage worker wraps an SGLang Scheduler subprocess via launch_partial_server(), supporting TP for the denoiser and NIXL RDMA @@ -45,6 +47,7 @@ HealthRequest, ) from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 +from worker_manager import WorkerManager # noqa: E402 logger = logging.getLogger(__name__) @@ -52,19 +55,7 @@ HOST = os.environ.get("HOST", "0.0.0.0") OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_videos") MAX_PIPELINE_DEPTH = int(os.environ.get("MAX_PIPELINE_DEPTH", "4")) - - -async def call_stage(client, request_json: str) -> dict: - result = None - stream = await client.round_robin(request_json) - async for chunk in stream: - data = chunk.data() if hasattr(chunk, "data") else chunk - if isinstance(data, str): - data = json.loads(data) - result = data - if result is None: - raise RuntimeError("Empty response from stage") - return result +STAGE_DISPATCH_RETRIES = int(os.environ.get("STAGE_DISPATCH_RETRIES", "2")) async def query_stage_health(health_client, stage_name: str) -> dict: @@ -136,14 +127,13 @@ async def status(self) -> dict: @dynamo_worker(enable_nats=False) async def worker(runtime: DistributedRuntime): - ns = runtime.namespace("disagg_diffusion") - encoder_client = await ns.component("encoder").endpoint("generate").client() - denoiser_client = await ns.component("denoiser").endpoint("generate").client() - vae_client = await ns.component("vae").endpoint("generate").client() + encoder_client = await runtime.endpoint("disagg_diffusion.encoder.generate").client() + denoiser_client = await runtime.endpoint("disagg_diffusion.denoiser.generate").client() + vae_client = await runtime.endpoint("disagg_diffusion.vae.generate").client() - encoder_health_client = await ns.component("encoder").endpoint("health").client() - denoiser_health_client = await ns.component("denoiser").endpoint("health").client() - vae_health_client = await ns.component("vae").endpoint("health").client() + encoder_health_client = await runtime.endpoint("disagg_diffusion.encoder.health").client() + denoiser_health_client = await runtime.endpoint("disagg_diffusion.denoiser.health").client() + vae_health_client = await runtime.endpoint("disagg_diffusion.vae.health").client() health_clients = { "encoder": encoder_health_client, @@ -160,30 +150,43 @@ async def worker(runtime: DistributedRuntime): enc_ids = encoder_client.instance_ids() den_ids = denoiser_client.instance_ids() vae_ids = vae_client.instance_ids() - n_enc, n_den, n_vae = len(enc_ids) or 1, len(den_ids) or 1, len(vae_ids) or 1 - logger.info("Workers: encoder=%d, denoiser=%d, vae=%d", n_enc, n_den, n_vae) + logger.info( + "Workers: encoder=%d %s, denoiser=%d %s, vae=%d %s", + len(enc_ids), enc_ids, len(den_ids), den_ids, len(vae_ids), vae_ids, + ) os.makedirs(OUTPUT_DIR, exist_ok=True) - stage_sems = { - "encoder": asyncio.Semaphore(n_enc), - "denoiser": asyncio.Semaphore(n_den), - "vae": asyncio.Semaphore(n_vae), + managers: Dict[str, WorkerManager] = { + "encoder": WorkerManager("encoder", encoder_client, enc_ids), + "denoiser": WorkerManager("denoiser", denoiser_client, den_ids), + "vae": WorkerManager("vae", vae_client, vae_ids), } - pipeline_depth = MAX_PIPELINE_DEPTH if MAX_PIPELINE_DEPTH > 0 else (n_enc + n_den + n_vae) + n_workers_total = sum(m.worker_count for m in managers.values()) + pipeline_depth = MAX_PIPELINE_DEPTH if MAX_PIPELINE_DEPTH > 0 else n_workers_total admission = asyncio.Semaphore(pipeline_depth) tracker = PipelineTracker() - async def run_stage(name: str, request_id: str, client, request_json: str) -> dict: - """Run a single stage with semaphore gating and tracking.""" - async with stage_sems[name]: - await tracker.enter(request_id, name) - t0 = time.monotonic() - result = await call_stage(client, request_json) - elapsed = time.monotonic() - t0 - await tracker.leave(request_id, name, elapsed) - logger.info("[%s] %s: %.2fs", request_id, name.capitalize(), elapsed) - return result, elapsed + async def dispatch_with_retry( + mgr: WorkerManager, request_id: str, request_json: str, + ) -> tuple: + """Dispatch to a stage worker; on failure retry on a different worker.""" + for attempt in range(STAGE_DISPATCH_RETRIES + 1): + wid = await mgr.acquire_worker() + try: + result, elapsed = await mgr.dispatch(wid, request_id, request_json) + if "error" in result and result["error"]: + raise RuntimeError(result["error"]) + return result, elapsed, wid + except Exception as e: + logger.warning( + "[%s] %s worker %d failed (attempt %d/%d): %s", + request_id, mgr.stage_name, wid, attempt + 1, + STAGE_DISPATCH_RETRIES + 1, e, + ) + if attempt >= STAGE_DISPATCH_RETRIES: + raise + raise RuntimeError("unreachable") async def handle_generate(request: dict) -> dict: request_id = str(uuid.uuid4())[:8] @@ -191,35 +194,53 @@ async def handle_generate(request: dict) -> dict: timings: Dict[str, float] = {} async with admission: - enc_req = EncoderRequest( - prompt=request["prompt"], - negative_prompt=request.get("negative_prompt", ""), - guidance_scale=request.get("guidance_scale", 1.0), - ) - enc_resp, timings["encoder_s"] = await run_stage( - "encoder", request_id, encoder_client, enc_req.model_dump_json(), - ) - - den_req = DenoiserRequest( - transfer_meta=enc_resp["transfer_meta"], - height=request.get("height", 544), - width=request.get("width", 960), - num_frames=request.get("num_frames", 61), - num_inference_steps=request.get("num_inference_steps", 50), - guidance_scale=request.get("guidance_scale", 1.0), - seed=seed, - ) - den_resp, timings["denoiser_s"] = await run_stage( - "denoiser", request_id, denoiser_client, den_req.model_dump_json(), - ) - - vae_req = VAEDecodeRequest( - transfer_meta=den_resp["transfer_meta"], - request_id=request_id, - ) - vae_resp, timings["vae_s"] = await run_stage( - "vae", request_id, vae_client, vae_req.model_dump_json(), - ) + try: + # Stage 1: Encoder + enc_req = EncoderRequest( + prompt=request["prompt"], + negative_prompt=request.get("negative_prompt", ""), + guidance_scale=request.get("guidance_scale", 1.0), + ) + await tracker.enter(request_id, "encoder") + enc_resp, timings["encoder_s"], enc_wid = await dispatch_with_retry( + managers["encoder"], request_id, enc_req.model_dump_json(), + ) + await tracker.leave(request_id, "encoder", timings["encoder_s"]) + logger.info("[%s] Encoder (worker %d): %.2fs", request_id, enc_wid, timings["encoder_s"]) + + # Stage 2: Denoiser + den_req = DenoiserRequest( + transfer_meta=enc_resp.get("transfer_meta", {}), + tensor_data=enc_resp.get("tensor_data", {}), + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=request.get("num_inference_steps", 50), + guidance_scale=request.get("guidance_scale", 1.0), + seed=seed, + ) + await tracker.enter(request_id, "denoiser") + den_resp, timings["denoiser_s"], den_wid = await dispatch_with_retry( + managers["denoiser"], request_id, den_req.model_dump_json(), + ) + await tracker.leave(request_id, "denoiser", timings["denoiser_s"]) + logger.info("[%s] Denoiser (worker %d): %.2fs", request_id, den_wid, timings["denoiser_s"]) + + # Stage 3: VAE + vae_req = VAEDecodeRequest( + transfer_meta=den_resp.get("transfer_meta", {}), + tensor_data=den_resp.get("tensor_data", {}), + request_id=request_id, + ) + await tracker.enter(request_id, "vae") + vae_resp, timings["vae_s"], vae_wid = await dispatch_with_retry( + managers["vae"], request_id, vae_req.model_dump_json(), + ) + await tracker.leave(request_id, "vae", timings["vae_s"]) + logger.info("[%s] VAE (worker %d): %.2fs", request_id, vae_wid, timings["vae_s"]) + except Exception: + await tracker.mark_failed(request_id) + raise timings["total_s"] = round(sum(timings.values()), 3) await tracker.mark_done(request_id) @@ -269,8 +290,14 @@ async def handle_stages_health(http_request: web.Request) -> web.Response: return web.json_response({"stages": list(results)}) async def handle_pipeline_status(http_request: web.Request) -> web.Response: - status = await tracker.status() - return web.json_response(status) + pipeline = await tracker.status() + stages = {name: mgr.status() for name, mgr in managers.items()} + pipeline["stages"] = stages + pipeline["pipeline_depth"] = { + "active": pipeline["active_count"], + "max": pipeline_depth, + } + return web.json_response(pipeline) async def handle_video(http_request: web.Request) -> web.Response: filename = http_request.match_info["filename"] diff --git a/examples/disagg_diffusion/orchestrator/worker_manager.py b/examples/disagg_diffusion/orchestrator/worker_manager.py new file mode 100644 index 000000000000..14ca2527a1ed --- /dev/null +++ b/examples/disagg_diffusion/orchestrator/worker_manager.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-stage worker pool manager with busy/idle tracking and targeted dispatch. + +Each pipeline stage (encoder, denoiser, vae) gets its own WorkerManager +that tracks which Dynamo worker instances are idle or busy and dispatches +requests to specific workers via ``client.direct()``. This replaces the +previous ``asyncio.Semaphore``-based flow control with explicit state that +is queryable via ``/pipeline/status``. +""" + +import asyncio +import json +import logging +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple + +logger = logging.getLogger(__name__) + + +@dataclass +class WorkerState: + """Runtime state for a single Dynamo worker instance.""" + + worker_id: int + status: Literal["idle", "busy"] = "idle" + current_request_id: Optional[str] = None + completed_count: int = 0 + total_latency_s: float = 0.0 + + +class WorkerManager: + """Manages the worker pool for one pipeline stage. + + Provides acquire/dispatch semantics: callers first ``acquire_worker()`` + (blocks until one is idle), then ``dispatch()`` to send the request to + that specific worker. On completion (or failure) the worker is + automatically returned to the idle pool. + """ + + def __init__(self, stage_name: str, client: Any, worker_ids: List[int]): + self._stage_name = stage_name + self._client = client + self._workers: Dict[int, WorkerState] = { + wid: WorkerState(worker_id=wid) for wid in worker_ids + } + self._idle_queue: asyncio.Queue[int] = asyncio.Queue() + for wid in worker_ids: + self._idle_queue.put_nowait(wid) + + # Observability counters + self._total_completed = 0 + self._total_failed = 0 + self._waiting_count = 0 # requests blocked waiting for an idle worker + + @property + def stage_name(self) -> str: + return self._stage_name + + @property + def worker_count(self) -> int: + return len(self._workers) + + async def acquire_worker(self) -> int: + """Block until a worker becomes idle. Returns worker_id.""" + self._waiting_count += 1 + try: + wid = await self._idle_queue.get() + finally: + self._waiting_count -= 1 + return wid + + async def dispatch( + self, worker_id: int, request_id: str, request_json: str + ) -> Tuple[dict, float]: + """Send request to a specific worker via ``client.direct()`` and track state. + + Returns (result_dict, elapsed_seconds). + """ + ws = self._workers[worker_id] + ws.status = "busy" + ws.current_request_id = request_id + t0 = time.monotonic() + try: + result = await self._call_direct(worker_id, request_json) + elapsed = time.monotonic() - t0 + ws.completed_count += 1 + ws.total_latency_s += elapsed + self._total_completed += 1 + return result, elapsed + except Exception: + self._total_failed += 1 + raise + finally: + ws.status = "idle" + ws.current_request_id = None + self._idle_queue.put_nowait(worker_id) + + async def _call_direct(self, worker_id: int, request_json: str) -> dict: + """Issue a Dynamo RPC to a specific worker instance.""" + result = None + stream = await self._client.direct(request_json, worker_id) + async for chunk in stream: + data = chunk.data() if hasattr(chunk, "data") else chunk + if isinstance(data, str): + data = json.loads(data) + result = data + if result is None: + raise RuntimeError( + f"Empty response from {self._stage_name} worker {worker_id}" + ) + return result + + def status(self) -> dict: + """Snapshot of per-worker state for ``/pipeline/status``.""" + workers = [] + for ws in self._workers.values(): + avg = ( + round(ws.total_latency_s / ws.completed_count, 3) + if ws.completed_count + else 0 + ) + workers.append( + { + "id": ws.worker_id, + "status": ws.status, + "request_id": ws.current_request_id, + "completed": ws.completed_count, + "avg_latency_s": avg, + } + ) + return { + "stage": self._stage_name, + "workers": workers, + "queue_depth": self._waiting_count, + "completed": self._total_completed, + "failed": self._total_failed, + } diff --git a/examples/disagg_diffusion/run_all.sh b/examples/disagg_diffusion/run_all.sh index b9fba7bd5015..328eff6d942b 100755 --- a/examples/disagg_diffusion/run_all.sh +++ b/examples/disagg_diffusion/run_all.sh @@ -7,7 +7,7 @@ # # Each stage supports multiple workers: use ';' to separate workers in GPU # specs. Each worker is an independent top-level process; Dynamo discovers -# them via etcd and the orchestrator round-robins requests automatically. +# them via etcd and the orchestrator dispatches to specific workers via client.direct(). # # Usage: # ./run_all.sh # launch all services (1 worker/stage) diff --git a/examples/disagg_diffusion/stress_test.sh b/examples/disagg_diffusion/stress_test.sh new file mode 100755 index 000000000000..6b3e01ad912b --- /dev/null +++ b/examples/disagg_diffusion/stress_test.sh @@ -0,0 +1,251 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Concurrent stress test for disaggregated diffusion pipeline. +# +# Sends N_REQUESTS concurrent requests, then validates: +# 1. All requests return HTTP 200 with valid JSON +# 2. All produced .mp4 files are non-zero size +# 3. /pipeline/status shows completed=N, failed=0 +# 4. All workers per stage were utilized +# +# Prerequisites: services must already be running (via run_all.sh). +# +# Usage: +# ./stress_test.sh # 20 requests, default port +# N_REQUESTS=50 ./stress_test.sh # 50 requests +# PORT=8081 ./stress_test.sh # custom port + +set -euo pipefail + +PORT="${PORT:-8080}" +N_REQUESTS="${N_REQUESTS:-20}" +BASE_URL="http://localhost:$PORT" +RESULT_DIR="/tmp/stress_test_$$" + +mkdir -p "$RESULT_DIR" + +echo "==========================================" +echo " Disagg Diffusion — Concurrent Stress Test" +echo "==========================================" +echo " Target: $BASE_URL" +echo " Requests: $N_REQUESTS" +echo " Results: $RESULT_DIR/" +echo "==========================================" + +# Verify the server is up +if ! curl -sf "$BASE_URL/health" > /dev/null 2>&1; then + echo "ERROR: Server not reachable at $BASE_URL/health" + echo " Start services first with run_all.sh" + exit 1 +fi + +# Record pre-test status +PRE_STATUS=$(curl -sf "$BASE_URL/pipeline/status") +PRE_COMPLETED=$(echo "$PRE_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['completed'])") +echo "Pre-test completed count: $PRE_COMPLETED" +echo "" + +PROMPTS=( + "A cat walking on green grass" + "A dog running on a sandy beach" + "A bird flying over a mountain lake" + "A fish swimming in a coral reef" + "A horse galloping through a meadow" + "A butterfly landing on a red flower" + "A wolf howling at the full moon" + "A dolphin jumping out of the ocean" + "A panda eating bamboo in the forest" + "A fox running through autumn leaves" + "A penguin sliding on ice in Antarctica" + "An eagle soaring above snowy peaks" + "A tiger walking through a jungle" + "A rabbit hopping through a garden" + "A deer drinking from a stream" + "An owl perched on a branch at night" + "A lion resting under an acacia tree" + "A whale breaching in the deep ocean" + "A parrot perched on a tropical branch" + "A turtle walking slowly on the sand" +) + +# Launch all requests concurrently +echo "Sending $N_REQUESTS concurrent requests..." +T0=$(date +%s) +PIDS=() + +for i in $(seq 0 $((N_REQUESTS - 1))); do + idx=$((i % ${#PROMPTS[@]})) + prompt="${PROMPTS[$idx]}" + ( + curl -sf -X POST "$BASE_URL/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d "{\"prompt\": \"$prompt\", \"num_frames\": 9, \"num_inference_steps\": 3}" \ + -o "$RESULT_DIR/resp_$i.json" \ + -w "%{http_code}" \ + > "$RESULT_DIR/status_$i.txt" 2>/dev/null + ) & + PIDS+=($!) +done + +echo "All $N_REQUESTS requests launched (PIDs: ${#PIDS[@]}). Waiting..." + +# Wait for all requests to finish +FAILURES=0 +for i in "${!PIDS[@]}"; do + if ! wait "${PIDS[$i]}"; then + FAILURES=$((FAILURES + 1)) + echo " Request $i: curl failed (PID ${PIDS[$i]})" + fi +done + +T1=$(date +%s) +ELAPSED=$((T1 - T0)) +echo "" +echo "All requests completed in ${ELAPSED}s" +echo "" + +# ── Validation ── + +PASS=0 +FAIL=0 +VIDEOS=() + +echo "── Validation ──" + +# Check 1: HTTP status codes +echo "" +echo "1) HTTP status codes:" +for i in $(seq 0 $((N_REQUESTS - 1))); do + status_file="$RESULT_DIR/status_$i.txt" + if [ -f "$status_file" ]; then + code=$(cat "$status_file") + if [ "$code" = "200" ]; then + PASS=$((PASS + 1)) + else + FAIL=$((FAIL + 1)) + echo " FAIL: request $i returned HTTP $code" + fi + else + FAIL=$((FAIL + 1)) + echo " FAIL: request $i — no status file (curl crashed)" + fi +done +echo " $PASS/$N_REQUESTS returned HTTP 200" + +# Check 2: Valid JSON with video URL +echo "" +echo "2) Response JSON + video files:" +VIDEO_PASS=0 +VIDEO_FAIL=0 +for i in $(seq 0 $((N_REQUESTS - 1))); do + resp="$RESULT_DIR/resp_$i.json" + if [ ! -f "$resp" ] || [ ! -s "$resp" ]; then + VIDEO_FAIL=$((VIDEO_FAIL + 1)) + continue + fi + url=$(python3 -c " +import json, sys +try: + d = json.load(open('$resp')) + print(d['data'][0].get('url', '')) +except Exception: + print('') +" 2>/dev/null) + if [ -n "$url" ]; then + # Extract filename from /videos/ + fname=$(basename "$url") + VIDEOS+=("$fname") + VIDEO_PASS=$((VIDEO_PASS + 1)) + else + VIDEO_FAIL=$((VIDEO_FAIL + 1)) + echo " FAIL: request $i — no video URL in response" + fi +done +echo " $VIDEO_PASS/$N_REQUESTS have valid video URLs" + +# Check 3: Video files exist and are non-zero +echo "" +echo "3) Video file integrity:" +VIDEO_DIR="/tmp/disagg_videos" +FILE_PASS=0 +FILE_FAIL=0 +for fname in "${VIDEOS[@]}"; do + fpath="$VIDEO_DIR/$fname" + if [ -f "$fpath" ] && [ -s "$fpath" ]; then + FILE_PASS=$((FILE_PASS + 1)) + else + FILE_FAIL=$((FILE_FAIL + 1)) + echo " FAIL: $fname missing or empty" + fi +done +echo " $FILE_PASS/${#VIDEOS[@]} video files valid (non-zero size)" + +# Check 4: Pipeline status — completed count and worker utilization +echo "" +echo "4) Pipeline status:" +POST_STATUS=$(curl -sf "$BASE_URL/pipeline/status") +POST_COMPLETED=$(echo "$POST_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['completed'])") +POST_FAILED=$(echo "$POST_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['failed'])") +NEW_COMPLETED=$((POST_COMPLETED - PRE_COMPLETED)) +echo " Completed: $NEW_COMPLETED (expected: $N_REQUESTS)" +echo " Failed: $POST_FAILED" + +# Check worker utilization +echo "" +echo "5) Worker utilization:" +echo "$POST_STATUS" | python3 -c " +import json, sys +status = json.load(sys.stdin) +all_used = True +for stage_name, stage in status.get('stages', {}).items(): + workers = stage.get('workers', []) + used = [w for w in workers if w['completed'] > 0] + total = len(workers) + print(f' {stage_name}: {len(used)}/{total} workers used', end='') + for w in workers: + print(f\" [id={w['id']} completed={w['completed']} avg={w['avg_latency_s']:.2f}s]\", end='') + print() + if len(used) < total: + all_used = False +if all_used: + print(' All workers utilized.') +else: + print(' WARNING: Not all workers were utilized.') +" + +# ── Summary ── +echo "" +echo "==========================================" +TOTAL_CHECKS=0 +TOTAL_PASS=0 + +# HTTP check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$PASS" -eq "$N_REQUESTS" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# Video URL check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$VIDEO_PASS" -eq "$N_REQUESTS" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# File integrity check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$FILE_PASS" -eq "${#VIDEOS[@]}" ] && [ "${#VIDEOS[@]}" -gt 0 ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# Pipeline completed count check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$NEW_COMPLETED" -eq "$N_REQUESTS" ] && [ "$POST_FAILED" -eq "0" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +if [ "$TOTAL_PASS" -eq "$TOTAL_CHECKS" ]; then + echo " RESULT: ALL PASSED ($TOTAL_PASS/$TOTAL_CHECKS checks)" + echo " $N_REQUESTS requests, ${ELAPSED}s total, $(echo "scale=1; $ELAPSED / $N_REQUESTS" | bc)s avg" +else + echo " RESULT: $TOTAL_PASS/$TOTAL_CHECKS checks passed" +fi +echo "==========================================" + +# Cleanup temp dir +rm -rf "$RESULT_DIR" + +[ "$TOTAL_PASS" -eq "$TOTAL_CHECKS" ] && exit 0 || exit 1 diff --git a/examples/disagg_diffusion/workers/denoiser_worker.py b/examples/disagg_diffusion/workers/denoiser_worker.py index 67fb0d6b8667..35774e461006 100755 --- a/examples/disagg_diffusion/workers/denoiser_worker.py +++ b/examples/disagg_diffusion/workers/denoiser_worker.py @@ -79,10 +79,23 @@ async def handle_generate(request, context): ) req.do_classifier_free_guidance = (req.guidance_scale > 1.0) - # Pass NIXL metadata for NixlReceiveStage to RDMA-pull embeddings + # Pass NIXL metadata or raw tensor data for NixlReceiveStage transfer_meta = request.get("transfer_meta", {}) + tensor_data = request.get("tensor_data", {}) if transfer_meta: req._nixl_transfer_meta = transfer_meta + elif tensor_data: + # ZMQ fallback: deserialize tensors and inject onto Req + import torch, base64, io + from sglang_utils import inject_tensors_to_req + tensors = {} + for k, v in tensor_data.items(): + if isinstance(v, list): + tensors[k] = [torch.load(io.BytesIO(base64.b64decode(b)), weights_only=True) for b in v] + else: + tensors[k] = torch.load(io.BytesIO(base64.b64decode(v)), weights_only=True) + inject_tensors_to_req(req, tensors) + logger.info("Injected %d tensor fields via ZMQ fallback", len(tensors)) output = await client.forward([req]) if output.error: @@ -91,8 +104,20 @@ async def handle_generate(request, context): result = output.output transfer_meta_out = result.get("_nixl_transfer_meta", {}) - logger.info("Denoised — NIXL latent metadata ready") - yield {"transfer_meta": transfer_meta_out, "shape": []} + if transfer_meta_out: + logger.info("Denoised — NIXL latent metadata ready") + yield {"transfer_meta": transfer_meta_out, "shape": []} + else: + # ZMQ fallback for latents + import torch, base64, io + td = {} + for k, v in result.items(): + if isinstance(v, torch.Tensor): + buf = io.BytesIO() + torch.save(v.cpu(), buf) + td[k] = base64.b64encode(buf.getvalue()).decode() + logger.info("Denoised — ZMQ fallback (%d tensor fields)", len(td)) + yield {"transfer_meta": {}, "tensor_data": td, "shape": []} except Exception as e: logger.error("Denoiser generate failed: %s", e, exc_info=True) @@ -106,9 +131,8 @@ async def handle_health(request, context): # ── Serve Dynamo endpoints ─────────────────────────────────────── - ns = runtime.namespace("disagg_diffusion") - gen_ep = ns.component("denoiser").endpoint("generate") - health_ep = ns.component("denoiser").endpoint("health") + gen_ep = runtime.endpoint("disagg_diffusion.denoiser.generate") + health_ep = runtime.endpoint("disagg_diffusion.denoiser.health") logger.info("Serving: disagg_diffusion.denoiser.generate + health") try: diff --git a/examples/disagg_diffusion/workers/encoder_worker.py b/examples/disagg_diffusion/workers/encoder_worker.py index 644d35e5e6a3..b4743b0fe7b3 100755 --- a/examples/disagg_diffusion/workers/encoder_worker.py +++ b/examples/disagg_diffusion/workers/encoder_worker.py @@ -74,8 +74,26 @@ async def handle_generate(request, context): result = output.output transfer_meta = result.get("_nixl_transfer_meta", {}) - logger.info("Encoded prompt — NIXL metadata ready") - yield {"transfer_meta": transfer_meta, "shapes": {}} + if transfer_meta: + logger.info("Encoded prompt — NIXL metadata ready") + yield {"transfer_meta": transfer_meta, "shapes": {}} + else: + # ZMQ fallback: forward raw tensors (when NIXL is disabled) + import torch, base64, io + tensor_data = {} + for k, v in result.items(): + if isinstance(v, torch.Tensor): + buf = io.BytesIO() + torch.save(v.cpu(), buf) + tensor_data[k] = base64.b64encode(buf.getvalue()).decode() + elif isinstance(v, list) and v and isinstance(v[0], torch.Tensor): + tensor_data[k] = [] + for t in v: + buf = io.BytesIO() + torch.save(t.cpu(), buf) + tensor_data[k].append(base64.b64encode(buf.getvalue()).decode()) + logger.info("Encoded prompt — ZMQ fallback (%d tensor fields)", len(tensor_data)) + yield {"transfer_meta": {}, "tensor_data": tensor_data, "shapes": {}} except Exception as e: logger.error("Encoder generate failed: %s", e, exc_info=True) @@ -86,9 +104,8 @@ async def handle_health(request, context): # ── Serve Dynamo endpoints ─────────────────────────────────────── - ns = runtime.namespace("disagg_diffusion") - gen_ep = ns.component("encoder").endpoint("generate") - health_ep = ns.component("encoder").endpoint("health") + gen_ep = runtime.endpoint("disagg_diffusion.encoder.generate") + health_ep = runtime.endpoint("disagg_diffusion.encoder.health") logger.info("Serving: disagg_diffusion.encoder.generate + health") try: diff --git a/examples/disagg_diffusion/workers/nixl_transfer.py b/examples/disagg_diffusion/workers/nixl_transfer.py index 3ee77d6e88b7..7b1cf2417328 100644 --- a/examples/disagg_diffusion/workers/nixl_transfer.py +++ b/examples/disagg_diffusion/workers/nixl_transfer.py @@ -22,6 +22,8 @@ import asyncio import logging +import os +import time from typing import Dict import torch @@ -30,7 +32,9 @@ try: import dynamo.nixl_connect as nixl_connect - NIXL_AVAILABLE = True + NIXL_AVAILABLE = not os.environ.get("DISABLE_NIXL", "").lower() in ("1", "true", "yes") + if not NIXL_AVAILABLE: + logger.info("NIXL disabled via DISABLE_NIXL env var") except ImportError: NIXL_AVAILABLE = False logger.info("NIXL not available — falling back to ZMQ tensor transfer") @@ -52,21 +56,50 @@ async def get(cls): class NixlTensorSender: """Register GPU tensors as NIXL-readable. Returns metadata for the receiver. - The readable is kept alive via a background task so the sender process - can return immediately after yielding metadata. + Buffers are held until the receiver completes the RDMA pull (detected + via synchronous ``readable.status`` polling) or a configurable timeout + expires. Previous implementation scheduled a ``_keep_alive`` background + task via ``asyncio.ensure_future``, but the event loop only runs during + ``run_until_complete`` and stops immediately after — so the background + task never executed, leaking GPU memory. """ + BUFFER_TIMEOUT_S = float(os.environ.get("NIXL_BUFFER_TIMEOUT_S", "120")) + def __init__(self): - self._pending: list = [] + # Each entry: (readable, flat_buffer_ref, creation_timestamp) + self._pending: list[tuple[object, torch.Tensor, float]] = [] def send(self, tensors: Dict[str, torch.Tensor]) -> dict: """Register tensors and return metadata dict (synchronous wrapper).""" + self._sweep() # release completed / timed-out buffers first return asyncio.get_event_loop().run_until_complete(self._async_send(tensors)) - async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: - # Clean completed tasks - self._pending = [t for t in self._pending if not t.done()] + def _sweep(self): + """Poll pending readables: release completed or timed-out buffers.""" + now = time.monotonic() + still_pending = [] + for readable, flat, created_at in self._pending: + try: + status = readable.status # synchronous — calls update_notifs() + except Exception: + # If status check fails, treat as completed to avoid leak + logger.debug("NIXL readable status check failed, releasing buffer") + continue + if hasattr(status, "name") and status.name == "COMPLETE": + logger.debug("NIXL readable completed, releasing buffer") + elif str(status) == "OperationStatus.COMPLETE": + logger.debug("NIXL readable completed, releasing buffer") + elif now - created_at > self.BUFFER_TIMEOUT_S: + logger.warning( + "NIXL readable timed out after %.0fs, force-releasing buffer", + now - created_at, + ) + else: + still_pending.append((readable, flat, created_at)) + self._pending = still_pending + async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: connector = await _PersistentConnector.get() # Flatten all tensors into a single contiguous buffer @@ -82,15 +115,8 @@ async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: "nixl_metadata": raw_meta.model_dump() if hasattr(raw_meta, "model_dump") else raw_meta, } - # Keep readable alive until the receiver has pulled the data - async def _keep_alive(): - try: - await readable.wait_for_completion() - except Exception as e: - logger.warning("NIXL readable wait failed: %s", e) - - task = asyncio.ensure_future(_keep_alive()) - self._pending.append(task) + # Hold (readable, flat_buffer, timestamp) — prevents GC until sweep releases + self._pending.append((readable, flat, time.monotonic())) return meta diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py index fabc0c406747..c2d4718173fc 100644 --- a/examples/disagg_diffusion/workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -57,7 +57,7 @@ get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch, RequestMetrics from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -96,11 +96,35 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: # Fallback: tensors arrived via ZMQ pickle — just move to GPU return self._device_move(batch) + _NIXL_PULL_MAX_RETRIES = int(os.environ.get("NIXL_PULL_MAX_RETRIES", "5")) + _NIXL_PULL_BACKOFF_S = float(os.environ.get("NIXL_PULL_BACKOFF_S", "0.5")) + def _nixl_pull(self, batch: Req, meta: dict) -> Req: from nixl_transfer import NixlTensorReceiver if self._receiver is None: self._receiver = NixlTensorReceiver() - tensors = self._receiver.recv(meta, device="cuda") + # Retry on REMOTE_DISCONNECT — NIXL 1:N fan-out can be flaky on + # initial peer connection when multiple receivers target one sender. + last_err = None + for attempt in range(self._NIXL_PULL_MAX_RETRIES + 1): + try: + tensors = self._receiver.recv(meta, device="cuda") + break + except Exception as e: + if "REMOTE_DISCONNECT" in str(e) and attempt < self._NIXL_PULL_MAX_RETRIES: + import time + wait = self._NIXL_PULL_BACKOFF_S * (attempt + 1) + logger.warning( + "NIXL pull attempt %d/%d failed: %s, retrying in %.1fs", + attempt + 1, self._NIXL_PULL_MAX_RETRIES, e, wait, + ) + self._receiver = NixlTensorReceiver() # fresh receiver + time.sleep(wait) + last_err = e + continue + raise + else: + raise RuntimeError(f"NIXL pull failed after {self._NIXL_PULL_MAX_RETRIES} retries") from last_err # Reconstruct indexed fields (e.g. prompt_embeds_0, prompt_embeds_1 # → prompt_embeds list) reconstructed = {} @@ -163,19 +187,10 @@ def __init__(self, output_fields: List[str]): self._output_fields = output_fields self._sender = None - @staticmethod - def _make_timings() -> object: - """Create a RequestTimings object compatible with sglang's executor.""" - try: - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import RequestTimings - return RequestTimings(request_id="nixl-send") - except Exception: - return None - def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: tensors = self._extract_tensors(batch) if not tensors: - return OutputBatch(output={}, timings=self._make_timings()) + return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) # Debug: log tensor stats before NIXL send for k, t in tensors.items(): @@ -213,7 +228,7 @@ def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: if self._sender is None: self._sender = NixlTensorSender() meta = self._sender.send(tensors) - return OutputBatch(output={"_nixl_transfer_meta": meta}, timings=self._make_timings()) + return OutputBatch(output={"_nixl_transfer_meta": meta}, metrics=RequestMetrics(request_id="nixl")) def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: """Fallback: send raw tensors via ZMQ pickle.""" @@ -231,7 +246,7 @@ def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: real_key = base[6:] output[real_key] = [idx_map[i] for i in sorted(idx_map)] del output[base] - return OutputBatch(output=output, timings=self._make_timings()) + return OutputBatch(output=output, metrics=RequestMetrics(request_id="nixl")) # ═══════════════════════════════════════════════════════════════════════ @@ -389,7 +404,7 @@ def init_device_and_model(self) -> None: stages = self._custom_stages_fn(self.pipeline, self.server_args) for stage in stages: name = type(stage).__name__ - self.pipeline.add_stage(name, stage) + self.pipeline.add_stage(stage, name) if getattr(self.server_args, "dit_layerwise_offload", False) and OffloadableDiTMixin is not None: for module_name in [ diff --git a/examples/disagg_diffusion/workers/protocol.py b/examples/disagg_diffusion/workers/protocol.py index ce935db9cbef..2a059ff2eb64 100644 --- a/examples/disagg_diffusion/workers/protocol.py +++ b/examples/disagg_diffusion/workers/protocol.py @@ -35,6 +35,7 @@ class EncoderResponse(BaseModel): class DenoiserRequest(BaseModel): transfer_meta: Dict[str, Any] + tensor_data: Dict[str, Any] = {} height: int = 544 width: int = 960 num_frames: int = 61 @@ -54,6 +55,7 @@ class DenoiserResponse(BaseModel): class VAEDecodeRequest(BaseModel): transfer_meta: Dict[str, Any] + tensor_data: Dict[str, Any] = {} request_id: str = "" diff --git a/examples/disagg_diffusion/workers/sglang_utils.py b/examples/disagg_diffusion/workers/sglang_utils.py index 2468cb791153..fac5aa70f69d 100644 --- a/examples/disagg_diffusion/workers/sglang_utils.py +++ b/examples/disagg_diffusion/workers/sglang_utils.py @@ -170,6 +170,8 @@ def build_req( seed=seed, generator=torch.Generator(device="cpu").manual_seed(seed), do_classifier_free_guidance=(guidance_scale > 1.0), + save_output=False, + return_file_paths_only=False, ) for k, v in extra_fields.items(): @@ -218,6 +220,8 @@ def inject_tensors_to_req( class StageClient: """Async ZMQ REQ client that talks to a SGLang Scheduler subprocess.""" + FORWARD_TIMEOUT_S = float(os.environ.get("STAGE_FORWARD_TIMEOUT_S", "120")) + def __init__(self, endpoint: str, name: str = ""): import zmq.asyncio self._name = name @@ -228,10 +232,13 @@ def __init__(self, endpoint: str, name: str = ""): logger.info("StageClient(%s) connected to %s", name, endpoint) async def forward(self, reqs): - """Send request(s) and receive response.""" + """Send request(s) and receive response (with timeout).""" async with self._lock: await self._sock.send_pyobj(reqs) - return await self._sock.recv_pyobj() + return await asyncio.wait_for( + self._sock.recv_pyobj(), + timeout=self.FORWARD_TIMEOUT_S, + ) def close(self): self._sock.close() diff --git a/examples/disagg_diffusion/workers/vae_worker.py b/examples/disagg_diffusion/workers/vae_worker.py index 6a65cdb383d7..95f37d1eeb22 100755 --- a/examples/disagg_diffusion/workers/vae_worker.py +++ b/examples/disagg_diffusion/workers/vae_worker.py @@ -73,10 +73,23 @@ async def handle_generate(request, context): seed=request.get("seed", 42), ) - # Pass NIXL metadata for NixlReceiveStage to RDMA-pull latents + # Pass NIXL metadata or raw tensor data for NixlReceiveStage transfer_meta = request.get("transfer_meta", {}) + tensor_data = request.get("tensor_data", {}) if transfer_meta: req._nixl_transfer_meta = transfer_meta + elif tensor_data: + # ZMQ fallback: deserialize latents and inject onto Req + import torch, base64, io + from sglang_utils import inject_tensors_to_req + tensors = {} + for k, v in tensor_data.items(): + if isinstance(v, list): + tensors[k] = [torch.load(io.BytesIO(base64.b64decode(b)), weights_only=True) for b in v] + else: + tensors[k] = torch.load(io.BytesIO(base64.b64decode(v)), weights_only=True) + inject_tensors_to_req(req, tensors) + logger.info("Injected %d tensor fields via ZMQ fallback", len(tensors)) output = await client.forward([req]) if output.error: @@ -105,9 +118,8 @@ async def handle_health(request, context): # ── Serve Dynamo endpoints ─────────────────────────────────────── - ns = runtime.namespace("disagg_diffusion") - gen_ep = ns.component("vae").endpoint("generate") - health_ep = ns.component("vae").endpoint("health") + gen_ep = runtime.endpoint("disagg_diffusion.vae.generate") + health_ep = runtime.endpoint("disagg_diffusion.vae.health") logger.info("Serving: disagg_diffusion.vae.generate + health") try: From 92453a97299ff0984dfdbe7c9e652e80f4bf7d15 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 18:16:18 +0800 Subject: [PATCH 19/22] fix: NIXL 1:N REMOTE_DISCONNECT in disagg diffusion pipeline Root causes and fixes: 1. PersistentConnector: replace _PersistentConnector singleton (which created a new Connection per operation) with a proper subclass that reuses one Connection/agent. Each Sender/Receiver owns its own instance, eagerly initialized at stage construction. 2. UCX transport: default UCX_TLS=all picks IB RDMA which fails across NUMA nodes. Must set UCX_TLS=cuda_ipc,tcp,self,cuda_copy,cma to force NVLink path for intra-node GPU-direct transfers. 3. Sender returns readable_op: NixlTensorSender.send() now returns (meta, readable_op). NixlSendStage holds readable_ops in _active_readables until receiver completes the RDMA pull, preventing premature descriptor deregistration. 4. TP broadcast: with TP>1, only rank 0 gets NIXL data. Added _tp_broadcast_fields to distribute tensors to all TP ranks via torch.distributed.broadcast after the NIXL pull. 5. ZMQ socket recovery: StageClient.forward() now resets the ZMQ REQ socket after timeout, preventing cascading EFSM errors. 6. Worker discovery: orchestrator now waits WORKER_SETTLE_S (default 30s) after initial discovery to pick up slow-registering workers. Verified: 20 concurrent requests, 3 TP=2 denoisers, 8 GPUs, 0 failures. Co-Authored-By: Claude Opus 4.6 --- .../orchestrator/run_disagg.py | 16 ++ .../disagg_diffusion/workers/nixl_transfer.py | 180 +++++++++++------- .../workers/partial_gpu_worker.py | 140 ++++++++++++-- .../disagg_diffusion/workers/sglang_utils.py | 29 ++- 4 files changed, 274 insertions(+), 91 deletions(-) diff --git a/examples/disagg_diffusion/orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py index 660f7c39070e..384e69090fef 100755 --- a/examples/disagg_diffusion/orchestrator/run_disagg.py +++ b/examples/disagg_diffusion/orchestrator/run_disagg.py @@ -146,6 +146,22 @@ async def worker(runtime: DistributedRuntime): await denoiser_client.wait_for_instances() await vae_client.wait_for_instances() + # Wait for additional workers that may still be registering. + # wait_for_instances() returns after the first instance; model loading + # times vary, so poll until the count stabilizes or timeout. + WORKER_SETTLE_S = int(os.environ.get("WORKER_SETTLE_S", "30")) + if WORKER_SETTLE_S > 0: + import time as _time + deadline = _time.monotonic() + WORKER_SETTLE_S + prev_count = 0 + while _time.monotonic() < deadline: + cur = len(denoiser_client.instance_ids()) + if cur > prev_count: + prev_count = cur + logger.info("Discovered %d denoiser(s) so far, waiting for more…", cur) + await asyncio.sleep(2) + logger.info("Worker settle period done (%ds)", WORKER_SETTLE_S) + # Discover registered worker instances per stage enc_ids = encoder_client.instance_ids() den_ids = denoiser_client.instance_ids() diff --git a/examples/disagg_diffusion/workers/nixl_transfer.py b/examples/disagg_diffusion/workers/nixl_transfer.py index 7b1cf2417328..64fe82f50679 100644 --- a/examples/disagg_diffusion/workers/nixl_transfer.py +++ b/examples/disagg_diffusion/workers/nixl_transfer.py @@ -8,13 +8,19 @@ control plane; actual tensor data (embeddings, latents) transfers GPU->GPU via NIXL RDMA. +Follows the PersistentConnector pattern from embedding_transfer.py: +- Each Sender/Receiver owns its PersistentConnector (one Connection/agent). +- Remote._release is nooped to keep agent pairs alive. +- The sender returns a ``readable_op`` handle that the caller must hold + until the receiver completes the RDMA pull. + Usage inside PipelineStage.forward() (synchronous context):: - sender = NixlTensorSender() - meta = sender.send({"latents": tensor}) # registers & returns metadata - # ... pass meta via ZMQ ... + sender = NixlTensorSender() # creates agent eagerly + meta, readable_op = sender.send({"latents": tensor}) + # ... pass meta via ZMQ, hold readable_op until COMPLETE ... - receiver = NixlTensorReceiver() + receiver = NixlTensorReceiver() # creates agent eagerly tensors = receiver.recv(meta, device="cuda") # RDMA pull """ @@ -23,8 +29,7 @@ import asyncio import logging import os -import time -from typing import Dict +from typing import Dict, Tuple import torch @@ -32,7 +37,12 @@ try: import dynamo.nixl_connect as nixl_connect - NIXL_AVAILABLE = not os.environ.get("DISABLE_NIXL", "").lower() in ("1", "true", "yes") + + NIXL_AVAILABLE = not os.environ.get("DISABLE_NIXL", "").lower() in ( + "1", + "true", + "yes", + ) if not NIXL_AVAILABLE: logger.info("NIXL disabled via DISABLE_NIXL env var") except ImportError: @@ -40,96 +50,122 @@ logger.info("NIXL not available — falling back to ZMQ tensor transfer") -class _PersistentConnector: - """Lazily-initialized NIXL Connector singleton per process.""" +# --------------------------------------------------------------------------- +# PersistentConnector + Remote._release noop +# --------------------------------------------------------------------------- +# Exact pattern from components/src/dynamo/common/multimodal/embedding_transfer.py + +if NIXL_AVAILABLE: + + class PersistentConnector(nixl_connect.Connector): + """Connector that reuses a single Connection for all operations.""" + + def __init__(self): + super().__init__() + self._connection = None + + async def _create_connection(self) -> nixl_connect.Connection: + if self._connection is None: + self._connection = nixl_connect.Connection(self, 1) + await self._connection.initialize() + return self._connection - _instance = None + # NOTE: We do NOT noop Remote._release here. Our recv() is synchronous + # (awaits wait_for_completion before returning), so the transfer is + # always complete before Remote is GC'd. Keeping the remote agent + # registered across requests causes NIXL_ERR_NOT_ALLOWED on the + # second add_remote_agent() call. - @classmethod - async def get(cls): - if cls._instance is None: - cls._instance = nixl_connect.Connector() - await cls._instance.initialize() - return cls._instance + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _run_coro(coro): + """Run a coroutine from synchronous context (sglang scheduler thread).""" + try: + loop = asyncio.get_event_loop() + if loop.is_running(): + # Shouldn't happen in sglang's scheduler, but be safe + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor() as pool: + return pool.submit(asyncio.run, coro).result(timeout=30) + return loop.run_until_complete(coro) + except RuntimeError: + return asyncio.run(coro) + + +# --------------------------------------------------------------------------- +# NixlTensorSender +# --------------------------------------------------------------------------- class NixlTensorSender: - """Register GPU tensors as NIXL-readable. Returns metadata for the receiver. - - Buffers are held until the receiver completes the RDMA pull (detected - via synchronous ``readable.status`` polling) or a configurable timeout - expires. Previous implementation scheduled a ``_keep_alive`` background - task via ``asyncio.ensure_future``, but the event loop only runs during - ``run_until_complete`` and stops immediately after — so the background - task never executed, leaking GPU memory. - """ + """Register GPU tensors as NIXL-readable. Returns (metadata, readable_op). - BUFFER_TIMEOUT_S = float(os.environ.get("NIXL_BUFFER_TIMEOUT_S", "120")) + Each instance owns a PersistentConnector whose Connection (nixl_agent) + is created eagerly in ``__init__`` so UCX is fully initialized before + any transfer occurs. + + The caller **must** hold ``readable_op`` until the receiver completes + the RDMA pull. + """ def __init__(self): - # Each entry: (readable, flat_buffer_ref, creation_timestamp) - self._pending: list[tuple[object, torch.Tensor, float]] = [] - - def send(self, tensors: Dict[str, torch.Tensor]) -> dict: - """Register tensors and return metadata dict (synchronous wrapper).""" - self._sweep() # release completed / timed-out buffers first - return asyncio.get_event_loop().run_until_complete(self._async_send(tensors)) - - def _sweep(self): - """Poll pending readables: release completed or timed-out buffers.""" - now = time.monotonic() - still_pending = [] - for readable, flat, created_at in self._pending: - try: - status = readable.status # synchronous — calls update_notifs() - except Exception: - # If status check fails, treat as completed to avoid leak - logger.debug("NIXL readable status check failed, releasing buffer") - continue - if hasattr(status, "name") and status.name == "COMPLETE": - logger.debug("NIXL readable completed, releasing buffer") - elif str(status) == "OperationStatus.COMPLETE": - logger.debug("NIXL readable completed, releasing buffer") - elif now - created_at > self.BUFFER_TIMEOUT_S: - logger.warning( - "NIXL readable timed out after %.0fs, force-releasing buffer", - now - created_at, - ) - else: - still_pending.append((readable, flat, created_at)) - self._pending = still_pending - - async def _async_send(self, tensors: Dict[str, torch.Tensor]) -> dict: - connector = await _PersistentConnector.get() + self.connector = PersistentConnector() + # Eagerly create the Connection / nixl_agent so UCX is ready + _run_coro(self.connector._create_connection()) + + def send(self, tensors: Dict[str, torch.Tensor]) -> Tuple[dict, object]: + """Register tensors and return (metadata_dict, readable_op).""" + return _run_coro(self._async_send(tensors)) + async def _async_send( + self, tensors: Dict[str, torch.Tensor] + ) -> Tuple[dict, object]: # Flatten all tensors into a single contiguous buffer flat = torch.cat([t.contiguous().view(-1) for t in tensors.values()]) descriptor = nixl_connect.Descriptor(flat) - readable = await connector.create_readable(descriptor) + readable = await self.connector.create_readable(descriptor) raw_meta = readable.metadata() meta = { "tensor_keys": list(tensors.keys()), "shapes": {k: list(t.shape) for k, t in tensors.items()}, - "dtypes": {k: str(t.dtype).removeprefix("torch.") for k, t in tensors.items()}, - "nixl_metadata": raw_meta.model_dump() if hasattr(raw_meta, "model_dump") else raw_meta, + "dtypes": { + k: str(t.dtype).removeprefix("torch.") for k, t in tensors.items() + }, + "nixl_metadata": raw_meta.model_dump() + if hasattr(raw_meta, "model_dump") + else raw_meta, } - # Hold (readable, flat_buffer, timestamp) — prevents GC until sweep releases - self._pending.append((readable, flat, time.monotonic())) - return meta + # Return both — caller holds readable to prevent GC + return meta, readable + + +# --------------------------------------------------------------------------- +# NixlTensorReceiver +# --------------------------------------------------------------------------- class NixlTensorReceiver: - """Pull tensors from a remote sender via NIXL RDMA.""" + """Pull tensors from a remote sender via NIXL RDMA. + + Each instance owns a PersistentConnector whose Connection (nixl_agent) + is created eagerly in ``__init__``. + """ + + def __init__(self): + self.connector = PersistentConnector() + # Eagerly create the Connection / nixl_agent so UCX is ready + _run_coro(self.connector._create_connection()) def recv(self, meta: dict, device: str = "cuda") -> Dict[str, torch.Tensor]: """Pull tensors described by metadata. Returns {name: tensor}.""" - return asyncio.get_event_loop().run_until_complete(self._async_recv(meta, device)) + return _run_coro(self._async_recv(meta, device)) async def _async_recv(self, meta: dict, device: str) -> Dict[str, torch.Tensor]: - connector = await _PersistentConnector.get() - # Calculate total size and per-tensor specs specs = [] total_bytes = 0 @@ -148,14 +184,14 @@ async def _async_recv(self, meta: dict, device: str) -> Dict[str, torch.Tensor]: descriptor = nixl_connect.Descriptor(flat) rdma_meta = nixl_connect.RdmaMetadata.model_validate(meta["nixl_metadata"]) - read_op = await connector.begin_read(rdma_meta, descriptor) + read_op = await self.connector.begin_read(rdma_meta, descriptor) await read_op.wait_for_completion() # Slice the flat buffer into individual tensors result = {} offset = 0 for key, shape, dtype, size in specs: - result[key] = flat[offset:offset + size].view(dtype=dtype).reshape(shape) + result[key] = flat[offset : offset + size].view(dtype=dtype).reshape(shape) offset += size return result diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py index c2d4718173fc..f75ea406243c 100644 --- a/examples/disagg_diffusion/workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -87,22 +87,48 @@ class NixlReceiveStage(PipelineStage): def __init__(self, tensor_fields: List[str]): super().__init__() self._tensor_fields = tensor_fields - self._receiver = None + # Eagerly create receiver so its NIXL agent/UCX endpoint is ready + # before any transfer occurs in this process. + # Only rank 0 uses NIXL — TP slaves skip to avoid UCX conflicts. + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE and get_tp_rank() == 0: + from nixl_transfer import NixlTensorReceiver + self._receiver = NixlTensorReceiver() + else: + self._receiver = None def forward(self, batch: Req, server_args: ServerArgs) -> Req: + tp_world = get_tp_world_size() + tp_rank = get_tp_rank() nixl_meta = getattr(batch, "_nixl_transfer_meta", None) - if nixl_meta is not None: - return self._nixl_pull(batch, nixl_meta) - # Fallback: tensors arrived via ZMQ pickle — just move to GPU + + if tp_world <= 1: + # No TP — simple path + if nixl_meta is not None and self._receiver is not None: + return self._nixl_pull(batch, nixl_meta) + return self._device_move(batch) + + # TP mode: rank 0 decides NIXL vs ZMQ, broadcast decision to all ranks + # (torch.distributed.broadcast requires ALL ranks to participate) + import torch.distributed as dist + flag = torch.zeros(1, dtype=torch.int32, device="cuda") + if tp_rank == 0 and nixl_meta is not None and self._receiver is not None: + flag[0] = 1 + dist.broadcast(flag, src=0) + + if flag.item(): + # NIXL path: rank 0 pulls, then broadcasts tensors to all TP ranks + if tp_rank == 0: + batch = self._nixl_pull(batch, nixl_meta) + self._tp_broadcast_fields(batch, src_rank=0) + return batch + return self._device_move(batch) _NIXL_PULL_MAX_RETRIES = int(os.environ.get("NIXL_PULL_MAX_RETRIES", "5")) _NIXL_PULL_BACKOFF_S = float(os.environ.get("NIXL_PULL_BACKOFF_S", "0.5")) def _nixl_pull(self, batch: Req, meta: dict) -> Req: - from nixl_transfer import NixlTensorReceiver - if self._receiver is None: - self._receiver = NixlTensorReceiver() # Retry on REMOTE_DISCONNECT — NIXL 1:N fan-out can be flaky on # initial peer connection when multiple receivers target one sender. last_err = None @@ -118,7 +144,9 @@ def _nixl_pull(self, batch: Req, meta: dict) -> Req: "NIXL pull attempt %d/%d failed: %s, retrying in %.1fs", attempt + 1, self._NIXL_PULL_MAX_RETRIES, e, wait, ) - self._receiver = NixlTensorReceiver() # fresh receiver + # Reuse the same receiver — PersistentConnector handles + # agent reuse. Creating a fresh receiver was itself a + # source of REMOTE_DISCONNECT (new agent overwhelms UCX). time.sleep(wait) last_err = e continue @@ -153,6 +181,55 @@ def _nixl_pull(self, batch: Req, meta: dict) -> Req: inject_tensors_to_req(batch, reconstructed) return batch + def _tp_broadcast_fields(self, batch: Req, src_rank: int = 0) -> None: + """Broadcast tensor fields from src_rank to all TP ranks. + + Uses broadcast_object_list for shape/dtype metadata, then + point-to-point broadcast for GPU tensor data. + """ + import torch.distributed as dist + tp_rank = get_tp_rank() + device = torch.device("cuda") + + # Collect field data on rank 0, None on others + field_data = [None] # single-element list for broadcast_object_list + if tp_rank == src_rank: + meta = {} + for field in self._tensor_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + meta[field] = [{"shape": list(t.shape), "dtype": str(t.dtype)} for t in val if isinstance(t, torch.Tensor)] + elif isinstance(val, torch.Tensor): + meta[field] = {"shape": list(val.shape), "dtype": str(val.dtype)} + field_data = [meta] + + dist.broadcast_object_list(field_data, src=src_rank) + meta = field_data[0] + if meta is None: + return + + # Allocate tensors on non-src ranks, then broadcast data + for field, info in meta.items(): + if isinstance(info, list): + # List of tensors + if tp_rank == src_rank: + tensors = getattr(batch, field) + else: + tensors = [torch.empty(m["shape"], dtype=getattr(torch, m["dtype"].removeprefix("torch.")), device=device) for m in info] + setattr(batch, field, tensors) + for t in tensors: + dist.broadcast(t, src=src_rank) + elif isinstance(info, dict): + # Single tensor + if tp_rank != src_rank: + t = torch.empty(info["shape"], dtype=getattr(torch, info["dtype"].removeprefix("torch.")), device=device) + setattr(batch, field, t) + else: + t = getattr(batch, field) + dist.broadcast(t, src=src_rank) + def _device_move(self, batch: Req) -> Req: """Fallback: move CPU tensors to GPU (ZMQ pickle path).""" device = torch.device("cuda") @@ -180,14 +257,34 @@ class NixlSendStage(PipelineStage): Actual tensor data stays on GPU — only metadata travels over ZMQ. Falls back to sending raw tensors when NIXL is unavailable. + + Holds ``readable_op`` handles in ``_active_readables`` until the + receiver completes the RDMA pull — prevents premature GC of the + Connection/Descriptor/GPU buffer that causes REMOTE_DISCONNECT. """ def __init__(self, output_fields: List[str]): super().__init__() self._output_fields = output_fields - self._sender = None + # Eagerly create sender so its NIXL agent/UCX endpoint is ready + # before any transfer occurs in this process. + # Only rank 0 uses NIXL — TP slaves skip to avoid UCX conflicts. + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE and get_tp_rank() == 0: + from nixl_transfer import NixlTensorSender + self._sender = NixlTensorSender() + else: + self._sender = None + # (readable_op, flat_buffer_ref) — kept alive until receiver completes + self._active_readables: list[tuple[object, object]] = [] def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + # TP rank > 0: only rank 0's output is returned to the orchestrator + if get_tp_rank() != 0: + return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) + + self._poll_completed() # release finished readables first + tensors = self._extract_tensors(batch) if not tensors: return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) @@ -199,10 +296,26 @@ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: k, t.shape, t.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) from nixl_transfer import NIXL_AVAILABLE - if NIXL_AVAILABLE: + if NIXL_AVAILABLE and self._sender is not None: return self._nixl_send(tensors) return self._fallback_send(tensors) + def _poll_completed(self): + """Release readable_ops whose RDMA pull has completed.""" + still_active = [] + for readable, buf_ref in self._active_readables: + try: + status = readable.status + name = getattr(status, "name", str(status)) + if "COMPLETE" in name: + logger.debug("NIXL readable completed, releasing buffer") + continue # drop — GC releases descriptor + buffer + still_active.append((readable, buf_ref)) + except Exception: + logger.warning("NIXL readable status check failed, keeping alive to be safe") + still_active.append((readable, buf_ref)) + self._active_readables = still_active + def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: """Flatten list-valued fields into individual tensors.""" result: Dict[str, torch.Tensor] = {} @@ -224,10 +337,9 @@ def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: return result def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: - from nixl_transfer import NixlTensorSender - if self._sender is None: - self._sender = NixlTensorSender() - meta = self._sender.send(tensors) + meta, readable_op = self._sender.send(tensors) + # Hold readable_op (+ implicit buffer ref) until receiver completes pull + self._active_readables.append((readable_op, None)) return OutputBatch(output={"_nixl_transfer_meta": meta}, metrics=RequestMetrics(request_id="nixl")) def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: diff --git a/examples/disagg_diffusion/workers/sglang_utils.py b/examples/disagg_diffusion/workers/sglang_utils.py index fac5aa70f69d..47e181bb121c 100644 --- a/examples/disagg_diffusion/workers/sglang_utils.py +++ b/examples/disagg_diffusion/workers/sglang_utils.py @@ -223,22 +223,41 @@ class StageClient: FORWARD_TIMEOUT_S = float(os.environ.get("STAGE_FORWARD_TIMEOUT_S", "120")) def __init__(self, endpoint: str, name: str = ""): - import zmq.asyncio + import zmq, zmq.asyncio + self._zmq = zmq self._name = name + self._endpoint = endpoint self._ctx = zmq.asyncio.Context() self._sock = self._ctx.socket(zmq.REQ) self._sock.connect(endpoint) self._lock = asyncio.Lock() logger.info("StageClient(%s) connected to %s", name, endpoint) + def _reset_socket(self): + """Recreate ZMQ socket after a timeout leaves it in a broken state. + + ZMQ REQ sockets enforce strict send/recv alternation. If recv + times out after send, the socket is stuck waiting for a reply and + all subsequent sends raise ``EFSM``. The only recovery is to + close and reconnect. + """ + self._sock.close(linger=0) + self._sock = self._ctx.socket(self._zmq.REQ) + self._sock.connect(self._endpoint) + logger.warning("StageClient(%s) reset ZMQ socket after timeout", self._name) + async def forward(self, reqs): """Send request(s) and receive response (with timeout).""" async with self._lock: await self._sock.send_pyobj(reqs) - return await asyncio.wait_for( - self._sock.recv_pyobj(), - timeout=self.FORWARD_TIMEOUT_S, - ) + try: + return await asyncio.wait_for( + self._sock.recv_pyobj(), + timeout=self.FORWARD_TIMEOUT_S, + ) + except (asyncio.TimeoutError, TimeoutError): + self._reset_socket() + raise def close(self): self._sock.close() From 7a0bda85e2b52dcafe79883065445f8605268a2a Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 19:48:02 +0800 Subject: [PATCH 20/22] refactor: add request validation, improve worker settle, update DESIGN.md MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add request_validation.py with pydantic-based HTTP body validation - Orchestrator: observe all stages (encoder/denoiser/vae) during settle period instead of only denoiser - Orchestrator: structured error handling (ValidationError → 400) - protocol.py: add Field annotations - DESIGN.md: streamline and align with current implementation Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/DESIGN.md | 713 ++++++------------ .../orchestrator/request_validation.py | 18 + .../orchestrator/run_disagg.py | 37 +- .../tests/test_protocol_models.py | 41 + .../tests/test_request_validation.py | 38 + examples/disagg_diffusion/workers/protocol.py | 16 +- 6 files changed, 366 insertions(+), 497 deletions(-) create mode 100644 examples/disagg_diffusion/orchestrator/request_validation.py create mode 100644 examples/disagg_diffusion/tests/test_protocol_models.py create mode 100644 examples/disagg_diffusion/tests/test_request_validation.py diff --git a/examples/disagg_diffusion/DESIGN.md b/examples/disagg_diffusion/DESIGN.md index 4b3685ea4321..e8bc85f474a4 100644 --- a/examples/disagg_diffusion/DESIGN.md +++ b/examples/disagg_diffusion/DESIGN.md @@ -1,518 +1,271 @@ -# Disaggregated Diffusion Pipeline — Design Document +# Disaggregated Diffusion Pipeline - Design Document -## 1. Overview +## 1. Goal and Scope -### Motivation +This design splits video generation into three independently scalable stages: +Encoder -> Denoiser -> VAE. -Modern diffusion pipelines (text-to-video, text-to-image, omni-modal) are -composed of heterogeneous stages — text encoding, iterative denoising, VAE -decoding — each with fundamentally different compute profiles: +The orchestrator handles control-plane scheduling via Dynamo RPC; tensors move +across stages through NIXL GPU-direct RDMA. -- **Different optimization strategies per stage.** Encoders are - memory-bound single-pass transforms; denoisers are compute-bound - multi-step loops that benefit from tensor parallelism; VAE decoders - are memory-intensive but run only once. Forcing all three into a - single process prevents stage-specific tuning (parallelism, batching, - memory management, quantization). +**Current design choices (implementation-aligned):** +- Multi-worker: each stage runs multiple worker instances. +- Workers can be added/removed at runtime through etcd discovery without + orchestrator restart. +- Orchestrator does stage routing; current policy is round-robin over idle + workers (`WorkerManager` queue order). +- Pipeline overlap is used to hide transfer overhead behind compute. +- Tensor buffers are currently temporary allocations per request; this is + acceptable with current latency profile and can be upgraded to pooled/buffer + management later. -- **Shifting compute balance.** As diffusion models mature, the DiT no - longer dominates the entire pipeline — faster denoisers (fewer steps, - distilled models) shift the bottleneck to encoding and decoding. - Multi-task pipelines are emerging (e.g. OneVideo: encode + denoise + - decode + audio in one model) where each task demands independent - scaling. - -- **Omni-modal future.** Models that jointly produce video, audio, - image, and text require stage-level separation so each modality's - compute can scale independently without wasting GPU resources. - -### Solution - -Decompose the pipeline into **N independent stages**, each running as a -Dynamo RPC worker on dedicated GPU(s). An orchestrator chains stages -together, using NIXL RDMA for GPU-direct tensor transfer between them. - -| Goal | Mechanism | +| Goal | Current mechanism | |---|---| -| Stage-level scaling | N workers per stage, auto-discovered via etcd | -| Pipeline parallelism | Semaphore admission, independent worker pools | -| GPU-direct transfer | NIXL RDMA — only ~1.5 KB metadata over RPC | -| Loose coupling | Workers are independent processes, no shared state | -| Dynamic scaling | Add/remove workers at runtime, no restart | -| Auto routing | Idle-queue dispatch, backpressure, retry on failure | - - -## 2. Architecture - -### 2.1 Architecture Diagram - -``` - ┌───────────┐ - │ etcd │ - │ registry │ - └─────┬─────┘ - register │ discover - ┌───────────────────────────────┼───────────────────────────────┐ - │ Orchestrator │ - │ │ - │ HTTP ──► handle_generate() ──► dispatch_with_retry() │ - │ PipelineTracker per-stage WorkerManager │ - │ Semaphore(depth) acquire → direct() → release │ - │ │ - └──────┬──────────────────────┬──────────────────────┬──────────┘ - │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC - │ (JSON ~1 KB) │ (NIXL meta ~1.5 KB) │ (NIXL meta ~1.5 KB) - ▼ ▼ ▼ - ┌──────────┐ ┌──────────┐ ┌──────────┐ - │Encoder-0 │──────────│Denoiser-0│────────────│ VAE-0 │ - │ GPU 0 │ NIXL │ GPU 1,2 │ NIXL │ GPU 5 │ - └──────────┘ RDMA └──────────┘ RDMA └──────────┘ - ┌──────────┐ ┌──────────┐ ┌──────────┐ - │Encoder-1 │──────────│Denoiser-1│────────────│ VAE-1 │ - │ GPU 3 │ NIXL │ GPU 4,5 │ NIXL │ GPU 7 │ - └──────────┘ RDMA └──────────┘ RDMA └──────────┘ - ...↕ ...↕ ...↕ - dynamic dynamic dynamic - add/remove add/remove add/remove -``` - -**Key points:** -- Workers self-register with etcd; the orchestrator discovers them - automatically. -- Dynamo RPC carries only small JSON payloads and NIXL metadata - (~1.5 KB). Actual tensors (embeddings, latents) transfer GPU-to-GPU - via NIXL RDMA, never touching the CPU or RPC channel. -- Each stage can have a different number of workers and GPU count. - The denoiser typically uses TP > 1 (multi-GPU), while encoder and VAE - each use a single GPU. - -### 2.2 Process Model - -Each worker is an independent OS process with no shared state: - -``` -Dynamo Worker Process (e.g. encoder_worker.py) -├── @dynamo_worker ← Dynamo runtime bootstrap -│ ├── serve_endpoint("generate") ← Dynamo RPC from orchestrator -│ │ └── handle_generate() -│ │ └── StageClient.forward() ← ZMQ to local backend subprocess -│ └── serve_endpoint("health") -│ -└── Backend subprocess (spawned at startup) - └── Inference engine (any backend) - ├── ReceiveStage ← RDMA-pull tensors from previous stage - ├── ComputeStage ← model-specific inference - └── SendStage ← register output tensors as RDMA-readable -``` - -- The Dynamo worker process handles RPC and control-plane logic. -- The backend subprocess runs the actual model inference. It can be any - inference engine — SGLang, vLLM, a custom PyTorch loop, etc. -- Communication between the two is via ZMQ REQ/REP (same-host IPC). - -### 2.3 Loose Coupling & Dynamic Scaling - -Workers are completely independent — they know nothing about each other -or the orchestrator: - -- **Startup:** Worker process starts → registers with etcd (via Dynamo - runtime) → orchestrator auto-discovers the new instance. -- **Shutdown:** Worker process exits → etcd lease expires → - orchestrator stops routing to it. -- **Add worker:** Start a new process on any available GPU → etcd → - orchestrator sees it within seconds. No restart, no reconfiguration. -- **Remove worker:** Kill the process → etcd lease expires → traffic - drains naturally. -- **Auto-scale:** An external controller can monitor queue depth per - stage (`GET /pipeline/status`) and spawn/kill workers as needed. - -### 2.4 Auto Routing - -Each stage has a `WorkerManager` that maintains an idle-pool queue: - -``` -WorkerManager("denoiser", client, [0, 1, 2]) -│ -├── _idle_queue: asyncio.Queue ← [0, 1, 2] initially -│ -├── acquire_worker() → int ← blocks until a worker is idle -├── dispatch(wid, rid, json) ← client.direct(json, wid) -│ └── on completion/failure → release wid back to _idle_queue -│ -└── status() → dict ← per-worker completed/latency/state -``` - -- `acquire_worker()` blocks the caller when all workers are busy - (backpressure). -- `dispatch()` sends to a specific worker via `client.direct()` and - tracks busy/idle state for observability. -- `dispatch_with_retry()` wraps this: on failure, acquires a different - worker and retries (configurable via `STAGE_DISPATCH_RETRIES`). - -### 2.5 Request Flow - -A single video generation request follows this path: - -``` -Client Orchestrator Encoder-k Denoiser-j VAE-i - │ │ │ │ │ - │─POST /v1/videos/──────►│ │ │ │ - │ generations │ │ │ │ - │ [acquire semaphore] │ │ │ - │ │ │ │ │ - │ │──EncoderRequest──────►│ │ │ - │ │ (prompt, cfg) [encode text] │ │ - │ │◄──NIXL metadata──────│ │ │ - │ │ (~1.5 KB) │ │ │ - │ │ │ │ │ - │ │──DenoiserRequest───────────────────►│ │ - │ │ (NIXL meta, params) [RDMA pull embeddings]│ - │ │ [denoise N steps] │ - │ │◄──NIXL metadata────────────────────│ │ - │ │ │ │ - │ │──VAEDecodeRequest──────────────────────────────►│ - │ │ (NIXL meta, req_id) [RDMA pull] │ - │ │ [decode] │ - │ │◄──{video_path}──────────────────────────────────│ - │ [release semaphore] │ │ │ - │◄──{url, timings}───────│ │ │ │ +| Stage-level scaling | Worker instances per stage, discovered from etcd | +| Throughput | Global admission semaphore + per-stage worker pools | +| Low transfer overhead | RPC sends metadata only (~1.5 KB), tensors use RDMA | +| Fault containment | Retry on another worker (`STAGE_DISPATCH_RETRIES`) | +| Runtime elasticity | Add/remove workers without orchestrator restart | + +## 2. Architecture and Flow + +### 2.1 Framework Diagram + +```mermaid +flowchart TB + classDef cp fill:#eaf2ff,stroke:#2b5ec8,stroke-width:1.2px,color:#0f172a; + classDef dp fill:#fff1f2,stroke:#c0392b,stroke-width:1.2px,color:#0f172a; + classDef stage fill:#f8fafc,stroke:#475569,stroke-width:1.2px,color:#0f172a; + classDef ext fill:#eef2f7,stroke:#64748b,stroke-width:1.2px,color:#0f172a; + + Client[Client] + Etcd[(etcd)] + Orch[Orchestrator] + + subgraph CP[Control Plane] + direction TB + Orch -->|Dynamo RPC + metadata| EncPool + Orch -->|Dynamo RPC + metadata| DenPool + Orch -->|Dynamo RPC + metadata| VaePool + end + + subgraph EncPool[Encoder Workers] + direction TB + E1[encoder-0] + E2[encoder-1] + end + + subgraph DenPool[Denoiser Workers] + direction TB + D1[denoiser-0] + D2[denoiser-1] + end + + subgraph VaePool[VAE Workers] + direction TB + V1[vae-0] + V2[vae-1] + end + + subgraph DP[Data Plane] + direction LR + EncPool -->|NIXL RDMA embeddings| DenPool + DenPool -->|NIXL RDMA latents| VaePool + end + + Client -->|HTTP| Orch + Etcd -->|register/discover| Orch + VaePool -->|video.mp4| Client + DenPool -. notify free buffer .-> EncPool + VaePool -. notify free buffer .-> DenPool + + class CP cp; + class DP dp; + class EncPool,DenPool,VaePool,Orch,Client,Etcd,E1,E2,D1,D2,V1,V2 stage; ``` -Only ~1.5 KB of NIXL metadata travels over Dynamo RPC between stages. -The actual tensor data (embeddings: ~tens of MB, latents: ~hundreds of -MB) transfers GPU-to-GPU via NIXL RDMA without CPU involvement. - -### 2.6 Pipeline Parallelism - -An `asyncio.Semaphore(pipeline_depth)` gates admission. Each stage has -its own independent worker pool, so multiple requests overlap: - +**Read this diagram as two planes:** +- **Control plane:** client request, worker discovery, and RPC dispatch. +- **Data plane:** embeddings and latents transferred GPU-to-GPU through NIXL. +- **Buffer lifecycle:** downstream stage notifies upstream to release sender buffers. + +### 2.2 Request Flow Diagram + +```mermaid +sequenceDiagram + autonumber + participant C as Client + participant O as Orchestrator + participant E as Encoder worker + participant D as Denoiser worker + participant V as VAE worker + + rect rgb(235, 245, 255) + Note over C,O,E,D,V: Control Plane (HTTP + Dynamo RPC + metadata) + C->>O: POST /v1/videos/generations + O->>E: EncoderRequest(prompt, cfg) + E-->>O: transfer_meta (embeddings, ~1.5KB) + O->>D: DenoiserRequest(meta, params) + D-->>O: transfer_meta (latents, ~1.5KB) + O->>V: VAEDecodeRequest(meta, request_id) + V-->>O: video_path + O-->>C: {url, timings} + end + + rect rgb(255, 241, 242) + Note over E,D,V: Data Plane (NIXL RDMA tensor transfer) + D->>E: RDMA pull embeddings (GPU->GPU) + V->>D: RDMA pull latents (GPU->GPU) + end + + Note over O,E,D,V: Routing: round-robin over idle workers + Note over O,E,D,V: Overlap: requests run concurrently across stages ``` -Time ──────────────────────────────────────────────────► -Req A: [Enc-0][======Den-0======][VAE-0] -Req B: [Enc-1][======Den-1======][VAE-1] -Req C: [Enc-0][======Den-0======][VAE-0] -Req D: [Enc-1][======Den-1======][VAE-1] -``` +**Single request path (with plane separation):** +1. Client sends `POST /v1/videos/generations`. +2. Orchestrator dispatches `EncoderRequest` and receives embedding metadata. +3. Orchestrator dispatches `DenoiserRequest` with metadata + inference params. +4. Orchestrator dispatches `VAEDecodeRequest` and receives `video_path`. +5. Client receives `{url, timings}`. -- `pipeline_depth` defaults to `MAX_PIPELINE_DEPTH` (env, default 4) - or the total number of workers across all stages. -- Each request independently acquires workers from each stage's pool. -- The denoiser is typically the bottleneck (50 diffusion steps), so - encoder and VAE workers are freed quickly to serve other requests. +**How routing and overlap work now:** +- Routing policy: round-robin among currently idle workers in each stage. +- Overlap behavior: requests can occupy different stages concurrently + (Encode/Denoise/Decode overlap), which hides most control/data transfer cost. +### 2.3 Runtime Model -## 3. Component Design +Each stage worker is an isolated process: +- **Dynamo worker process:** serves `generate` and `health` endpoints. +- **Backend subprocess:** executes stage pipeline through `StageClient` (ZMQ REQ/REP). +- **No shared memory/state between workers:** scaling and failure domains stay clean. -### 3.1 Orchestrator +## 3. Key Components -`orchestrator/run_disagg.py` — aiohttp HTTP server that chains stages. +### 3.1 Orchestrator (`orchestrator/run_disagg.py`) -**Endpoints:** +Responsibilities: +- Initialize stage clients and discover worker instance IDs. +- Enforce global concurrency via `asyncio.Semaphore(pipeline_depth)`. +- Chain Encoder -> Denoiser -> VAE with per-stage timing. +- Route each stage call to an idle worker using round-robin queue order. +- Expose API and observability endpoints. -| Method | Path | Description | +| Method | Path | Purpose | |---|---|---| -| `POST` | `/v1/videos/generations` | Submit generation request | +| `POST` | `/v1/videos/generations` | Run full pipeline | | `GET` | `/health` | Orchestrator liveness | -| `GET` | `/health/stages` | Per-stage health (queries each worker) | -| `GET` | `/pipeline/status` | Active requests, queue depth, latencies | -| `GET` | `/videos/` | Serve generated video file | - -**Key components:** +| `GET` | `/health/stages` | Fan-out health to stage workers | +| `GET` | `/pipeline/status` | Active requests, queue depth, worker stats | +| `GET` | `/videos/` | Return generated MP4 | -- `PipelineTracker` — Tracks active requests per stage, completed/failed - counts, and rolling average stage latencies. -- `dispatch_with_retry(mgr, request_id, json)` — Acquire worker → - dispatch → on failure retry on a different worker (up to - `STAGE_DISPATCH_RETRIES` attempts). -- `admission = asyncio.Semaphore(pipeline_depth)` — Limits concurrent - in-flight requests across the entire pipeline. +### 3.2 Worker Pooling (`orchestrator/worker_manager.py`) -### 3.2 WorkerManager +`WorkerManager` maintains: +- idle queue (`acquire_worker()` blocks on backpressure), +- direct dispatch to worker ID (`client.direct(...)`), +- per-worker runtime stats (`status`, `completed`, `avg_latency_s`, `queue_depth`). -`orchestrator/worker_manager.py` — Per-stage worker pool with -busy/idle tracking. +### 3.3 Stage Workers (`workers/*_worker.py`) -```python -class WorkerManager: - def __init__(self, stage_name: str, client, worker_ids: List[int]): ... - async def acquire_worker(self) -> int: ... # blocks until idle - async def dispatch(self, worker_id, request_id, request_json) -> (dict, float): ... - def status(self) -> dict: ... # per-worker stats -``` - -- Backed by `asyncio.Queue` (idle pool) — `acquire_worker()` awaits - the queue, `dispatch()` returns the worker to it on completion. -- `_call_direct()` sends to a specific worker via - `client.direct(json, worker_id)`. -- `status()` returns per-worker `{id, status, request_id, completed, - avg_latency_s}` plus stage-level `{queue_depth, completed, failed}`. +Common pattern: +- launch backend scheduler subprocess (`launch_stage_server(...)`), +- run Dynamo RPC handlers (`generate`, `health`), +- forward requests through `StageClient.forward(...)`. -### 3.3 Worker Interface +Design direction: keep worker/orchestrator logic stable and swap only backend +handler implementation (`sglang`, `diffusers`, or others). -All three workers follow an identical pattern: +Stage IO contract: -```python -@dynamo_worker(enable_nats=False) -async def worker(runtime: DistributedRuntime): - # 1. Launch backend subprocess - processes, client, server_args = launch_stage_server( - MODEL_PATH, required_modules, build_stage_fn, SCHEDULER_PORT, - ) - - # 2. Define Dynamo RPC handlers - async def handle_generate(request, context): - output = await client.forward([build_req(...)]) - yield result_dict # JSON response with NIXL metadata or output - - async def handle_health(request, context): - yield {"status": "ok", "stage": "..."} - - # 3. Serve endpoints - gen_ep = runtime.endpoint("disagg_diffusion..generate") - health_ep = runtime.endpoint("disagg_diffusion..health") - await asyncio.gather( - gen_ep.serve_endpoint(handle_generate), - health_ep.serve_endpoint(handle_health), - ) -``` - -Each worker: -- Spawns a backend subprocess (which loads model weights and runs - inference). -- Bridges Dynamo RPC ↔ backend via `StageClient` (async ZMQ). -- Handles NIXL metadata forwarding: output from one stage's send - becomes the next stage's receive metadata. -- Includes a ZMQ fallback path: when NIXL is unavailable, tensors are - serialized via `torch.save()` + base64 over the RPC channel. - -**Stage-specific behavior:** - -| Stage | Input | Compute | Output | -|---|---|---|---| -| Encoder | prompt text | Text encoding | NIXL metadata (embeddings) | -| Denoiser | NIXL metadata (embeddings) + params | RDMA pull → N denoise steps | NIXL metadata (latents) | -| VAE | NIXL metadata (latents) | RDMA pull → VAE decode | video file path | - -### 3.4 Tensor Transfer (NIXL) - -`workers/nixl_transfer.py` — GPU-direct RDMA transfer between stages. - -**Sender (`NixlTensorSender`):** - -```python -sender = NixlTensorSender() -meta = sender.send({"latents": tensor}) # → ~1.5 KB metadata dict -``` - -1. Flatten all tensors into a single contiguous GPU buffer - (`torch.cat`). -2. Create a NIXL `Descriptor` wrapping the flat buffer. -3. Register as `readable` via `connector.create_readable(descriptor)`. -4. Return metadata: tensor keys, shapes, dtypes, NIXL descriptor. -5. Hold `(readable, flat_buffer, timestamp)` in `_pending` list. -6. `_sweep()` polls `readable.status` on each subsequent `send()` — - releases buffers on `COMPLETE` or timeout (`NIXL_BUFFER_TIMEOUT_S`, - default 120s). - -**Receiver (`NixlTensorReceiver`):** - -```python -receiver = NixlTensorReceiver() -tensors = receiver.recv(meta, device="cuda") # → {"latents": tensor} -``` - -1. Parse metadata to compute total byte size and per-tensor specs. -2. Allocate a flat `torch.uint8` buffer directly on the target GPU. -3. `connector.begin_read(rdma_meta, descriptor)` → RDMA pull from - sender's GPU. -4. `read_op.wait_for_completion()` — blocks until transfer finishes. -5. Slice the flat buffer into individual tensors using stored shapes - and dtypes. - -### 3.5 Protocol Types - -`workers/protocol.py` — Pydantic models for Dynamo RPC serialization. - -```python -class EncoderRequest(BaseModel): - prompt: str - negative_prompt: str = "" - guidance_scale: float = 1.0 - -class DenoiserRequest(BaseModel): - transfer_meta: Dict[str, Any] # NIXL metadata from encoder - tensor_data: Dict[str, Any] = {} # ZMQ fallback - height: int = 544 - width: int = 960 - num_frames: int = 61 - num_inference_steps: int = 50 - guidance_scale: float = 1.0 - seed: int = 42 - -class VAEDecodeRequest(BaseModel): - transfer_meta: Dict[str, Any] # NIXL metadata from denoiser - tensor_data: Dict[str, Any] = {} # ZMQ fallback - request_id: str = "" -``` - -Responses carry either `transfer_meta` (NIXL path) or `tensor_data` -(ZMQ fallback) but never both. - - -## 4. Implementation: SGLang Backend - -This section describes the current backend implementation using SGLang's -multimodal generation runtime, with HunyuanVideo as the reference model. -The generic architecture (Sections 1-3) is backend-agnostic — any -inference engine that can run pipeline stages can replace SGLang. - -### 4.1 PartialGPUWorker - -`workers/partial_gpu_worker.py` — Extends SGLang's `GPUWorker` to load -only the modules each stage needs. - -```python -class PartialGPUWorker(GPUWorker): - def __init__(self, required_modules, custom_stages_fn, **kwargs): ... - def init_device_and_model(self): - # 1. Set up distributed environment (TP, SP, CFG parallel) - # 2. build_partial_pipeline() — load only required_modules - # 3. custom_stages_fn(pipeline, server_args) — build stage list - # 4. Register stages with pipeline -``` - -- Overrides only `init_device_and_model()`. All other `GPUWorker` - behavior (forward execution, memory analysis, LoRA) is inherited. -- `build_partial_pipeline()` (`sglang_utils.py`) dynamically creates a - subclass of the model's pipeline that suppresses automatic stage - creation and LoRA initialization, loading only the specified modules. - -**Custom pipeline stages (NIXL integration):** +| Stage | Input | Output | +|---|---|---| +| Encoder | prompt + guidance | `transfer_meta` for embeddings | +| Denoiser | embedding metadata + denoise params | `transfer_meta` for latents | +| VAE | latent metadata + request ID | `video_path` | -- `NixlReceiveStage(PipelineStage)` — Prepended at the start of - denoiser/VAE pipelines. Reads `_nixl_transfer_meta` from the `Req` - and RDMA-pulls tensors. Falls back to device-move for ZMQ path. - Includes retry logic for `REMOTE_DISCONNECT` errors. -- `NixlSendStage(PipelineStage)` — Appended as the last stage in - encoder/denoiser pipelines. Extracts tensors from `Req`, registers - with NIXL, returns `OutputBatch` containing only metadata. +Fallback path: if NIXL is unavailable, tensors are serialized over RPC (`tensor_data`). -**Stage builders** (picklable functions passed to subprocess): +### 3.4 Tensor Transport (`workers/nixl_transfer.py`) -| Function | Stages | -|---|---| -| `build_encoder_stages()` | `TextEncodingStage` → `NixlSendStage` | -| `build_denoiser_stages()` | `NixlReceiveStage` → `LatentPreparationStage` → `TimestepPreparationStage` → `DenoisingStage` → `NixlSendStage` | -| `build_vae_stages()` | `NixlReceiveStage` → `DecodingStage` | +`NixlTensorSender`: +- flattens tensors into one GPU buffer, +- registers a readable descriptor, +- returns compact metadata and tracks pending buffers. -### 4.2 Subprocess Launcher +`NixlTensorReceiver`: +- allocates destination GPU buffer, +- performs RDMA read, +- reconstructs tensors from metadata. -`workers/partial_gpu_worker.py:launch_partial_server()` — Spawns -SGLang Scheduler subprocess(es) with `PartialGPUWorker` monkey-patched -in place of the default `GPUWorker`. +**Buffer strategy (current vs future):** +- Current: request-scoped temporary buffer allocation; simple and sufficient + because measured overhead share is small. +- Future: buffer pooling/manager for tighter latency control and memory reuse + under higher concurrency. -``` -launch_partial_server(server_args, required_modules, custom_stages_fn) -│ -├── For each GPU (rank 0..N-1): -│ ├── Create readiness pipe -│ ├── mp.Process(target=_run_partial_scheduler_process) -│ │ ├── Monkey-patch: sched_mod.GPUWorker = _PatchedGPUWorker -│ │ └── run_scheduler_process(...) ← standard SGLang entry point -│ └── Start process -│ -├── Wire master/slave pipes (TP > 1: rank 0 is master, ranks 1..N are slaves) -├── Wait for all readiness signals -└── Return process list -``` +## 4. SGLang-Specific Integration -`launch_stage_server()` (`sglang_utils.py`) wraps this with config -setup: `patch_hunyuan_config()` → `ServerArgs.from_kwargs()` → -`launch_partial_server()` → `StageClient(endpoint)`. +The back half should be backend-agnostic. The recommended design is a generic +`StageHandler` contract, with framework-specific implementations. -### 4.3 StageClient - -`workers/sglang_utils.py:StageClient` — Async ZMQ REQ/REP client -connecting the Dynamo worker main process to the SGLang Scheduler -subprocess. +### 4.1 Generic StageHandler Abstraction ```python -class StageClient: - def __init__(self, endpoint: str, name: str = ""): ... - async def forward(self, reqs): # send_pyobj → recv_pyobj with timeout - def close(self): ... +class StageHandler: + async def start(self) -> None: ... + async def generate(self, request: dict) -> dict: ... + async def health(self) -> dict: ... + async def shutdown(self) -> None: ... ``` -- `asyncio.Lock` serializes concurrent calls (ZMQ REQ socket is - single-flight). -- Configurable timeout: `STAGE_FORWARD_TIMEOUT_S` (default 120s). - -### 4.4 HunyuanVideo Specifics - -- **Dual encoder detection:** `detect_encoder_modules()` reads - `model_index.json` and auto-detects `text_encoder_2` / `tokenizer_2` - (HunyuanVideo uses Llama + CLIP). Falls back to heuristic for known - model names. -- **HunyuanConfig patching:** `patch_hunyuan_config()` wraps - `HunyuanConfig.__init__` to supply `task_type=T2V` when omitted - (the base class requires it but HunyuanConfig doesn't default it). -- **Triton norm contiguous workaround:** - `_patch_triton_norm_contiguous()` wraps SGLang's triton - `norm_infer` to call `.contiguous()` on non-contiguous tensors from - HunyuanVideo's attention reshapes, avoiding the triton kernel - assertion `x.stride(-1) == 1`. -- **Component config sync:** `_sync_all_component_configs()` reads - `config.json` for every component in `model_index.json` and updates - `server_args.pipeline_config`, ensuring correct parameters (e.g. - `z_dim`) even for components whose weights are not loaded by the - current stage. - - -## 5. Roadmap - -- [x] Multi-worker scaling — N workers per stage, etcd auto-discovery -- [x] Pipeline parallelism — overlapping requests across stages -- [x] NIXL GPU-direct transfer — GPU-to-GPU RDMA, only metadata over RPC -- [ ] Runtime scaling — add/remove workers without restart; - external auto-scaler integration based on queue depth -- [ ] Smart routing — load-aware dispatch (not just idle-queue), - affinity-based routing, request priority -- [ ] Fault tolerance — worker health-check + eviction, - dead worker detection, graceful degradation -- [ ] Streaming output — stream decoded frames as produced -- [ ] Metrics & observability — per-stage latency histograms, - GPU utilization, NIXL throughput, Prometheus export -- [ ] Request cancellation — cancel in-flight requests, free GPU - immediately -- [ ] Multi-model support — OneVideo, omni-modal pipelines, - heterogeneous stage graphs (not just linear 3-stage) - - -## 6. File Map - -``` -examples/disagg_diffusion/ -├── DESIGN.md ← this document -├── README.md ← usage guide and quick start -├── run_all.sh ← launch etcd + all workers + orchestrator -├── stress_test.sh ← concurrent load testing script -│ -├── orchestrator/ -│ ├── run_disagg.py ← HTTP server, PipelineTracker, dispatch_with_retry -│ └── worker_manager.py ← WorkerManager: idle-pool, acquire/dispatch/status -│ -├── workers/ -│ ├── __init__.py -│ ├── protocol.py ← Pydantic request/response models (Dynamo RPC) -│ ├── encoder_worker.py ← Encoder Dynamo worker (text encoding → NIXL send) -│ ├── denoiser_worker.py ← Denoiser Dynamo worker (NIXL recv → denoise → NIXL send) -│ ├── vae_worker.py ← VAE Dynamo worker (NIXL recv → decode → save video) -│ ├── nixl_transfer.py ← NixlTensorSender / NixlTensorReceiver (RDMA) -│ ├── partial_gpu_worker.py ← PartialGPUWorker, NixlSend/ReceiveStage, subprocess launcher -│ └── sglang_utils.py ← StageClient, build_partial_pipeline, launch_stage_server -│ -└── validate/ - └── validate_split.py ← validation script for split correctness -``` +Worker responsibility stays unchanged: +- parse Dynamo request/response protocol, +- call `handler.generate(...)`, +- return stage output (`transfer_meta` or `video_path`). + +Backend-specific logic moves into handlers: +- process launch and model init, +- stage graph / pipeline execution, +- tensor extraction/injection details. + +### 4.2 Current Handler: SGLang + +Current implementation maps to an SGLang handler using +`workers/partial_gpu_worker.py`: +- `NixlReceiveStage` at denoiser/VAE entry, +- `NixlSendStage` at encoder/denoiser exit. + +Stage builders: +- `build_encoder_stages()`: `TextEncodingStage -> NixlSendStage` +- `build_denoiser_stages()`: `NixlReceiveStage -> ... -> DenoisingStage -> NixlSendStage` +- `build_vae_stages()`: `NixlReceiveStage -> DecodingStage` + +### 4.3 Alternative Handler: Diffusers (Planned) + +A diffusers handler should implement the same contract and keep the same +orchestrator protocol: +- input/output request schema unchanged, +- same NIXL metadata fields for stage handoff, +- same health/reporting behavior. + +This allows swapping `sglang` -> `diffusers` without changing orchestrator +routing, worker manager, or external API. + +Decision note: keep orchestrator as a standalone component for now; evaluate +merging into a router layer only after smart-routing requirements justify it. + +## 5. Roadmap (Condensed) + +- [x] Multi-worker per stage with etcd discovery +- [x] Pipeline overlap with global admission control +- [x] NIXL metadata-only RPC + GPU-direct tensor transfer +- [ ] Runtime auto-scaling from `/pipeline/status` +- [ ] Smarter routing (affinity/priority/load-aware) +- [ ] Fault tolerance hardening and graceful degradation +- [ ] Streaming decode output + richer observability +- [ ] Evaluate merging orchestrator into router layer +- [ ] Explore diffusion-based smart router for stage/worker selection +- [ ] Formalize `StageHandler` interface and migrate SGLang to handler plugin +- [ ] Add Diffusers handler with parity tests against SGLang outputs diff --git a/examples/disagg_diffusion/orchestrator/request_validation.py b/examples/disagg_diffusion/orchestrator/request_validation.py new file mode 100644 index 000000000000..bcb4dd18d81e --- /dev/null +++ b/examples/disagg_diffusion/orchestrator/request_validation.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request validation helpers for disaggregated diffusion orchestrator.""" + +from protocol import GenerateRequest + + +def validate_generate_http_body(body: dict) -> dict: + """Validate and normalize POST /v1/videos/generations request body.""" + req = GenerateRequest.model_validate(body) + req_dict = req.model_dump() + resp_format = body.get("response_format", "url") + if resp_format not in {"url", "b64_json"}: + raise ValueError("response_format must be 'url' or 'b64_json'") + req_dict["response_format"] = resp_format + return req_dict diff --git a/examples/disagg_diffusion/orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py index 384e69090fef..05b777bdaf6f 100755 --- a/examples/disagg_diffusion/orchestrator/run_disagg.py +++ b/examples/disagg_diffusion/orchestrator/run_disagg.py @@ -48,6 +48,8 @@ ) from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 from worker_manager import WorkerManager # noqa: E402 +from pydantic import ValidationError # noqa: E402 +from request_validation import validate_generate_http_body # noqa: E402 logger = logging.getLogger(__name__) @@ -148,17 +150,24 @@ async def worker(runtime: DistributedRuntime): # Wait for additional workers that may still be registering. # wait_for_instances() returns after the first instance; model loading - # times vary, so poll until the count stabilizes or timeout. + # times vary, so observe all stages for a short settle window. WORKER_SETTLE_S = int(os.environ.get("WORKER_SETTLE_S", "30")) if WORKER_SETTLE_S > 0: import time as _time deadline = _time.monotonic() + WORKER_SETTLE_S - prev_count = 0 + prev_counts = (-1, -1, -1) while _time.monotonic() < deadline: - cur = len(denoiser_client.instance_ids()) - if cur > prev_count: - prev_count = cur - logger.info("Discovered %d denoiser(s) so far, waiting for more…", cur) + cur_counts = ( + len(encoder_client.instance_ids()), + len(denoiser_client.instance_ids()), + len(vae_client.instance_ids()), + ) + if cur_counts != prev_counts: + logger.info( + "Discovered workers so far: encoder=%d denoiser=%d vae=%d", + cur_counts[0], cur_counts[1], cur_counts[2], + ) + prev_counts = cur_counts await asyncio.sleep(2) logger.info("Worker settle period done (%ds)", WORKER_SETTLE_S) @@ -285,10 +294,20 @@ async def handle_generate(request: dict) -> dict: async def handle_post(http_request: web.Request) -> web.Response: try: body = await http_request.json() - if "prompt" not in body: - return web.json_response({"error": "missing 'prompt' field"}, status=400) - result = await handle_generate(body) + except Exception: + return web.json_response({"error": "invalid JSON body"}, status=400) + + try: + req_dict = validate_generate_http_body(body) + result = await handle_generate(req_dict) return web.json_response(result) + except ValidationError as e: + return web.json_response( + {"error": "invalid request", "details": e.errors()}, + status=400, + ) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) except Exception as e: logger.error("Request failed: %s", e, exc_info=True) return web.json_response({"error": str(e)}, status=500) diff --git a/examples/disagg_diffusion/tests/test_protocol_models.py b/examples/disagg_diffusion/tests/test_protocol_models.py new file mode 100644 index 000000000000..993f45fab85a --- /dev/null +++ b/examples/disagg_diffusion/tests/test_protocol_models.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +from pathlib import Path +import sys + + +ROOT = Path(__file__).resolve().parents[1] +WORKERS_DIR = ROOT / "workers" +sys.path.insert(0, str(WORKERS_DIR)) + +from protocol import DenoiserResponse, EncoderResponse, GenerateRequest, VAEDecodeRequest # noqa: E402 + + +def test_protocol_mutable_defaults_are_isolated(): + a = EncoderResponse() + b = EncoderResponse() + a.transfer_meta["x"] = 1 + a.shapes["latents"] = [1, 2, 3] + + assert b.transfer_meta == {} + assert b.shapes == {} + + +def test_protocol_other_dict_defaults_are_isolated(): + a = DenoiserResponse() + b = DenoiserResponse() + a.transfer_meta["k"] = "v" + a.shape.append(42) + + assert b.transfer_meta == {} + assert b.shape == [] + + x = VAEDecodeRequest(transfer_meta={}) + y = VAEDecodeRequest(transfer_meta={}) + x.tensor_data["t"] = "blob" + assert y.tensor_data == {} + + +def test_generate_request_seed_is_optional(): + req = GenerateRequest(prompt="hello") + assert req.seed is None diff --git a/examples/disagg_diffusion/tests/test_request_validation.py b/examples/disagg_diffusion/tests/test_request_validation.py new file mode 100644 index 000000000000..084232fbf0ae --- /dev/null +++ b/examples/disagg_diffusion/tests/test_request_validation.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +from pathlib import Path +import sys + +import pytest +from pydantic import ValidationError + + +ROOT = Path(__file__).resolve().parents[1] +WORKERS_DIR = ROOT / "workers" +ORCH_DIR = ROOT / "orchestrator" +sys.path.insert(0, str(WORKERS_DIR)) +sys.path.insert(0, str(ORCH_DIR)) + +from request_validation import validate_generate_http_body # noqa: E402 + + +def test_validate_generate_http_body_accepts_minimal_payload(): + out = validate_generate_http_body({"prompt": "a cat"}) + assert out["prompt"] == "a cat" + assert out["response_format"] == "url" + assert out["seed"] is None + + +def test_validate_generate_http_body_rejects_invalid_response_format(): + with pytest.raises(ValueError, match="response_format"): + validate_generate_http_body( + { + "prompt": "a cat", + "response_format": "invalid", + } + ) + + +def test_validate_generate_http_body_rejects_invalid_schema(): + with pytest.raises(ValidationError): + validate_generate_http_body({"height": 544}) diff --git a/examples/disagg_diffusion/workers/protocol.py b/examples/disagg_diffusion/workers/protocol.py index 2a059ff2eb64..c86082f2af34 100644 --- a/examples/disagg_diffusion/workers/protocol.py +++ b/examples/disagg_diffusion/workers/protocol.py @@ -11,7 +11,7 @@ from typing import Any, Dict, List, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field # --------------------------------------------------------------------------- @@ -25,8 +25,8 @@ class EncoderRequest(BaseModel): class EncoderResponse(BaseModel): - transfer_meta: Dict[str, Any] = {} - shapes: Dict[str, List[int]] = {} + transfer_meta: Dict[str, Any] = Field(default_factory=dict) + shapes: Dict[str, List[int]] = Field(default_factory=dict) # --------------------------------------------------------------------------- @@ -35,7 +35,7 @@ class EncoderResponse(BaseModel): class DenoiserRequest(BaseModel): transfer_meta: Dict[str, Any] - tensor_data: Dict[str, Any] = {} + tensor_data: Dict[str, Any] = Field(default_factory=dict) height: int = 544 width: int = 960 num_frames: int = 61 @@ -45,8 +45,8 @@ class DenoiserRequest(BaseModel): class DenoiserResponse(BaseModel): - transfer_meta: Dict[str, Any] = {} - shape: List[int] = [] + transfer_meta: Dict[str, Any] = Field(default_factory=dict) + shape: List[int] = Field(default_factory=list) # --------------------------------------------------------------------------- @@ -55,7 +55,7 @@ class DenoiserResponse(BaseModel): class VAEDecodeRequest(BaseModel): transfer_meta: Dict[str, Any] - tensor_data: Dict[str, Any] = {} + tensor_data: Dict[str, Any] = Field(default_factory=dict) request_id: str = "" @@ -76,7 +76,7 @@ class GenerateRequest(BaseModel): num_frames: int = 61 num_inference_steps: int = 50 guidance_scale: float = 1.0 - seed: int = 42 + seed: Optional[int] = None # --------------------------------------------------------------------------- From e2e92ebb5470b4c3702d5c9d1c9a495075870c2b Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 19:53:14 +0800 Subject: [PATCH 21/22] docs: rewrite DESIGN.md with ASCII diagrams and current implementation Restore ASCII art architecture/flow diagrams (more portable than mermaid). Update to reflect current implementation: - PersistentConnector pattern and agent lifecycle - One-sided RDMA pull with buffer lifecycle diagram - TP broadcast for multi-GPU denoisers - UCX_TLS configuration for intra-node transfers - StageClient ZMQ socket recovery - Worker settle period for discovery - Measured performance numbers (8-GPU HunyuanVideo) Co-Authored-By: Claude Opus 4.6 --- examples/disagg_diffusion/DESIGN.md | 648 ++++++++++++++++++---------- 1 file changed, 426 insertions(+), 222 deletions(-) diff --git a/examples/disagg_diffusion/DESIGN.md b/examples/disagg_diffusion/DESIGN.md index e8bc85f474a4..5bdefa6a0c9e 100644 --- a/examples/disagg_diffusion/DESIGN.md +++ b/examples/disagg_diffusion/DESIGN.md @@ -1,271 +1,475 @@ -# Disaggregated Diffusion Pipeline - Design Document +# Disaggregated Diffusion Pipeline — Design Document -## 1. Goal and Scope +## 1. Overview -This design splits video generation into three independently scalable stages: -Encoder -> Denoiser -> VAE. +### Motivation -The orchestrator handles control-plane scheduling via Dynamo RPC; tensors move -across stages through NIXL GPU-direct RDMA. +Modern diffusion pipelines (text-to-video, text-to-image, omni-modal) are +composed of heterogeneous stages — text encoding, iterative denoising, VAE +decoding — each with fundamentally different compute profiles: -**Current design choices (implementation-aligned):** -- Multi-worker: each stage runs multiple worker instances. -- Workers can be added/removed at runtime through etcd discovery without - orchestrator restart. -- Orchestrator does stage routing; current policy is round-robin over idle - workers (`WorkerManager` queue order). -- Pipeline overlap is used to hide transfer overhead behind compute. -- Tensor buffers are currently temporary allocations per request; this is - acceptable with current latency profile and can be upgraded to pooled/buffer - management later. +- **Different optimization strategies per stage.** Encoders are + memory-bound single-pass transforms; denoisers are compute-bound + multi-step loops that benefit from tensor parallelism; VAE decoders + are memory-intensive but run only once. Forcing all three into a + single process prevents stage-specific tuning (parallelism, batching, + memory management, quantization). -| Goal | Current mechanism | +- **Shifting compute balance.** As diffusion models mature, the DiT no + longer dominates the entire pipeline — faster denoisers (fewer steps, + distilled models) shift the bottleneck to encoding and decoding. + Multi-task pipelines are emerging (e.g. OneVideo: encode + denoise + + decode + audio in one model) where each task demands independent + scaling. + +- **Omni-modal future.** Models that jointly produce video, audio, + image, and text require stage-level separation so each modality's + compute can scale independently without wasting GPU resources. + +### Solution + +Decompose the pipeline into **N independent stages**, each running as a +Dynamo RPC worker on dedicated GPU(s). An orchestrator chains stages +together, using NIXL RDMA for GPU-direct tensor transfer between them. + +| Goal | Mechanism | |---|---| -| Stage-level scaling | Worker instances per stage, discovered from etcd | -| Throughput | Global admission semaphore + per-stage worker pools | -| Low transfer overhead | RPC sends metadata only (~1.5 KB), tensors use RDMA | -| Fault containment | Retry on another worker (`STAGE_DISPATCH_RETRIES`) | -| Runtime elasticity | Add/remove workers without orchestrator restart | - -## 2. Architecture and Flow - -### 2.1 Framework Diagram - -```mermaid -flowchart TB - classDef cp fill:#eaf2ff,stroke:#2b5ec8,stroke-width:1.2px,color:#0f172a; - classDef dp fill:#fff1f2,stroke:#c0392b,stroke-width:1.2px,color:#0f172a; - classDef stage fill:#f8fafc,stroke:#475569,stroke-width:1.2px,color:#0f172a; - classDef ext fill:#eef2f7,stroke:#64748b,stroke-width:1.2px,color:#0f172a; - - Client[Client] - Etcd[(etcd)] - Orch[Orchestrator] - - subgraph CP[Control Plane] - direction TB - Orch -->|Dynamo RPC + metadata| EncPool - Orch -->|Dynamo RPC + metadata| DenPool - Orch -->|Dynamo RPC + metadata| VaePool - end - - subgraph EncPool[Encoder Workers] - direction TB - E1[encoder-0] - E2[encoder-1] - end - - subgraph DenPool[Denoiser Workers] - direction TB - D1[denoiser-0] - D2[denoiser-1] - end - - subgraph VaePool[VAE Workers] - direction TB - V1[vae-0] - V2[vae-1] - end - - subgraph DP[Data Plane] - direction LR - EncPool -->|NIXL RDMA embeddings| DenPool - DenPool -->|NIXL RDMA latents| VaePool - end - - Client -->|HTTP| Orch - Etcd -->|register/discover| Orch - VaePool -->|video.mp4| Client - DenPool -. notify free buffer .-> EncPool - VaePool -. notify free buffer .-> DenPool - - class CP cp; - class DP dp; - class EncPool,DenPool,VaePool,Orch,Client,Etcd,E1,E2,D1,D2,V1,V2 stage; +| Stage-level scaling | N workers per stage, auto-discovered via etcd | +| Pipeline parallelism | Semaphore admission, independent worker pools | +| GPU-direct transfer | NIXL RDMA — only ~1.5 KB metadata over RPC | +| Loose coupling | Workers are independent processes, no shared state | +| Dynamic scaling | Add/remove workers at runtime, no restart | +| Auto routing | Idle-queue dispatch, backpressure, retry on failure | + + +## 2. Architecture + +### 2.1 Architecture Diagram + +``` + ┌───────────┐ + │ etcd │ + │ registry │ + └─────┬─────┘ + register │ discover + ┌───────────────────────────────┼───────────────────────────────┐ + │ Orchestrator │ + │ │ + │ HTTP ──► handle_generate() ──► dispatch_with_retry() │ + │ PipelineTracker per-stage WorkerManager │ + │ Semaphore(depth) acquire → direct() → release │ + │ │ + └──────┬──────────────────────┬──────────────────────┬──────────┘ + │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC + │ (JSON ~1 KB) │ (NIXL meta ~1.5 KB) │ (NIXL meta ~1.5 KB) + ▼ ▼ ▼ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-0 │ │Denoiser-0│ │ VAE-0 │ + │ GPU 0 │ │ GPU 1,2 │ │ GPU 7 │ + └────┬─────┘ └──┬───┬───┘ └────┬─────┘ + │ NIXL RDMA pull │ │ NIXL RDMA pull │ + │◄──────────────────┘ └────────────────────►│ + │ (embeddings) (latents) │ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-1 │ │Denoiser-1│ │ VAE-1 │ + │ GPU 3 │ │ GPU 3,4 │ │ ... │ + └──────────┘ └──────────┘ └──────────┘ + ...↕ ┌──────────┐ ...↕ + dynamic │Denoiser-2│ dynamic + add/remove │ GPU 5,6 │ add/remove + └──────────┘ + ...↕ + dynamic + add/remove +``` + +**Key points:** +- Workers self-register with etcd; the orchestrator discovers them + automatically. A settle period (`WORKER_SETTLE_S`, default 30s) + waits for slow-starting workers after the first instance appears. +- Dynamo RPC carries only small JSON payloads and NIXL metadata + (~1.5 KB). Actual tensors (embeddings, latents) transfer GPU-to-GPU + via NIXL RDMA, never touching the CPU or RPC channel. +- **RDMA is receiver-initiated (one-sided pull):** the denoiser pulls + embeddings from the encoder, the VAE pulls latents from the denoiser. + After the pull completes, the sender receives a completion + notification and releases the GPU buffer. +- Each stage can have a different number of workers and GPU count. + The denoiser typically uses TP > 1 (multi-GPU), while encoder and VAE + each use a single GPU. + +### 2.2 Process Model + +Each worker is an independent OS process with no shared state: + +``` +Dynamo Worker Process (e.g. denoiser_worker.py) +├── @dynamo_worker ← Dynamo runtime bootstrap +│ ├── serve_endpoint("generate") ← Dynamo RPC from orchestrator +│ │ └── handle_generate() +│ │ └── StageClient.forward() ← ZMQ to local backend subprocess +│ └── serve_endpoint("health") +│ +└── Backend subprocess(es) (spawned at startup, TP=2 → 2 processes) + ├── Rank 0 (master, handles ZMQ + NIXL) + │ ├── NixlReceiveStage ← RDMA-pull from previous stage + │ ├── ComputeStage(s) ← model inference (TP all-reduce) + │ └── NixlSendStage ← register output as RDMA-readable + └── Rank 1 (slave, TP compute only) + ├── NixlReceiveStage ← receives tensors via TP broadcast + ├── ComputeStage(s) ← model inference (TP all-reduce) + └── NixlSendStage ← no-op (only rank 0 sends) +``` + +- The Dynamo worker process handles RPC and control-plane logic. +- The backend subprocess runs the actual model inference. +- Communication between the two is via ZMQ REQ/REP (same-host IPC). +- With TP > 1: rank 0 does the NIXL pull then broadcasts tensors to + other ranks via `torch.distributed.broadcast`. This is transparent + to the compute stages — TP is internal to the worker. + +### 2.3 Request Flow + +A single video generation request follows this path: + +``` +Client Orchestrator Encoder-k Denoiser-j VAE-i + │ │ │ │ │ + │─POST /v1/videos/──────►│ │ │ │ + │ generations │ │ │ │ + │ [acquire semaphore] │ │ │ + │ │ │ │ │ + │ │──EncoderRequest──────►│ │ │ + │ │ (prompt, cfg) [encode text] │ │ + │ │ [register readable] │ │ + │ │◄──NIXL metadata──────│ │ │ + │ │ (~1.5 KB) [hold buffer] │ │ + │ │ │ │ │ + │ │──DenoiserRequest───────────────────►│ │ + │ │ (NIXL meta, params) │ │ + │ │ │ [RDMA pull embeddings] │ + │ │ │─────────────►│ │ + │ │ [completion notif] │ │ + │ │ [release buffer] │ │ + │ │ │ [denoise N steps] │ + │ │ │ [register readable] │ + │ │◄──NIXL metadata────────────────────│ │ + │ │ │ [hold buffer] │ + │ │ │ │ │ + │ │──VAEDecodeRequest──────────────────────────────►│ + │ │ (NIXL meta, req_id) │ │ + │ │ │ │ [RDMA pull latents] + │ │ │ │───────────►│ + │ │ │ [completion notif] │ + │ │ │ [release buffer] [decode] + │ │◄──{video_path}──────────────────────────────────│ + │ [release semaphore] │ │ │ + │◄──{url, timings}───────│ │ │ │ ``` -**Read this diagram as two planes:** -- **Control plane:** client request, worker discovery, and RPC dispatch. -- **Data plane:** embeddings and latents transferred GPU-to-GPU through NIXL. -- **Buffer lifecycle:** downstream stage notifies upstream to release sender buffers. - -### 2.2 Request Flow Diagram - -```mermaid -sequenceDiagram - autonumber - participant C as Client - participant O as Orchestrator - participant E as Encoder worker - participant D as Denoiser worker - participant V as VAE worker - - rect rgb(235, 245, 255) - Note over C,O,E,D,V: Control Plane (HTTP + Dynamo RPC + metadata) - C->>O: POST /v1/videos/generations - O->>E: EncoderRequest(prompt, cfg) - E-->>O: transfer_meta (embeddings, ~1.5KB) - O->>D: DenoiserRequest(meta, params) - D-->>O: transfer_meta (latents, ~1.5KB) - O->>V: VAEDecodeRequest(meta, request_id) - V-->>O: video_path - O-->>C: {url, timings} - end - - rect rgb(255, 241, 242) - Note over E,D,V: Data Plane (NIXL RDMA tensor transfer) - D->>E: RDMA pull embeddings (GPU->GPU) - V->>D: RDMA pull latents (GPU->GPU) - end - - Note over O,E,D,V: Routing: round-robin over idle workers - Note over O,E,D,V: Overlap: requests run concurrently across stages +Only ~1.5 KB of NIXL metadata travels over Dynamo RPC between stages. +The actual tensor data (embeddings: ~tens of MB, latents: ~hundreds of +MB) transfers GPU-to-GPU via NIXL RDMA without CPU involvement. + +**Buffer lifecycle:** The sender holds the GPU buffer (via `readable_op` +in `_active_readables`) until the receiver completes the RDMA pull and +the sender receives a completion notification. `_poll_completed()` +checks and releases finished buffers at the start of each new request. + +### 2.4 Pipeline Parallelism + +An `asyncio.Semaphore(pipeline_depth)` gates admission. Each stage has +its own independent worker pool, so multiple requests overlap: + ``` +Time ──────────────────────────────────────────────────► -**Single request path (with plane separation):** -1. Client sends `POST /v1/videos/generations`. -2. Orchestrator dispatches `EncoderRequest` and receives embedding metadata. -3. Orchestrator dispatches `DenoiserRequest` with metadata + inference params. -4. Orchestrator dispatches `VAEDecodeRequest` and receives `video_path`. -5. Client receives `{url, timings}`. +Req A: [Enc-0][======Den-0======][VAE-0] +Req B: [Enc-0][======Den-1======][VAE-0] +Req C: [Enc-0][======Den-2======][VAE-0] +Req D: [Enc-0][======Den-0======][VAE-0] +``` -**How routing and overlap work now:** -- Routing policy: round-robin among currently idle workers in each stage. -- Overlap behavior: requests can occupy different stages concurrently - (Encode/Denoise/Decode overlap), which hides most control/data transfer cost. +- `pipeline_depth` defaults to `MAX_PIPELINE_DEPTH` (env, default 4) + or the total number of workers across all stages. +- Each request independently acquires workers from each stage's pool. +- The denoiser is typically the bottleneck (many diffusion steps), so + encoder and VAE workers are freed quickly to serve other requests. -### 2.3 Runtime Model -Each stage worker is an isolated process: -- **Dynamo worker process:** serves `generate` and `health` endpoints. -- **Backend subprocess:** executes stage pipeline through `StageClient` (ZMQ REQ/REP). -- **No shared memory/state between workers:** scaling and failure domains stay clean. +## 3. Component Design -## 3. Key Components +### 3.1 Orchestrator -### 3.1 Orchestrator (`orchestrator/run_disagg.py`) +`orchestrator/run_disagg.py` — aiohttp HTTP server that chains stages. -Responsibilities: -- Initialize stage clients and discover worker instance IDs. -- Enforce global concurrency via `asyncio.Semaphore(pipeline_depth)`. -- Chain Encoder -> Denoiser -> VAE with per-stage timing. -- Route each stage call to an idle worker using round-robin queue order. -- Expose API and observability endpoints. +**Endpoints:** -| Method | Path | Purpose | +| Method | Path | Description | |---|---|---| -| `POST` | `/v1/videos/generations` | Run full pipeline | +| `POST` | `/v1/videos/generations` | Submit generation request | | `GET` | `/health` | Orchestrator liveness | -| `GET` | `/health/stages` | Fan-out health to stage workers | -| `GET` | `/pipeline/status` | Active requests, queue depth, worker stats | -| `GET` | `/videos/` | Return generated MP4 | +| `GET` | `/health/stages` | Per-stage health (queries each worker) | +| `GET` | `/pipeline/status` | Active requests, queue depth, latencies | +| `GET` | `/videos/` | Serve generated video file | -### 3.2 Worker Pooling (`orchestrator/worker_manager.py`) +**Key components:** -`WorkerManager` maintains: -- idle queue (`acquire_worker()` blocks on backpressure), -- direct dispatch to worker ID (`client.direct(...)`), -- per-worker runtime stats (`status`, `completed`, `avg_latency_s`, `queue_depth`). +- `PipelineTracker` — Tracks active requests per stage, completed/failed + counts, and rolling average stage latencies. +- `dispatch_with_retry(mgr, request_id, json)` — Acquire worker → + dispatch → on failure retry on a different worker (up to + `STAGE_DISPATCH_RETRIES` attempts). +- `admission = asyncio.Semaphore(pipeline_depth)` — Limits concurrent + in-flight requests across the entire pipeline. -### 3.3 Stage Workers (`workers/*_worker.py`) +### 3.2 WorkerManager -Common pattern: -- launch backend scheduler subprocess (`launch_stage_server(...)`), -- run Dynamo RPC handlers (`generate`, `health`), -- forward requests through `StageClient.forward(...)`. +`orchestrator/worker_manager.py` — Per-stage worker pool with +busy/idle tracking. -Design direction: keep worker/orchestrator logic stable and swap only backend -handler implementation (`sglang`, `diffusers`, or others). +- Backed by `asyncio.Queue` (idle pool) — `acquire_worker()` awaits + the queue, `dispatch()` returns the worker to it on completion. +- `_call_direct()` sends to a specific worker via + `client.direct(json, worker_id)`. +- `status()` returns per-worker `{id, status, request_id, completed, + avg_latency_s}` plus stage-level `{queue_depth, completed, failed}`. -Stage IO contract: +### 3.3 Tensor Transfer — NIXL -| Stage | Input | Output | -|---|---|---| -| Encoder | prompt + guidance | `transfer_meta` for embeddings | -| Denoiser | embedding metadata + denoise params | `transfer_meta` for latents | -| VAE | latent metadata + request ID | `video_path` | +`workers/nixl_transfer.py` — GPU-direct RDMA transfer between stages. + +**Connection management — `PersistentConnector`:** + +Each `NixlTensorSender` and `NixlTensorReceiver` owns a +`PersistentConnector` instance (subclass of `nixl_connect.Connector`) +that overrides `_create_connection()` to reuse a single +`Connection` (= one `nixl_agent` / UCX endpoint) across all operations. + +``` +Per-process NIXL agents (e.g. denoiser with TP=2): + + Rank 0 subprocess: + ├── NixlReceiveStage → NixlTensorReceiver → PersistentConnector → agent A + └── NixlSendStage → NixlTensorSender → PersistentConnector → agent B + + Rank 1 subprocess: + └── (no NIXL agents — receives tensors via TP broadcast from rank 0) +``` + +- Agents are created eagerly at stage construction time, before any + request arrives. This ensures UCX endpoints are fully initialized. +- Each sender/receiver has its own `PersistentConnector` → its own + agent. This avoids UCX state conflicts between read and write roles. +- Agents are reused across all requests — never recreated. -Fallback path: if NIXL is unavailable, tensors are serialized over RPC (`tensor_data`). +**UCX transport configuration:** -### 3.4 Tensor Transport (`workers/nixl_transfer.py`) +For intra-node transfers, set `UCX_TLS=cuda_ipc,tcp,self,cuda_copy,cma` +to force NVLink (cuda_ipc) instead of IB RDMA. The default +`UCX_TLS=all` may select IB which fails across NUMA boundaries on +multi-socket systems. -`NixlTensorSender`: -- flattens tensors into one GPU buffer, -- registers a readable descriptor, -- returns compact metadata and tracks pending buffers. +**Sender (`NixlTensorSender`):** -`NixlTensorReceiver`: -- allocates destination GPU buffer, -- performs RDMA read, -- reconstructs tensors from metadata. +```python +sender = NixlTensorSender() # creates agent eagerly +meta, readable_op = sender.send({"latents": t}) # → (metadata, handle) +# caller holds readable_op until receiver completes RDMA pull +``` + +1. Flatten all tensors into a single contiguous GPU buffer. +2. Create NIXL `Descriptor` and register as readable. +3. Return `(metadata_dict, readable_op)` — caller must hold + `readable_op` to prevent GC of the descriptor and GPU buffer. -**Buffer strategy (current vs future):** -- Current: request-scoped temporary buffer allocation; simple and sufficient - because measured overhead share is small. -- Future: buffer pooling/manager for tighter latency control and memory reuse - under higher concurrency. +**Receiver (`NixlTensorReceiver`):** -## 4. SGLang-Specific Integration +```python +receiver = NixlTensorReceiver() # creates agent eagerly +tensors = receiver.recv(meta, device="cuda") # → {"latents": tensor} +``` -The back half should be backend-agnostic. The recommended design is a generic -`StageHandler` contract, with framework-specific implementations. +1. Allocate flat GPU buffer on target device. +2. `begin_read(rdma_meta, descriptor)` → one-sided RDMA pull. +3. `wait_for_completion()` — blocks until transfer finishes. +4. Slice flat buffer into individual tensors. + +**Buffer lifecycle:** + +``` +Sender (Encoder/Denoiser) Receiver (Denoiser/VAE) +───────────────────────── ─────────────────────── +create_readable(descriptor) + → register GPU buffer + → return (meta, readable_op) + begin_read(meta) + → RDMA pull from sender GPU + wait_for_completion() + → transfer done + ← completion notification +_poll_completed() + → status == COMPLETE + → drop readable_op → GC buffer +``` -### 4.1 Generic StageHandler Abstraction +### 3.4 Worker Interface + +All three workers follow an identical pattern: ```python -class StageHandler: - async def start(self) -> None: ... - async def generate(self, request: dict) -> dict: ... - async def health(self) -> dict: ... - async def shutdown(self) -> None: ... +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + # 1. Launch backend subprocess + processes, client, server_args = launch_stage_server(...) + + # 2. Define Dynamo RPC handlers + async def handle_generate(request, context): + output = await client.forward([build_req(...)]) + yield result_dict # JSON with NIXL metadata or video path + + # 3. Serve endpoints + gen_ep = runtime.endpoint("disagg_diffusion..generate") + await gen_ep.serve_endpoint(handle_generate) ``` -Worker responsibility stays unchanged: -- parse Dynamo request/response protocol, -- call `handler.generate(...)`, -- return stage output (`transfer_meta` or `video_path`). +**Stage IO contract:** + +| Stage | Input | Output | +|---|---|---| +| Encoder | prompt, guidance_scale | `transfer_meta` (embedding metadata) | +| Denoiser | `transfer_meta` + inference params | `transfer_meta` (latent metadata) | +| VAE | `transfer_meta` + request_id | `video_path` | + +Fallback: when NIXL is unavailable, tensors are serialized via +`torch.save()` + base64 in the `tensor_data` field. + + +## 4. Implementation: SGLang Backend + +### 4.1 PartialGPUWorker -Backend-specific logic moves into handlers: -- process launch and model init, -- stage graph / pipeline execution, -- tensor extraction/injection details. +`workers/partial_gpu_worker.py` — Extends SGLang's `GPUWorker` to load +only the modules each stage needs. -### 4.2 Current Handler: SGLang +- Overrides only `init_device_and_model()`. All other `GPUWorker` + behavior (forward execution, memory analysis, LoRA) is inherited. +- `build_partial_pipeline()` dynamically loads only the specified + modules (e.g. `["transformer", "scheduler"]` for denoiser). -Current implementation maps to an SGLang handler using -`workers/partial_gpu_worker.py`: -- `NixlReceiveStage` at denoiser/VAE entry, -- `NixlSendStage` at encoder/denoiser exit. +**NIXL pipeline stages:** -Stage builders: -- `build_encoder_stages()`: `TextEncodingStage -> NixlSendStage` -- `build_denoiser_stages()`: `NixlReceiveStage -> ... -> DenoisingStage -> NixlSendStage` -- `build_vae_stages()`: `NixlReceiveStage -> DecodingStage` +- `NixlReceiveStage(PipelineStage)` — Prepended at denoiser/VAE entry. + Rank 0: RDMA-pulls tensors via `NixlTensorReceiver`. + TP > 1: broadcasts pulled tensors to other ranks via + `torch.distributed.broadcast`. + Falls back to `_device_move()` for ZMQ path. -### 4.3 Alternative Handler: Diffusers (Planned) +- `NixlSendStage(PipelineStage)` — Appended at encoder/denoiser exit. + Rank 0 only: registers tensors as NIXL-readable, stores `readable_op` + in `_active_readables`, returns metadata in `OutputBatch`. + Rank > 0: returns empty `OutputBatch` (only rank 0's output is used). -A diffusers handler should implement the same contract and keep the same -orchestrator protocol: -- input/output request schema unchanged, -- same NIXL metadata fields for stage handoff, -- same health/reporting behavior. +**Stage builders:** -This allows swapping `sglang` -> `diffusers` without changing orchestrator -routing, worker manager, or external API. +| Function | Stages | +|---|---| +| `build_encoder_stages()` | `TextEncodingStage` → `NixlSendStage` | +| `build_denoiser_stages()` | `NixlReceiveStage` → `LatentPrep` → `TimestepPrep` → `DenoisingStage` → `NixlSendStage` | +| `build_vae_stages()` | `NixlReceiveStage` → `DecodingStage` | + +### 4.2 StageClient + +`workers/sglang_utils.py:StageClient` — Async ZMQ REQ/REP client. + +- `asyncio.Lock` serializes concurrent calls (ZMQ REQ requires strict + send/recv alternation). +- Configurable timeout: `STAGE_FORWARD_TIMEOUT_S` (default 120s). +- On timeout: `_reset_socket()` closes and reconnects the ZMQ socket + to recover from the broken REQ state (prevents cascading EFSM errors). + +### 4.3 HunyuanVideo Specifics -Decision note: keep orchestrator as a standalone component for now; evaluate -merging into a router layer only after smart-routing requirements justify it. +- **Dual encoder:** auto-detects `text_encoder` + `text_encoder_2` + (Llama + CLIP) from `model_index.json`. +- **HunyuanConfig patching:** wraps `HunyuanConfig.__init__` to supply + `task_type=T2V` when omitted. +- **Triton norm workaround:** wraps `norm_infer` to call + `.contiguous()` on non-contiguous tensors from attention reshapes. +- **Component config sync:** reads all component `config.json` files + to ensure correct parameters even for non-loaded components. -## 5. Roadmap (Condensed) -- [x] Multi-worker per stage with etcd discovery -- [x] Pipeline overlap with global admission control -- [x] NIXL metadata-only RPC + GPU-direct tensor transfer -- [ ] Runtime auto-scaling from `/pipeline/status` -- [ ] Smarter routing (affinity/priority/load-aware) -- [ ] Fault tolerance hardening and graceful degradation +## 5. Deployment + +### 5.1 Example: 8-GPU HunyuanVideo + +```bash +# 1 encoder (GPU 0) + 3 denoisers TP=2 (GPU 1,2 / 3,4 / 5,6) + 1 VAE (GPU 7) +UCX_TLS=cuda_ipc,tcp,self,cuda_copy,cma \ +GPU_ENC=0 GPU_DEN="1,2;3,4;5,6" GPU_VAE=7 PORT=8091 \ + ./run_all.sh +``` + +### 5.2 Environment Variables + +| Variable | Default | Description | +|---|---|---| +| `GPU_ENC` | `0` | GPU(s) for encoder workers (`;`-separated for multi-worker) | +| `GPU_DEN` | `1,2` | GPU(s) for denoiser workers (`,` = TP within worker, `;` = multiple workers) | +| `GPU_VAE` | `3` | GPU(s) for VAE workers | +| `PORT` | `8080` | HTTP port | +| `UCX_TLS` | `all` | UCX transport list — set to `cuda_ipc,tcp,self,cuda_copy,cma` for intra-node | +| `MAX_PIPELINE_DEPTH` | `4` | Max concurrent requests in pipeline | +| `STAGE_FORWARD_TIMEOUT_S` | `120` | ZMQ timeout per stage request (seconds) | +| `WORKER_SETTLE_S` | `30` | Seconds to wait for slow workers after first discovery | +| `STAGE_DISPATCH_RETRIES` | `2` | Retry count on stage dispatch failure | +| `DISABLE_NIXL` | `false` | Force ZMQ fallback (disable NIXL) | + +### 5.3 Measured Performance (HunyuanVideo, 9 frames, 3 steps) + +| Stage | Workers | Avg Latency | +|---|---|---| +| Encoder | 1 × GPU | 0.33s | +| Denoiser | 3 × TP=2 | 3.7s | +| VAE | 1 × GPU | 3.8s | +| **End-to-end** | | **~7.9s** | +| **Throughput (20 concurrent)** | | **~0.25 req/s** (VAE-bound) | + + +## 6. Roadmap + +- [x] Multi-worker scaling — N workers per stage, etcd auto-discovery +- [x] Pipeline parallelism — overlapping requests across stages +- [x] NIXL GPU-direct transfer — GPU-to-GPU RDMA, only metadata over RPC +- [x] TP support — TP broadcast in NixlReceiveStage, rank-0-only send +- [x] PersistentConnector — stable NIXL connections, no per-request agent churn +- [ ] Runtime scaling — external auto-scaler based on `/pipeline/status` +- [ ] Smart routing — load-aware dispatch, affinity, request priority +- [ ] Fault tolerance — health-check + eviction, graceful degradation - [ ] Streaming decode output + richer observability -- [ ] Evaluate merging orchestrator into router layer -- [ ] Explore diffusion-based smart router for stage/worker selection -- [ ] Formalize `StageHandler` interface and migrate SGLang to handler plugin -- [ ] Add Diffusers handler with parity tests against SGLang outputs +- [ ] Multi-model support — omni-modal pipelines, heterogeneous stage graphs + + +## 7. File Map + +``` +examples/disagg_diffusion/ +├── DESIGN.md ← this document +├── run_all.sh ← launch etcd + all workers + orchestrator +├── stress_test.sh ← concurrent load testing script +│ +├── orchestrator/ +│ ├── run_disagg.py ← HTTP server, pipeline dispatch, admission control +│ ├── worker_manager.py ← WorkerManager: idle-pool, acquire/dispatch/status +│ └── request_validation.py ← Pydantic-based HTTP request validation +│ +└── workers/ + ├── protocol.py ← Pydantic request/response models (Dynamo RPC) + ├── encoder_worker.py ← Encoder Dynamo worker + ├── denoiser_worker.py ← Denoiser Dynamo worker (supports TP) + ├── vae_worker.py ← VAE Dynamo worker + ├── nixl_transfer.py ← NixlTensorSender / NixlTensorReceiver (RDMA) + ├── partial_gpu_worker.py ← PartialGPUWorker, NixlSend/ReceiveStage, launcher + └── sglang_utils.py ← StageClient, build_partial_pipeline, launch_stage_server +``` From 1e394647c9f72f57473a3c003bea81a67e56aee1 Mon Sep 17 00:00:00 2001 From: Hongli Mi Date: Fri, 20 Mar 2026 20:16:32 +0800 Subject: [PATCH 22/22] =?UTF-8?q?perf:=20optimize=20NIXL=20transfer=20path?= =?UTF-8?q?=20=E2=80=94=20NixlSendStage=20640ms=20=E2=86=92=203ms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove GPU tensor debug stats logging (.float() copy was 600ms+ per request even at DEBUG level due to GPU synchronization) - Skip torch.cat when only one tensor (avoid unnecessary GPU copy) - Reuse event loop in _run_coro instead of creating new one each call - Add NIXL send timing breakdown log (flatten/descriptor/create/meta) NixlSendStage: 640ms → 3ms (213x faster) N=10 stress: 53s → 41s (23% faster) Co-Authored-By: Claude Opus 4.6 --- .../disagg_diffusion/workers/nixl_transfer.py | 40 ++++++++++++++----- .../workers/partial_gpu_worker.py | 18 --------- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/examples/disagg_diffusion/workers/nixl_transfer.py b/examples/disagg_diffusion/workers/nixl_transfer.py index 64fe82f50679..ce284363273b 100644 --- a/examples/disagg_diffusion/workers/nixl_transfer.py +++ b/examples/disagg_diffusion/workers/nixl_transfer.py @@ -81,18 +81,14 @@ async def _create_connection(self) -> nixl_connect.Connection: # Helpers # --------------------------------------------------------------------------- +_event_loop = None + def _run_coro(coro): """Run a coroutine from synchronous context (sglang scheduler thread).""" - try: - loop = asyncio.get_event_loop() - if loop.is_running(): - # Shouldn't happen in sglang's scheduler, but be safe - import concurrent.futures - with concurrent.futures.ThreadPoolExecutor() as pool: - return pool.submit(asyncio.run, coro).result(timeout=30) - return loop.run_until_complete(coro) - except RuntimeError: - return asyncio.run(coro) + global _event_loop + if _event_loop is None or _event_loop.is_closed(): + _event_loop = asyncio.new_event_loop() + return _event_loop.run_until_complete(coro) # --------------------------------------------------------------------------- @@ -123,11 +119,33 @@ def send(self, tensors: Dict[str, torch.Tensor]) -> Tuple[dict, object]: async def _async_send( self, tensors: Dict[str, torch.Tensor] ) -> Tuple[dict, object]: + import time as _time + t0 = _time.monotonic() + # Flatten all tensors into a single contiguous buffer - flat = torch.cat([t.contiguous().view(-1) for t in tensors.values()]) + vals = list(tensors.values()) + if len(vals) == 1: + flat = vals[0].contiguous().view(-1) + else: + flat = torch.cat([t.contiguous().view(-1) for t in vals]) + t1 = _time.monotonic() + descriptor = nixl_connect.Descriptor(flat) + t2 = _time.monotonic() + readable = await self.connector.create_readable(descriptor) + t3 = _time.monotonic() + raw_meta = readable.metadata() + t4 = _time.monotonic() + + logger.info( + "NIXL send breakdown: flatten=%.1fms descriptor=%.1fms " + "create_readable=%.1fms metadata=%.1fms total=%.1fms " + "(buf=%s %.2fMB)", + (t1-t0)*1000, (t2-t1)*1000, (t3-t2)*1000, (t4-t3)*1000, + (t4-t0)*1000, list(flat.shape), flat.nbytes/1e6, + ) meta = { "tensor_keys": list(tensors.keys()), diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py index f75ea406243c..0966b6499fc0 100644 --- a/examples/disagg_diffusion/workers/partial_gpu_worker.py +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -165,18 +165,6 @@ def _nixl_pull(self, batch: Req, meta: dict) -> Req: reconstructed[k] = v for base, idx_map in indexed.items(): reconstructed[base] = [idx_map[i] for i in sorted(idx_map)] - # Debug: log tensor stats after NIXL pull - for k, v in reconstructed.items(): - if isinstance(v, list): - for i, t in enumerate(v): - if hasattr(t, 'float'): - f = t.float() - logger.debug("NIXL_RECV %s[%d]: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", - k, i, t.shape, t.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) - elif hasattr(v, 'float'): - f = v.float() - logger.debug("NIXL_RECV %s: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", - k, v.shape, v.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) from sglang_utils import inject_tensors_to_req inject_tensors_to_req(batch, reconstructed) return batch @@ -289,12 +277,6 @@ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: if not tensors: return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) - # Debug: log tensor stats before NIXL send - for k, t in tensors.items(): - f = t.float() - logger.debug("NIXL_SEND %s: shape=%s dtype=%s mean=%.6f std=%.6f min=%.6f max=%.6f", - k, t.shape, t.dtype, f.mean().item(), f.std().item(), f.min().item(), f.max().item()) - from nixl_transfer import NIXL_AVAILABLE if NIXL_AVAILABLE and self._sender is not None: return self._nixl_send(tensors)