chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/.claude/CLAUDE.md

@@ -0,0 +1,112 @@
+# CLAUDE.md — sglang-diffusion (multimodal_gen)
+
+## What is this?
+
+SGLang's diffusion/multimodal generation subsystem. Separate from the LLM runtime (`srt`). Supports 20+ image/video diffusion models (Wan, FLUX, HunyuanVideo, LTX, Qwen-Image, etc.) with distributed inference, LoRA, and multiple attention backends.
+
+## Quick Start
+
+```bash
+# One-shot generation
+sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --prompt "A curious raccoon" --save-output
+
+# Start server
+sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --num-gpus 4
+
+# Python API
+from sglang import DiffGenerator
+gen = DiffGenerator.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+result = gen.generate(sampling_params_kwargs={"prompt": "A curious raccoon"})
+```
+
+## Architecture
+
+```
+CLI / Python API / HTTP Server (FastAPI + OpenAI-compatible)
+    ↓ ZMQ
+Scheduler (request queue, batching, dispatch)
+    ↓ multiprocessing pipes
+GPU Worker(s) → ComposedPipeline (stages: TextEncode → Denoise → Decode)
+```
+
+### Key Directories
+
+```
+runtime/
+├── entrypoints/        # CLI (generate/serve), HTTP server, Python API (DiffGenerator)
+├── managers/           # scheduler.py, gpu_worker.py
+├── pipelines_core/     # ComposedPipelineBase, stages/, schedule_batch.py (Req/OutputBatch)
+├── pipelines/          # Model-specific pipelines (wan, flux, hunyuan, ltx, qwen_image, ...)
+├── models/             # encoders/, dits/, vaes/, schedulers/
+├── layers/             # attention/, lora/, quantization/
+├── loader/             # Model loading, weight utils
+├── server_args.py      # ServerArgs (all CLI/config params)
+└── distributed/        # Multi-GPU (TP, SP: ulysses/ring)
+configs/
+├── pipeline_configs/    # Per-model pipeline configs
+├── sample/             # SamplingParams
+└── models/             # DiT, VAE, Encoder configs
+```
+
+### Key Classes
+
+| Class | Location | Purpose |
+|-------|----------|---------|
+| `DiffGenerator` | `runtime/entrypoints/diffusion_generator.py` | Python API entry point |
+| `ComposedPipelineBase` | `runtime/pipelines_core/composed_pipeline_base.py` | Pipeline orchestrator (stages) |
+| `Scheduler` | `runtime/managers/scheduler.py` | ZMQ event loop, request dispatch |
+| `GPUWorker` | `runtime/managers/gpu_worker.py` | GPU inference worker |
+| `Req` / `OutputBatch` | `runtime/pipelines_core/schedule_batch.py` | Request/output containers |
+| `ServerArgs` | `runtime/server_args.py` | All config params |
+| `SamplingParams` | `configs/sample/sampling_params.py` | Generation params |
+| `PipelineConfig` | `configs/pipeline_configs/base.py` | Model structure config |
+
+### Key Functions
+
+| Function | Module | Purpose |
+|----------|--------|---------|
+| `build_pipeline()` | `runtime/pipelines_core/__init__.py` | Instantiate pipeline from model_path |
+| `get_model_info()` | `registry.py` | Resolve pipeline + config classes |
+| `launch_server()` | `runtime/launch_server.py` | Start multi-process server |
+
+## Adding a New Model
+
+1. Create pipeline in `runtime/pipelines/` extending `ComposedPipelineBase`
+2. Define stages via `create_pipeline_stages()` (TextEncoding → Denoising → Decoding)
+3. Add config in `configs/pipeline_configs/`
+4. Register in `registry.py` via `register_configs()`
+
+## Multi-GPU
+
+```bash
+# Sequence parallelism (video frames across GPUs)
+sglang serve --model-path ... --num-gpus 4 --ulysses-degree 2 --ring-degree 2
+
+# Tensor parallelism (model layers across GPUs)
+sglang serve --model-path ... --num-gpus 2 --tp-size 2
+```
+
+## Testing
+
+```bash
+# Tests live in test/ subdirectory
+python -m pytest python/sglang/multimodal_gen/test/
+
+# No need to pre-download models — auto-downloaded at runtime
+# Dependencies assumed already installed via `pip install -e "python[diffusion]"`
+```
+
+## Performance Tuning
+
+For questions about optimal performance, fastest commands, VRAM reduction, or best flag combinations for a given model/GPU setup, **read the [sglang-diffusion-performance skill](skills/sglang-diffusion-performance/SKILL.md)**. It contains a complete table of lossless and lossy optimization flags with trade-offs, quick recipes, and tuning tips.
+
+### Perf Measurement
+
+Look for `Pixel data generated successfully in xxxx seconds` in console output. With warmup enabled, use the line containing `warmup excluded` for accurate timing.
+
+## Env Vars
+
+Defined in `envs.py` (300+ vars). Key ones:
+- `SGLANG_DIFFUSION_ATTENTION_BACKEND` — attention backend override
+- `SGLANG_CACHE_DIT_ENABLED` — enable Cache-DiT acceleration
+- `SGLANG_CLOUD_STORAGE_TYPE` — cloud output storage (s3, etc.)

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/SKILL.md

@@ -0,0 +1,626 @@
+---
+name: sglang-diffusion-add-model
+description: Use when adding a new diffusion model or Diffusers pipeline to SGLang.
+---
+
+# Add a Diffusion Model to SGLang
+
+Use this skill when adding a new diffusion model or pipeline variant to `sglang.multimodal_gen`.
+
+## Two Pipeline Styles
+
+### Style A: Hybrid Monolithic Pipeline (Recommended)
+
+The recommended default for most new models. Uses a three-stage structure:
+
+```
+BeforeDenoisingStage (model-specific)  -->  DenoisingStage (standard)  -->  DecodingStage (standard)
+```
+
+- **BeforeDenoisingStage**: A single, model-specific stage that consolidates all pre-processing logic: input validation, text encoding, image encoding, latent preparation, timestep setup. This stage is unique per model.
+- **DenoisingStage**: Framework-standard stage for the denoising loop (DiT/UNet forward passes). Shared across models.
+- **DecodingStage**: Framework-standard stage for VAE decoding. Shared across models.
+
+**Why recommended?** Modern diffusion models have highly heterogeneous pre-processing requirements (different text encoders, different latent formats, different conditioning mechanisms). The Hybrid approach keeps pre-processing isolated per model, avoids fragile shared stages with excessive conditional logic, and lets developers port Diffusers reference code quickly.
+
+### Style B: Modular Composition Style
+
+Uses the framework's fine-grained standard stages (`TextEncodingStage`, `LatentPreparationStage`, `TimestepPreparationStage`, etc.) to build the pipeline by composition.
+
+This style is appropriate when:
+- **The new model's pre-processing can largely reuse existing stages** — e.g., a model that uses standard CLIP/T5 text encoding + standard latent preparation with minimal customization. In this case, `add_standard_t2i_stages()` or `add_standard_ti2i_stages()` may be all you need.
+- **A model-specific optimization needs to be extracted as a standalone stage** — e.g., a specialized encoding or conditioning step that benefits from being a separate stage for profiling, parallelism control, or reuse across multiple pipeline variants.
+
+See existing Modular examples: `QwenImagePipeline` (uses `add_standard_t2i_stages`), `FluxPipeline`, `WanPipeline`, `SanaPipeline`, `StableDiffusion3Pipeline`, and `ZImagePipeline`.
+
+### How to Choose
+
+| Situation | Recommended Style |
+|-----------|-------------------|
+| Model has unique/complex pre-processing (VLM captioning, AR token generation, custom latent packing, etc.) | **Hybrid** — consolidate into a BeforeDenoisingStage |
+| Model fits neatly into standard text-to-image or text+image-to-image pattern | **Modular** — use `add_standard_t2i_stages()` / `add_standard_ti2i_stages()` |
+| Porting a Diffusers pipeline with many custom steps | **Hybrid** — copy the `__call__` logic into a single stage |
+| Adding a variant of an existing model that shares most logic | **Modular** — reuse existing stages, customize via PipelineConfig callbacks |
+| A specific pre-processing step needs special parallelism or profiling isolation | **Modular** — extract that step as a dedicated stage |
+
+**Key principle (both styles)**: The stage(s) before `DenoisingStage` must produce a `Req` batch object with all the standard tensor fields that `DenoisingStage` expects (latents, timesteps, prompt_embeds, etc.). As long as this contract is met, the pipeline remains composable regardless of which style you use.
+
+---
+
+## Key Files and Directories
+
+| Purpose | Path |
+|---------|------|
+| Pipeline classes | `python/sglang/multimodal_gen/runtime/pipelines/` |
+| Model-specific stages | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/` |
+| PipelineStage base class | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py` |
+| Pipeline base class | `python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py` |
+| Standard stages (Denoising, Decoding) | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/` |
+| Pipeline configs | `python/sglang/multimodal_gen/configs/pipeline_configs/` |
+| Sampling params | `python/sglang/multimodal_gen/configs/sample/` |
+| DiT model implementations | `python/sglang/multimodal_gen/runtime/models/dits/` |
+| VAE implementations | `python/sglang/multimodal_gen/runtime/models/vaes/` |
+| Encoder implementations | `python/sglang/multimodal_gen/runtime/models/encoders/` |
+| Scheduler implementations | `python/sglang/multimodal_gen/runtime/models/schedulers/` |
+| Model/VAE/DiT configs | `python/sglang/multimodal_gen/configs/models/dits/`, `vaes/`, `encoders/` |
+| Central registry | `python/sglang/multimodal_gen/registry.py` |
+| Model component registry | `python/sglang/multimodal_gen/runtime/models/registry.py` |
+| Current support list | `docs/diffusion/compatibility_matrix.md` |
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Obtain and Study the Reference Implementation
+
+**Before writing any code, obtain the model's reference implementation or Diffusers pipeline code.** You need the actual source code to work from — do not guess or assume the model's architecture. If the user already gave a HuggingFace model ID or repo, inspect that yourself first. Ask the user only when the reference implementation is private, ambiguous, or otherwise unavailable. Typical sources are:
+- The model's Diffusers pipeline source (e.g., the `pipeline_*.py` file from the `diffusers` library or HuggingFace repo)
+- Or the model's official reference implementation (e.g., from the model author's GitHub repo)
+- Or the HuggingFace model ID so you can look up `model_index.json` and the associated pipeline class
+
+Once you have the reference code, study it thoroughly:
+
+1. Find the model's `model_index.json` to identify required modules (text_encoder, vae, transformer, scheduler, etc.)
+2. Read the Diffusers pipeline's `__call__` method end-to-end. Identify:
+   - How text prompts are encoded
+   - How latents are prepared (shape, dtype, scaling)
+   - How timesteps/sigmas are computed
+   - What conditioning kwargs the DiT/UNet expects
+   - How the denoising loop works (classifier-free guidance, etc.)
+   - How VAE decoding is done (scaling factors, tiling, etc.)
+
+### Step 2: Evaluate Reuse of Existing Pipelines and Stages
+
+**Before creating any new files, check whether an existing pipeline or stage can be reused or extended.** Only create new pipelines/stages when the existing ones would require extensive modifications or when no similar implementation exists.
+
+Specifically:
+1. **Compare the new model's architecture against existing pipelines** before creating files. Current native families include LTX-2/2.3, HunyuanVideo/FastHunyuan, Wan/FastWan/TurboWan/LingBot World, MOVA, FLUX/FLUX.2/Klein, Z-Image, Qwen-Image/edit/layered, GLM-Image, SD3, Hunyuan3D, Helios, Cosmos3, SANA/SANA-WM, FireRed, ERNIE-Image, JoyAI, and Ideogram4. If the new model shares most of its structure with an existing one (e.g., same text encoders, similar latent format, compatible denoising loop), prefer:
+   - Adding a new config variant to the existing pipeline rather than creating a new pipeline class
+   - Reusing the existing `BeforeDenoisingStage` with minor parameter differences
+   - Using `add_standard_t2i_stages()` / `add_standard_ti2i_stages()` / `add_standard_ti2v_stages()` if the model fits standard patterns
+2. **Check existing stages** in `runtime/pipelines_core/stages/` and `stages/model_specific_stages/`. If an existing stage handles 80%+ of what the new model needs, extend it rather than duplicating it.
+3. **Check existing model components** — many models share VAEs (e.g., `AutoencoderKL`), text encoders (CLIP, T5), and schedulers. Reuse these directly instead of re-implementing.
+
+**Rule of thumb**: Only create a new file when the existing implementation would need substantial structural changes to accommodate the new model, or when no architecturally similar implementation exists.
+
+### Step 3: Implement Model Components
+
+Adapt or implement the model's core components in the appropriate directories.
+
+**DiT/Transformer** (`runtime/models/dits/{model_name}.py`):
+
+```python
+# python/sglang/multimodal_gen/runtime/models/dits/my_model.py
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    LayerNormScaleShift,
+    RMSNormScaleShift,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    get_attn_backend,
+)
+
+
+class MyModelTransformer2DModel(nn.Module):
+    """DiT model for MyModel.
+
+    Adapt from the Diffusers/reference implementation. Key points:
+    - Use SGLang's fused LayerNorm/RMSNorm ops (see `existing-fast-paths.md` under the benchmark/profile skill)
+    - Use SGLang's attention backend selector
+    - Keep the same parameter naming as Diffusers for weight loading compatibility
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # ... model layers ...
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        # ... model-specific kwargs ...
+    ) -> torch.Tensor:
+        # ... forward pass ...
+        return output
+```
+
+**Tensor Parallel (TP) and Sequence Parallel (SP)**: For multi-GPU deployment, it is recommended to add TP/SP support to the DiT model. This can be done incrementally after the single-GPU implementation is verified. Reference existing implementations and adapt to your model's architecture:
+
+- **Wan model** (`runtime/models/dits/wanvideo.py`) — Full TP + SP reference:
+  - TP: Uses `ColumnParallelLinear` for Q/K/V projections, `RowParallelLinear` for output projections, attention heads divided by `tp_size`
+  - SP: Sequence dimension sharding via `get_sp_world_size()`, padding for alignment, `sequence_model_parallel_all_gather` for aggregation
+  - Cross-attention skips SP (`skip_sequence_parallel=is_cross_attention`)
+- **Qwen-Image model** (`runtime/models/dits/qwen_image.py`) — SP + USPAttention reference:
+  - SP: Uses `USPAttention` (Ulysses + Ring Attention), configured via `--ulysses-degree` / `--ring-degree`
+  - TP: Uses `MergedColumnParallelLinear` for QKV (with Nunchaku quantization), `ReplicatedLinear` otherwise
+
+**Important**: These are references only — each model has its own architecture and parallelism requirements. Consider:
+- How attention heads can be divided across TP ranks
+- Whether the model's sequence dimension is naturally shardable for SP
+- Which linear layers benefit from column/row parallel sharding vs. replication
+- Whether cross-attention or other special modules need SP exclusion
+
+Key imports for distributed support:
+```python
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_sp_group,
+    get_sp_world_size,
+    get_tp_world_size,
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    ReplicatedLinear,
+)
+```
+
+**VAE** (`runtime/models/vaes/{model_name}.py`): Implement if the model uses a non-standard VAE. Many models reuse existing VAEs.
+
+**Encoders** (`runtime/models/encoders/{model_name}.py`): Implement if the model uses custom text/image encoders.
+
+**Schedulers** (`runtime/models/schedulers/{scheduler_name}.py`): Implement if the model requires a custom scheduler not available in Diffusers.
+
+### Step 4: Create Model Configs
+
+**DiT Config** (`configs/models/dits/{model_name}.py`):
+
+```python
+# python/sglang/multimodal_gen/configs/models/dits/mymodel.py
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTConfig
+
+
+@dataclass
+class MyModelDitConfig(DiTConfig):
+    arch_config: dict = field(default_factory=lambda: {
+        "in_channels": 16,
+        "num_layers": 24,
+        "patch_size": 2,
+        # ... model-specific architecture params ...
+    })
+```
+
+**VAE Config** (`configs/models/vaes/{model_name}.py`):
+
+```python
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
+
+
+@dataclass
+class MyModelVAEConfig(VAEConfig):
+    vae_scale_factor: int = 8
+    # ... VAE-specific params ...
+```
+
+**Sampling Params** (`configs/sample/{model_name}.py`):
+
+```python
+from dataclasses import dataclass
+
+from sglang.multimodal_gen.configs.sample.base import SamplingParams
+
+
+@dataclass
+class MyModelSamplingParams(SamplingParams):
+    num_inference_steps: int = 50
+    guidance_scale: float = 7.5
+    height: int = 1024
+    width: int = 1024
+    # ... model-specific defaults ...
+```
+
+### Step 5: Create PipelineConfig
+
+The `PipelineConfig` holds static model configuration and defines callback methods used by the standard `DenoisingStage` and `DecodingStage`.
+
+```python
+# python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py
+
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ImagePipelineConfig,
+    ModelTaskType,
+    # PipelineConfig,              # common base for many video pipelines
+    # SpatialImagePipelineConfig,  # alternative base for spatial image models
+)
+from sglang.multimodal_gen.configs.models.dits.mymodel import MyModelDitConfig
+from sglang.multimodal_gen.configs.models.vaes.mymodel import MyModelVAEConfig
+
+
+@dataclass
+class MyModelPipelineConfig(ImagePipelineConfig):
+    """Pipeline config for MyModel.
+
+    This config provides callbacks that the standard DenoisingStage and
+    DecodingStage use during execution. The BeforeDenoisingStage handles
+    all model-specific pre-processing independently.
+    """
+
+    task_type: ModelTaskType = ModelTaskType.T2I
+    vae_precision: str = "bf16"
+    should_use_guidance: bool = True
+    vae_tiling: bool = False
+    enable_autocast: bool = False
+
+    dit_config: DiTConfig = field(default_factory=MyModelDitConfig)
+    vae_config: VAEConfig = field(default_factory=MyModelVAEConfig)
+
+    # --- Callbacks used by DenoisingStage ---
+
+    def get_freqs_cis(self, batch, device, rotary_emb, dtype):
+        """Prepare rotary position embeddings for the DiT."""
+        # Model-specific RoPE computation
+        ...
+        return freqs_cis
+
+    def prepare_pos_cond_kwargs(self, batch, latent_model_input, t, **kwargs):
+        """Build positive conditioning kwargs for each denoising step."""
+        return {
+            "hidden_states": latent_model_input,
+            "encoder_hidden_states": batch.prompt_embeds[0],
+            "timestep": t,
+            # ... model-specific kwargs ...
+        }
+
+    def prepare_neg_cond_kwargs(self, batch, latent_model_input, t, **kwargs):
+        """Build negative conditioning kwargs for CFG."""
+        return {
+            "hidden_states": latent_model_input,
+            "encoder_hidden_states": batch.negative_prompt_embeds[0],
+            "timestep": t,
+            # ... model-specific kwargs ...
+        }
+
+    # --- Callbacks used by DecodingStage ---
+
+    def get_decode_scale_and_shift(self):
+        """Return (scale, shift) for latent denormalization before VAE decode."""
+        return self.vae_config.latents_std, self.vae_config.latents_mean
+
+    def post_denoising_loop(self, latents, batch):
+        """Optional post-processing after the denoising loop finishes."""
+        return latents.to(torch.bfloat16)
+
+    def post_decoding(self, frames, server_args):
+        """Optional post-processing after VAE decoding."""
+        return frames
+```
+
+There is no separate `VideoPipelineConfig` base class. For video models, choose
+`ModelTaskType.T2V`, `ModelTaskType.I2V`, or `ModelTaskType.TI2V`, and follow
+existing video configs such as Wan, LTX, Hunyuan, Helios, or MOVA when deciding
+whether to subclass `PipelineConfig` directly or use a model-specific base.
+
+**Important**: The `prepare_pos_cond_kwargs` / `prepare_neg_cond_kwargs` methods define what the DiT receives at each denoising step. These must match the DiT's `forward()` signature.
+
+### Step 6: Implement the BeforeDenoisingStage (Core Step)
+
+This is the heart of the Hybrid pattern. Create a single stage that handles ALL pre-processing.
+
+```python
+# python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/my_model.py
+
+import torch
+from typing import List, Optional, Union
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class MyModelBeforeDenoisingStage(PipelineStage):
+    """Monolithic pre-processing stage for MyModel.
+
+    Consolidates all logic before the denoising loop:
+    - Input validation
+    - Text/image encoding
+    - Latent preparation
+    - Timestep/sigma computation
+
+    This stage produces a Req batch with all fields required by
+    the standard DenoisingStage.
+    """
+
+    def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler):
+        super().__init__()
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.transformer = transformer
+        self.scheduler = scheduler
+        # ... other initialization (image processors, scale factors, etc.) ...
+
+    # --- Internal helper methods ---
+    # Copy/adapt directly from the Diffusers reference pipeline.
+    # These are private to this stage; no need to make them reusable.
+
+    def _encode_prompt(self, prompt, device, dtype):
+        """Encode text prompt into embeddings."""
+        # ... model-specific text encoding logic ...
+        return prompt_embeds, negative_prompt_embeds
+
+    def _prepare_latents(self, batch_size, height, width, dtype, device, generator):
+        """Create initial noisy latents."""
+        # ... model-specific latent preparation ...
+        return latents
+
+    def _prepare_timesteps(self, num_inference_steps, device):
+        """Compute the timestep/sigma schedule."""
+        # ... model-specific timestep computation ...
+        return timesteps, sigmas
+
+    # --- Main forward method ---
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Execute all pre-processing and populate batch for DenoisingStage.
+
+        This method mirrors the first half of a Diffusers pipeline __call__,
+        up to (but not including) the denoising loop.
+        """
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        generator = torch.Generator(device=device).manual_seed(batch.seed)
+
+        # 1. Encode prompt
+        prompt_embeds, negative_prompt_embeds = self._encode_prompt(
+            batch.prompt, device, dtype
+        )
+
+        # 2. Prepare latents
+        latents = self._prepare_latents(
+            batch_size=1,
+            height=batch.height,
+            width=batch.width,
+            dtype=dtype,
+            device=device,
+            generator=generator,
+        )
+
+        # 3. Prepare timesteps
+        timesteps, sigmas = self._prepare_timesteps(
+            batch.num_inference_steps, device
+        )
+
+        # 4. Populate batch with everything DenoisingStage needs
+        batch.prompt_embeds = [prompt_embeds]
+        batch.negative_prompt_embeds = [negative_prompt_embeds]
+        batch.latents = latents
+        batch.timesteps = timesteps
+        batch.num_inference_steps = len(timesteps)
+        batch.sigmas = sigmas
+        batch.generator = generator
+        batch.raw_latent_shape = latents.shape
+        batch.height = batch.height
+        batch.width = batch.width
+
+        return batch
+```
+
+**Key fields that `DenoisingStage` expects on the batch** (set these in your `forward`):
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `batch.latents` | `torch.Tensor` | Initial noisy latent tensor |
+| `batch.timesteps` | `torch.Tensor` | Timestep schedule |
+| `batch.num_inference_steps` | `int` | Number of denoising steps |
+| `batch.sigmas` | `list[float]` | Sigma schedule (as a list, not numpy) |
+| `batch.prompt_embeds` | `list[torch.Tensor]` | Positive prompt embeddings (wrapped in list) |
+| `batch.negative_prompt_embeds` | `list[torch.Tensor]` | Negative prompt embeddings (wrapped in list) |
+| `batch.generator` | `torch.Generator` | RNG generator for reproducibility |
+| `batch.raw_latent_shape` | `tuple` | Original latent shape before any packing |
+| `batch.height` / `batch.width` | `int` | Output dimensions |
+
+### Step 7: Define the Pipeline Class
+
+The pipeline class is minimal -- it just wires the stages together.
+
+```python
+# python/sglang/multimodal_gen/runtime/pipelines/my_model.py
+
+from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.my_model import (
+    MyModelBeforeDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
+    pipeline_name = "MyModelPipeline"  # Must match model_index.json _class_name
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+        # ... list all modules from model_index.json ...
+    ]
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        # 1. Monolithic pre-processing (model-specific)
+        self.add_stage(
+            MyModelBeforeDenoisingStage(
+                vae=self.get_module("vae"),
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        # 2. Standard denoising loop (framework-provided)
+        self.add_stage(
+            DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        # 3. Standard VAE decoding (framework-provided)
+        self.add_standard_decoding_stage()
+
+
+# REQUIRED: This is how the registry discovers the pipeline
+EntryClass = [MyModelPipeline]
+```
+
+### Step 8: Register the Model
+
+In `python/sglang/multimodal_gen/registry.py`, register your configs:
+
+```python
+register_configs(
+    sampling_param_cls=MyModelSamplingParams,
+    pipeline_config_cls=MyModelPipelineConfig,
+    hf_model_paths=[
+        "org/my-model-name",  # HuggingFace model ID(s)
+    ],
+    model_detectors=[
+        lambda path: "my-model" in path.lower(),
+    ],
+)
+```
+
+`register_configs()` does not take a `model_family` argument. It registers the
+sampling and pipeline config classes, then resolves models by exact
+`hf_model_paths` or optional detector predicates. Prefer exact `hf_model_paths`
+for public checkpoints used in docs or tests; use detector predicates only for
+families where local mirrors, renamed repos, or generated paths are common.
+
+The `EntryClass` in your pipeline file is automatically discovered by the registry's `_discover_and_register_pipelines()` function -- no additional registration needed for the pipeline class itself.
+
+### Step 9: Verify Output Quality
+
+After implementation, **you must verify that the generated output is not noise**. A noisy or garbled output image/video is the most common sign of an incorrect implementation. Common causes include:
+
+- Incorrect latent scale/shift factors (`get_decode_scale_and_shift` returning wrong values)
+- Wrong timestep/sigma schedule (order, dtype, or value range)
+- Mismatched conditioning kwargs (fields not matching the DiT's `forward()` signature)
+- Incorrect VAE decoder configuration (wrong `vae_scale_factor`, missing denormalization)
+- Rotary embedding style mismatch (`is_neox_style` set incorrectly)
+- Wrong prompt embedding format (missing list wrapping, wrong encoder output selection)
+
+**If the output is noise, the implementation is incorrect — do not ship it.** Debug by:
+1. Comparing intermediate tensor values (latents, prompt_embeds, timesteps) against the Diffusers reference pipeline
+2. Running the Diffusers pipeline and SGLang pipeline side-by-side with the same seed
+3. Checking each stage's output shape and value range independently
+
+## Reference Implementations
+
+### Hybrid Style (recommended for most new models)
+
+| Model | Pipeline | BeforeDenoisingStage | PipelineConfig |
+|-------|----------|---------------------|----------------|
+| GLM-Image | `runtime/pipelines/glm_image.py` | `stages/model_specific_stages/glm_image.py` | `configs/pipeline_configs/glm_image.py` |
+| Qwen-Image-Layered | `runtime/pipelines/qwen_image.py` (`QwenImageLayeredPipeline`) | `stages/model_specific_stages/qwen_image_layered.py` | `configs/pipeline_configs/qwen_image.py` (`QwenImageLayeredPipelineConfig`) |
+| Cosmos3 | `runtime/pipelines/cosmos3_pipeline.py` | `stages/model_specific_stages/cosmos3.py` | `configs/pipeline_configs/cosmos3.py` |
+| ErnieImage | `runtime/pipelines/ernie_image.py` | `stages/model_specific_stages/ernie_image_pe.py` | `configs/pipeline_configs/ernie_image.py` |
+| Hunyuan3D | `runtime/pipelines/hunyuan3d_pipeline.py` | `stages/model_specific_stages/hunyuan3d/` | `configs/pipeline_configs/hunyuan3d.py` |
+| SANA-WM | `runtime/pipelines/sana_wm_pipeline.py`, `sana_wm_realtime_pipeline.py` | `stages/model_specific_stages/sana_wm/` | `configs/pipeline_configs/sana_wm.py` |
+| LingBot World realtime | `runtime/pipelines/lingbot_world_causal_dmd_pipeline.py` | `stages/model_specific_stages/lingbot_world/` | `configs/pipeline_configs/lingbot_world.py` |
+
+### Modular Style (when standard stages fit well)
+
+| Model | Pipeline | Notes |
+|-------|----------|-------|
+| Qwen-Image (T2I) | `runtime/pipelines/qwen_image.py` | Uses `add_standard_t2i_stages()` — standard text encoding + latent prep fits this model |
+| Qwen-Image-Edit | `runtime/pipelines/qwen_image.py` | Uses `add_standard_ti2i_stages()` — standard image-to-image flow |
+| Flux | `runtime/pipelines/flux.py` | Uses `add_standard_t2i_stages()` with custom `prepare_mu` |
+| FLUX.2 / FLUX.2 Klein | `runtime/pipelines/flux_2.py`, `flux_2_klein.py` | Reuses FLUX.2 stages; Klein differences live in config and sampling params |
+| Z-Image | `runtime/pipelines/zimage_pipeline.py` | Uses standard image pipeline stages plus Z-Image-specific config/model code |
+| Ideogram4 | `runtime/pipelines/ideogram.py` | Uses dedicated text encoding and denoising stages while keeping standard latent prep |
+| SANA | `runtime/pipelines/sana.py` | Spatial image pipeline; reuse the spatial image config pattern |
+| Stable Diffusion 3/3.5 | `runtime/pipelines/stable_diffusion_3.py` | Spatial image pipeline; compare scheduler, VAE scale, and conditioning layout |
+| LTX-2 / LTX-2.3 | `runtime/pipelines/ltx_2_pipeline.py` | Video pipeline family with one-stage, two-stage, and HQ variants |
+| Helios | `runtime/pipelines/helios_pipeline.py` | Video pipeline family with custom denoising and decoding stages |
+| FireRed/JoyAI image edit | `runtime/pipelines/qwen_image.py`, `runtime/pipelines/joy_image.py` | FireRed reuses Qwen edit-plus config; JoyAI has its own edit pipeline |
+| Wan | `runtime/pipelines/wan_pipeline.py` | Uses `add_standard_ti2v_stages()` |
+
+---
+
+## Checklist
+
+Before submitting, verify:
+
+**Common (both styles):**
+- [ ] **Pipeline file** exists at `runtime/pipelines/{model_name}.py` with `EntryClass`
+- [ ] **PipelineConfig** at `configs/pipeline_configs/{model_name}.py`
+- [ ] **SamplingParams** at `configs/sample/{model_name}.py`
+- [ ] **DiT model** at `runtime/models/dits/{model_name}.py`
+- [ ] **DiT config** at `configs/models/dits/{model_name}.py`
+- [ ] **VAE** — reuse existing (e.g., `AutoencoderKL`) or create new at `runtime/models/vaes/`
+- [ ] **VAE config** — reuse existing or create new at `configs/models/vaes/{model_name}.py`
+- [ ] **Registry entry** in `registry.py` via `register_configs()`
+- [ ] `pipeline_name` matches Diffusers `model_index.json` `_class_name`
+- [ ] `_required_config_modules` lists all modules from `model_index.json`
+- [ ] `PipelineConfig` callbacks (`prepare_pos_cond_kwargs`, `get_freqs_cis`, etc.) match DiT's `forward()` signature
+- [ ] Latent scale/shift factors are correctly configured
+- [ ] Use fused kernels where possible (see `existing-fast-paths.md` under the benchmark/profile skill)
+- [ ] Weight names match Diffusers for automatic loading
+- [ ] **TP/SP support** considered for DiT model (recommended; reference `wanvideo.py` for TP+SP, `qwen_image.py` for USPAttention)
+- [ ] **Output quality verified** — generated images/videos are not noise; compared against Diffusers reference output
+
+**Hybrid style only:**
+- [ ] **BeforeDenoisingStage** at `stages/model_specific_stages/{model_name}.py`
+- [ ] `BeforeDenoisingStage.forward()` populates all fields needed by `DenoisingStage`
+
+## Common Pitfalls
+
+1. **`batch.sigmas` must be a Python list**, not a numpy array. Use `.tolist()` to convert.
+2. **`batch.prompt_embeds` is a list of tensors** (one per encoder), not a single tensor. Wrap with `[tensor]`.
+3. **Don't forget `batch.raw_latent_shape`** -- `DecodingStage` uses it to unpack latents.
+4. **Rotary embedding style matters**: `is_neox_style=True` = split-half rotation, `is_neox_style=False` = interleaved. Check the reference model carefully.
+5. **VAE precision**: Many VAEs need fp32 or bf16 for numerical stability. Set `vae_precision` in the PipelineConfig accordingly.
+6. **Avoid forcing model-specific logic into shared stages**: If your model's pre-processing doesn't naturally fit the existing standard stages, prefer the Hybrid pattern with a dedicated BeforeDenoisingStage rather than adding conditional branches to shared stages.
+
+## After Implementation: Tests and Performance Data
+
+After the model produces non-noise output, read
+[references/testing-and-accuracy.md](references/testing-and-accuracy.md) before
+adding GPU cases, component-accuracy skips/hooks, suite entries, or benchmark
+claims. That reference tracks the current `gpu_cases.py` / `testcase_configs.py`
+/ `accuracy_testcase_configs.py` / `run_suite.py` split and the component-accuracy
+decision rules.

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-add-model/references/testing-and-accuracy.md

@@ -0,0 +1,80 @@
+# Testing And Accuracy
+
+Use this reference after a new diffusion model or pipeline variant can already
+produce a non-noise image or video.
+
+## Test Placement
+
+- Add concrete GPU integration cases in `python/sglang/multimodal_gen/test/server/gpu_cases.py`.
+- Keep reusable dataclasses, constants, thresholds, and testcase factory helpers in `python/sglang/multimodal_gen/test/server/testcase_configs.py`.
+- Add the case id to `python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py`
+  only when it should be part of component-accuracy coverage. Adding a GPU case
+  alone does not enroll it there.
+- Let `python/sglang/multimodal_gen/test/run_suite.py` own suite selection, runtime-based partitioning, and standalone test files. Do not hard-code CI shard lists elsewhere.
+- If a new standalone test file is added to a suite, update `STANDALONE_FILE_EST_TIMES` after the first measured CI/runtime value is known.
+
+Useful local entrypoints from repo root:
+
+```bash
+PYTHONPATH=python python3 python/sglang/multimodal_gen/test/run_suite.py --suite unit
+PYTHONPATH=python python3 python/sglang/multimodal_gen/test/run_suite.py --suite component-accuracy-1-gpu -k <case_id>
+PYTHONPATH=python python3 python/sglang/multimodal_gen/test/run_suite.py --suite 1-gpu --total-partitions 1 --partition-id 0 -k <case_id>
+```
+
+## Component Accuracy When Adding A GPU Case
+
+If you add a new entry to `ONE_GPU_CASES`, `TWO_GPU_CASES`, or a B200-specific
+case group in `gpu_cases.py`, treat component accuracy as part of the
+model-adding workflow. Do not assume the new testcase will automatically fit or
+enter the existing component-accuracy harness.
+
+The component-accuracy harness compares SGLang components against Diffusers/HF
+reference components. This is stricter than pipeline-level inference. New GPU
+cases commonly fail here for one of three reasons:
+
+1. The model family needs explicit hook wiring in `python/sglang/multimodal_gen/test/server/accuracy_hooks.py`.
+   - Add hook logic only when the harness cannot call the raw component correctly without it.
+   - Valid reasons include missing required forward arguments, required autocast/runtime context, or family-specific input preparation for the same component contract.
+   - Do not change the compared output mode or add harness-side behavior that changes the component contract just to make the test pass.
+
+2. The component is already covered by another testcase with the same source component and topology.
+   - Do not add redundant component-accuracy coverage.
+   - Add a skip entry in `python/sglang/multimodal_gen/test/server/accuracy_config.py` with a concrete reason such as `Representative VAE accuracy is already covered by ... for the same source component and topology`.
+   - This is the preferred path for variant-only cases such as LoRA, Cache-DiT, upscaling, or other cases that reuse the same underlying component weights and topology.
+
+3. The HF/Diffusers reference component cannot be loaded or compared faithfully in the harness.
+   - Add a skip entry in `accuracy_config.py` with the exact technical failure.
+   - Good reasons include missing/unsupported HF component layout, incomplete checkpoints, unsupported raw component contract, or proven divergence after matched weight transfer and matching output shape.
+   - Keep the skip reason concrete and technical. Do not write vague reasons like "component accuracy flaky" or "needs investigation."
+
+When adding a new GPU case, make this decision explicitly:
+
+- if the case should have component-accuracy coverage, add its case id to
+  `accuracy_testcase_configs.py`
+- if the family needs minimal harness wiring, add the smallest possible change in `accuracy_hooks.py`
+- if the case is only a variant of an already covered source component and topology, add a skip in `accuracy_config.py`
+- if the HF/Diffusers reference component cannot be compared faithfully, add a skip in `accuracy_config.py`
+- if the case is intentionally GPU-smoke-only, leave it out of `accuracy_testcase_configs.py` and keep that choice explicit in the PR notes
+
+Do not add a new GPU case and wait for CI to discover missing component-accuracy
+wiring.
+
+## Follow-up Scope
+
+Once the model is working and output quality is verified, cover the follow-up
+scope the user requested. If the user did not specify test or benchmark depth,
+propose the smallest useful validation set before launching long GPU runs.
+
+Tests should cover:
+
+- pipeline construction and stage wiring
+- single-GPU inference producing non-noise output
+- multi-GPU inference if TP/SP is supported
+- relevant unit tests for new math, parsing, scheduling, or loader behavior
+
+For performance data:
+
+- use the `warmup excluded` latency line for command-line generation
+- keep prompt, seed, shape, step count, model path, backend, and GPU topology fixed
+- use `sglang-diffusion-benchmark-profile` for denoise perf dumps and profiler traces
+- use `python/sglang/multimodal_gen/benchmarks/bench_serving.py` for serving benchmarks

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md

@@ -0,0 +1,80 @@
+---
+name: sglang-diffusion-benchmark-profile
+description: Use when benchmarking denoise latency or profiling a diffusion bottleneck in SGLang.
+---
+
+# SGLang Diffusion Benchmark and Profile
+
+Use this skill when measuring denoise performance, finding the slow op, checking whether an existing fast path can solve it, or verifying that a hotspot is real before any kernel work in `sglang.multimodal_gen`.
+
+This skill is diagnosis-first. It owns:
+- checked-in denoise benchmark presets
+- perf dump collection and before/after comparison
+- `torch.profiler` trace capture and quick hotspot ranking
+- mapping hot kernels back to known fast paths and fusion families
+- packaging confirmed kernel work with enough evidence for the appropriate kernel, Nsight, or framework-specific optimization workflow
+
+This skill does not own low-level kernel authoring or standalone Nsight workflows.
+
+## Preflight
+
+Before running any benchmark, profiler, or kernel-validation command:
+- use `scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__`
+- verify the repo is writable
+- export `HF_TOKEN` before using gated Hugging Face models such as `black-forest-labs/FLUX.*`
+- export `FLASHINFER_DISABLE_VERSION_CHECK=1`
+- choose idle GPU(s) before starting perf work
+
+## Native Backend Gate
+
+All diffusion benchmark and profiling results owned by this skill must come from the native SGLang diffusion backend.
+
+Treat any of the following as a hard stop condition:
+- `Falling back to diffusers backend`
+- `Using diffusers backend`
+- `Loaded diffusers pipeline`
+
+If any benchmark, perf-dump, or `torch.profiler` command prints one of those signals:
+- stop the workflow immediately
+- do not keep the generated numbers or traces as SGLang benchmark evidence
+- do not continue to hotspot classification or kernel work
+- first fix model resolution, pipeline selection, overlay/materialization, or other backend-selection issues so the model runs on the native SGLang diffusion path
+
+## Main Reference
+
+- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark, perf dump, and `torch.profiler` workflow; uses checked-in nightly-aligned presets plus current-source extras such as FLUX.2 Klein, Cosmos3, Ideogram4, ERNIE/GLM/SANA image models, FastWan2.2, `LTX-2.3` one-stage/two-stage/HQ, HunyuanVideo, MOVA, Helios, JoyAI/FireRed image edit, and Hunyuan3D shape
+- [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels, packed QKV paths, fused `QK norm + RoPE`, distributed overlap patterns, and open optimization PRs before proposing new code
+- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection
+- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; supports `--no-torch-compile`, validates nightly preset drift with `--validate-nightly-alignment`, and saves perf dumps by label for `compare_perf.py`
+
+## Opportunity Discovery Rule
+
+Before calling a diffusion hotspot "new", first classify it with `existing-fast-paths.md`.
+
+Always rule out these existing families first:
+- HunyuanVideo VAE GroupNorm+SiLU
+- LTX upsampler GroupNorm+SiLU
+- Z-Image residual-form modulation
+- SANA packed self-attention Q/K/V and cross-attention K/V GEMMs
+- fused diffusion `QK norm + RoPE`
+- LTX2 split RoPE
+- LTX2 residual-gate add
+- varlen USP attention pack/scatter
+- NVFP4 / Nunchaku packed QKV
+- Nunchaku fused GELU MLP
+- Ulysses / USP attention overlap
+- turbo-layer async all-to-all overlap
+- `torch.compile` compute / communication reorder
+- dual-stream diffusion execution
+
+If the user explicitly requires `torch.compile` to stay off, do not use the
+default benchmark preset invocation unchanged. Either pass the checked-in
+benchmark helper its no-compile switch or run the equivalent manual command
+without `--enable-torch-compile`.
+
+For FLUX-family manual profiling runs with a quantized transformer override:
+- use `sglang generate` directly
+- pass the override as `--transformer-path <dir>`
+- prefer `--prompt-path <file>` when also fixing `--output-file-name`
+- if the base model is already cached locally and the machine has unreliable HF access, use the local cached `--model-path` plus `HF_HUB_OFFLINE=1`
+- remember that `--profile` changes latency substantially; use the non-profile perf dump for the real before/after benchmark claim

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md

@@ -0,0 +1,523 @@
+---
+name: benchmark-and-profile-reference
+description: Reference commands and workflow for denoise benchmarks, perf dumps, and torch.profiler analysis in SGLang Diffusion.
+---
+
+# SGLang Diffusion Benchmark and Profile Guide
+
+**Primary Metric: Denoise Latency**
+- Denoise latency is the total DiT forward-pass time across all inference steps.
+- It is the dominant cost for diffusion inference and the main optimization target.
+- End-to-end latency and peak memory are secondary sanity checks.
+
+> **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images or videos against a reference baseline before and after any change.
+
+## Scope
+
+This guide intentionally stops at:
+- checked-in denoise benchmarks
+- structured perf dumps
+- `torch.profiler` trace capture
+- hotspot ranking
+- mapping hotspots to known fast paths
+
+If the hotspot survives this checklist, package the perf dump, profiler trace,
+exact command, and shape/topology notes for the appropriate kernel, Nsight, or
+framework-specific optimization workflow. Do not grow this skill back into a
+general Nsight or kernel-authoring guide.
+
+## Prerequisites
+
+```bash
+ENV_PY=python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py
+BENCH_PY=python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py
+ROOT=$(python3 "$ENV_PY" print-root)
+cd "$ROOT"
+python3 "$ENV_PY" check-write-access >/dev/null
+
+export HF_TOKEN=<your_hf_token>  # required for gated repos such as black-forest-labs/FLUX.*
+export FLASHINFER_DISABLE_VERSION_CHECK=1
+export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 1)
+
+ASSET_DIR=$(python3 "$ENV_PY" print-assets-dir --mkdir)
+BENCH_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir)
+PROFILE_DIR=$(python3 "$ENV_PY" print-output-dir --kind profiles --mkdir)
+export PROFILE_DIR
+
+check() {
+  local label="$1"
+  shift
+  "$@" &>/dev/null && echo "[OK]  $label" || echo "[MISS] $label"
+}
+
+check "sglang" python3 -c "import sglang"
+check "torch+CUDA" python3 -c "import torch; assert torch.cuda.is_available()"
+check "torch.profiler" python3 -c "import torch.profiler"
+```
+
+## Native Backend Gate
+
+Every benchmark and profile result in this guide must come from the native SGLang diffusion backend.
+
+If the command log contains any of:
+- `Falling back to diffusers backend`
+- `Using diffusers backend`
+- `Loaded diffusers pipeline`
+
+then stop immediately:
+- do not record the perf dump or trace as valid benchmark evidence
+- do not compare it against other runs
+- do not continue to hotspot ranking or kernel optimization
+- first fix backend selection so the model stays on the native SGLang diffusion path
+
+The checked-in benchmark helper pins `--backend=sglang` so native presets fail
+fast instead of silently falling back through `--backend=auto`. Do the same for
+manual native profiling commands unless you are intentionally collecting a
+diffusers baseline.
+
+Environment notes:
+- all commands below assume you are inside the configured diffusion container shell
+- export `HF_TOKEN` before any gated Hugging Face model run
+- export `FLASHINFER_DISABLE_VERSION_CHECK=1` before any benchmark or profiler run
+- re-run `print-idle-gpus` before each perf command if GPU availability may have changed
+- keep benchmark commands within 4 GPUs or fewer
+
+Download input images required by some presets:
+
+```bash
+wget -O "${ASSET_DIR}/cat.png" \
+  https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
+wget -O "${ASSET_DIR}/mova_single_person.jpg" \
+  https://github.com/OpenMOSS/MOVA/raw/main/assets/single_person.jpg
+```
+
+## Benchmark Presets
+
+Treat `"$BENCH_PY"` as the source of truth for preset order.
+
+Nightly diffusion comparison is server/API based (`sglang serve` plus requests).
+This skill stays on `sglang generate` for local benchmarking and profiling, but
+the nightly-aligned presets in `bench_diffusion_denoise.py` mirror
+`scripts/ci/utils/diffusion/comparison_configs.json` on model, task, prompt,
+reference image, size, frames, seed, GPU count, and SGLang serve args. If
+`comparison_configs.json` omits sampling params such as steps or guidance, the
+nightly-aligned `sglang generate` preset omits them too and relies on the same
+runtime defaults.
+When in doubt, re-check that JSON before trusting this reference.
+
+List the current preset order:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" --list-models
+```
+
+Check that the nightly presets still match the Nvidia nightly comparison config:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" --validate-nightly-alignment
+```
+
+Run one preset and save a perf dump:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model ltx2 \
+  --label baseline \
+  --output-dir "${BENCH_DIR}"
+```
+
+Keep `torch.compile` off when the task requires it:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model flux \
+  --label baseline \
+  --output-dir "${BENCH_DIR}" \
+  --no-torch-compile
+```
+
+Run the `LTX-2.3` one-stage skill preset:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model ltx23-one-stage \
+  --label baseline \
+  --output-dir "${BENCH_DIR}"
+```
+
+Run the nightly-aligned `LTX-2.3` TI2V two-stage preset:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model ltx23-ti2v-two-stage \
+  --label baseline \
+  --output-dir "${BENCH_DIR}"
+```
+
+Run the `LTX-2.3` two-stage skill preset:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model ltx23-two-stage \
+  --label baseline \
+  --output-dir "${BENCH_DIR}"
+```
+
+Run the full preset sweep only when you have enough GPU time for both the
+nightly-aligned cases and the source-tracked extras:
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --all \
+  --label prXXXX \
+  --output-dir "${BENCH_DIR}"
+```
+
+Nightly-aligned presets come first, followed by current-source extras from the
+registry / GPU test cases, then broader skill-only stress presets.
+
+Use the preset categories this way:
+- **Nightly-aligned**: exact mirrors of
+  `scripts/ci/utils/diffusion/comparison_configs.json`; use these when the goal
+  is apples-to-apples comparison with CI / nightly coverage.
+- **Current-source extras**: models or request shapes with explicit support
+  evidence in the current registry, GPU cases, compatibility matrix, pipeline
+  files, or unit tests, but without a nightly comparison case yet.
+- **Skill-only stress / coverage presets**: extra profiling scenarios kept by
+  this skill to stress a topology, high-resolution path, multi-GPU mode, or
+  model-specific stage. These may be older than the latest registry additions,
+  so re-check the active source tree before treating them as support-matrix
+  commitments.
+
+| Preset | Model | Nightly | Notes |
+| --- | --- | --- | --- |
+| `flux` | `black-forest-labs/FLUX.1-dev` | Yes: `flux1_dev_t2i_1024` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2, `--dit-layerwise-offload false`; no explicit steps/guidance override |
+| `flux2` | `black-forest-labs/FLUX.2-dev` | Yes: `flux2_dev_t2i_1024` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2, `--dit-layerwise-offload false`; no explicit steps/guidance override |
+| `qwen` | `Qwen/Qwen-Image-2512` | Yes: `qwen_image_2512_t2i_1024` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2; no explicit steps/guidance override |
+| `qwen-edit` | `Qwen/Qwen-Image-Edit-2511` | Yes: `qwen_image_edit_2511` | Uses the nightly cat image and edit prompt, 2 GPUs, TP size 2 |
+| `zimage` | `Tongyi-MAI/Z-Image-Turbo` | Yes: `zimage_turbo_t2i_1024` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2; no explicit steps/guidance override |
+| `wan-t2v` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | Yes: `wan22_t2v_a14b_720p` | 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
+| `wan-ti2v` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Yes: `wan22_ti2v_5b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, seed 42 |
+| `ltx2` | `Lightricks/LTX-2` | Yes: `ltx2_twostage_t2v` | `LTX2TwoStagePipeline`, 2 GPUs, CFG parallel, 768x512, 121 frames, seed 42 |
+| `ltx23-ti2v-two-stage` | `Lightricks/LTX-2.3` | Yes: `ltx2.3_twostage_ti2v_2gpus` | Nightly cat image, motion prompt, `LTX2TwoStagePipeline`, 2 GPUs, `--cfg-parallel-size 2`, 768x512, 121 frames, seed 42 |
+| `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
+| `qwen-image` | `Qwen/Qwen-Image` | No | Current-source extra covering the base Qwen-Image native path, separate from the nightly `Qwen-Image-2512` case |
+| `qwen-edit-2509` | `Qwen/Qwen-Image-Edit-2509` | No | Current-source extra for the pre-2511 edit-plus path; uses the cat image, 1024x1024 |
+| `zimage-base` | `Tongyi-MAI/Z-Image` | No | Current-source extra for non-turbo Z-Image; keep it separate from `zimage` / `Z-Image-Turbo` |
+| `flux2-klein` | `black-forest-labs/FLUX.2-klein-4B` | No | Current-source extra for the distilled FLUX.2 Klein path; gated repo, 1024x1024, DiT layerwise offload disabled |
+| `flux2-klein-base` | `black-forest-labs/FLUX.2-klein-base-4B` | No | Current-source extra for the undistilled FLUX.2 Klein Base path; gated repo, 1024x1024, DiT layerwise offload disabled |
+| `cosmos3-nano-t2i` | `nvidia/Cosmos3-Nano` | No | Current-source extra for the single-frame Cosmos3 image path; sets `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in the helper environment |
+| `cosmos3-nano-t2v` | `nvidia/Cosmos3-Nano` | No | Current-source extra for a short Cosmos3 video path; sets `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in the helper environment |
+| `ideogram4-fp8` | `ideogram-ai/ideogram-4-fp8` | No | Current-source extra matching the native Ideogram 4 FP8 pipeline; do not override steps/guidance directly because the sampling preset owns them |
+| `ernie-image-turbo` | `baidu/ERNIE-Image-Turbo` | No | Current-source extra for ERNIE-Image Turbo |
+| `glm-image` | `zai-org/GLM-Image` | No | Current-source extra for GLM-Image |
+| `sana-1.5-1.6b` | `Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers` | No | Current-source extra for a SANA native image path |
+| `fastwan22-ti2v-5b` | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | No | Current-source extra matching the FastWan2.2 TI2V registered path |
+| `ltx23-hq-two-stage` | `Lightricks/LTX-2.3` | No | Current-source extra for `LTX2TwoStageHQPipeline` with `--ltx2-two-stage-device-mode=original`; high-resolution and VRAM-heavy |
+| `ltx23-one-stage` | `Lightricks/LTX-2.3` | No | Skill-only extra preset for the native `LTX-2.3` one-stage baseline; 2 GPUs, 768x512, 121 frames, fps 24, 30 steps, guidance 3.0, seed 1234 |
+| `ltx23-two-stage` | `Lightricks/LTX-2.3` | No | Skill-only high-resolution stress preset for the native `LTX-2.3` two-stage path; uses `LTX2TwoStagePipeline`, 2 GPUs, 1536x1024, 121 frames, fps 24, 30 steps, guidance 3.0, seed 1234 |
+| `ltx23-two-stage-cfg-parallel` | `Lightricks/LTX-2.3` | No | Skill-only high-resolution CFG-parallel stress preset matching `ltx23-two-stage` plus `--cfg-parallel-size 2` |
+| `hunyuanvideo` | `hunyuanvideo-community/HunyuanVideo` | No | Skill-only extra preset |
+| `mova-720p` | `OpenMOSS-Team/MOVA-720p` | No | Skill-only extra preset |
+| `helios` | `BestWishYsh/Helios-Base` | No | Skill-only extra preset |
+| `joyai-edit` | `jdopensource/JoyAI-Image-Edit-Diffusers` | No | Skill-only JoyAI image-edit preset; uses the cat image, 1024x1024, 40 steps, guidance 4.0, 2-GPU CFG parallel |
+| `firered-edit-1.0` | `FireRedTeam/FireRed-Image-Edit-1.0` | No | Skill-only FireRed 1.0 image-edit preset; QwenImageEditPlus native path; uses 2-GPU CFG parallel |
+| `firered-edit-1.1` | `FireRedTeam/FireRed-Image-Edit-1.1` | No | Skill-only FireRed 1.1 image-edit preset; QwenImageEditPlus native path; uses 2-GPU CFG parallel |
+| `hunyuan3d-shape` | `tencent/Hunyuan3D-2` | No | Skill-only Hunyuan3D shape-generation preset; primary metric is `Hunyuan3DShapeDenoisingStage` |
+
+For Wan2.2 video models, remember the difference between **nightly alignment**
+and **best latency tuning**:
+- the nightly-aligned 4-GPU commands intentionally keep `--enable-cfg-parallel --ulysses-degree=2` so CFG and ring behavior stay covered
+- do not assume that is the fastest topology
+- for pure latency tuning, benchmark pure Ulysses too, for example `--ulysses-degree=4 --ring-degree=1` on 4 GPUs, and on 8 GPUs compare pure `--ulysses-degree=8` against `--enable-cfg-parallel --ulysses-degree=4`
+
+### Manual command example: LTX-2 Two-Stage
+
+```bash
+sglang generate \
+  --model-path=Lightricks/LTX-2 \
+  --pipeline-class-name=LTX2TwoStagePipeline \
+  --prompt="A cat and a dog baking a cake together in a kitchen." \
+  --width=768 --height=512 \
+  --num-frames=121 \
+  --seed=42 --num-gpus=2 --enable-cfg-parallel \
+  --save-output --enable-torch-compile --warmup
+```
+
+`LTX2TwoStagePipeline` is a native path. The spatial upsampler and distilled
+LoRA are auto-resolved from the same model snapshot unless you override them.
+
+### Manual command example: LTX-2.3 TI2V Two-Stage
+
+```bash
+sglang generate \
+  --model-path=Lightricks/LTX-2.3 \
+  --pipeline-class-name=LTX2TwoStagePipeline \
+  --prompt="The cat starts walking slowly towards the camera." \
+  --image-path="${ASSET_DIR}/cat.png" \
+  --width=768 --height=512 \
+  --num-frames=121 \
+  --seed=42 --num-gpus=2 --cfg-parallel-size=2 \
+  --save-output --enable-torch-compile --warmup
+```
+
+This matches the nightly comparison case `ltx2.3_twostage_ti2v_2gpus`.
+
+### Manual command example: LTX-2.3 One-Stage
+
+```bash
+sglang generate \
+  --model-path=Lightricks/LTX-2.3 \
+  --prompt="A beautiful sunset over the ocean" \
+  --negative-prompt="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
+  --width=768 --height=512 \
+  --num-frames=121 --fps=24 \
+  --num-inference-steps=30 --guidance-scale=3.0 \
+  --seed=1234 --num-gpus=2 \
+  --save-output --enable-torch-compile --warmup
+```
+
+Use this when you want the native `LTX2Pipeline` baseline for `LTX-2.3` at the
+validated one-stage resolution.
+
+### Manual command example: LTX-2.3 Two-Stage High-Resolution Stress
+
+```bash
+sglang generate \
+  --model-path=Lightricks/LTX-2.3 \
+  --pipeline-class-name=LTX2TwoStagePipeline \
+  --prompt="A beautiful sunset over the ocean" \
+  --negative-prompt="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
+  --width=1536 --height=1024 \
+  --num-frames=121 --fps=24 \
+  --num-inference-steps=30 --guidance-scale=3.0 \
+  --seed=1234 --num-gpus=2 \
+  --save-output --enable-torch-compile --warmup
+```
+
+This matches the skill-only `ltx23-two-stage` preset. Use it as a
+high-resolution stress target, not as a nightly comparison case.
+
+### Manual command example: JoyAI Image Edit
+
+```bash
+sglang generate \
+  --backend=sglang \
+  --model-path=jdopensource/JoyAI-Image-Edit-Diffusers \
+  --prompt="Make the cat wear a red hat" \
+  --image-path="${ASSET_DIR}/cat.png" \
+  --width=1024 --height=1024 \
+  --num-inference-steps=40 --guidance-scale=4.0 \
+  --num-gpus=2 --enable-cfg-parallel --ulysses-degree=1 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --save-output --enable-torch-compile --warmup
+```
+
+### Manual command example: FireRed Image Edit
+
+```bash
+sglang generate \
+  --backend=sglang \
+  --model-path=FireRedTeam/FireRed-Image-Edit-1.1 \
+  --prompt="Make the cat wear a red hat" \
+  --image-path="${ASSET_DIR}/cat.png" \
+  --width=1024 --height=1024 \
+  --num-inference-steps=40 --guidance-scale=4.0 \
+  --num-gpus=2 --enable-cfg-parallel --ulysses-degree=1 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --save-output --enable-torch-compile --warmup
+```
+
+Use `FireRedTeam/FireRed-Image-Edit-1.0` in the same command when comparing the
+1.0 checkpoint. Both FireRed presets use the native `QwenImageEditPlusPipeline`
+path. On H100, 2-GPU CFG parallel reduced 40-step denoise latency versus the
+otherwise matching 2-GPU Ulysses command: FireRed 1.0 from 13419.15 ms to
+10955.90 ms, and FireRed 1.1 from 13414.72 ms to 10934.21 ms.
+
+### Manual command example: Hunyuan3D Shape
+
+```bash
+OUTPUT_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir)
+CONFIG_DIR="${OUTPUT_DIR}/generated_configs"
+mkdir -p "${CONFIG_DIR}"
+printf '{"paint_enable": false}\n' > "${CONFIG_DIR}/hunyuan3d-shape.json"
+
+sglang generate \
+  --backend=sglang \
+  --model-path=tencent/Hunyuan3D-2 \
+  --prompt="generate 3d mesh" \
+  --image-path="${ASSET_DIR}/cat.png" \
+  --config="${CONFIG_DIR}/hunyuan3d-shape.json" \
+  --num-inference-steps=50 --guidance-scale=5.0 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --save-output --enable-torch-compile --warmup
+```
+
+For Hunyuan3D, compare the denoise stage separately from mesh export and paint
+stages. The benchmark helper reports `Hunyuan3DShapeDenoisingStage` as the
+primary denoise metric.
+
+### Manual command example: Wan2.2-I2V-A14B 720P
+
+```bash
+# Select four idle GPUs first:
+# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4)
+sglang generate \
+  --model-path=Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+  --prompt="The cat starts walking slowly towards the camera." \
+  --image-path="${ASSET_DIR}/cat.png" \
+  --width=1280 --height=720 --num-frames=81 \
+  --seed=42 --save-output \
+  --num-gpus=4 --enable-cfg-parallel --ulysses-degree=2 \
+  --text-encoder-cpu-offload --pin-cpu-memory \
+  --warmup --enable-torch-compile
+```
+
+`Wan2.2-I2V-A14B` uses the 720p max-area config by default, and explicit
+`--width/--height` overrides control the target area while preserving the
+reference-image aspect ratio.
+
+## Perf Dump Workflow
+
+For every benchmark run, write a perf dump JSON:
+
+```bash
+sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/<result>.json"
+```
+
+Before/after comparison:
+
+```bash
+python3 python/sglang/multimodal_gen/benchmarks/compare_perf.py \
+  "${BENCH_DIR}/baseline.json" \
+  "${BENCH_DIR}/new.json"
+```
+
+Always keep:
+- denoise latency
+- end-to-end latency
+- peak GPU memory
+- exact command line, model shape, dtype, and GPU topology
+
+Never keep a perf dump produced after a diffusers-backend fallback.
+
+## `torch.profiler` Workflow
+
+### 1. Establish the baseline
+
+```bash
+PYTHONPATH=python python3 "$BENCH_PY" \
+  --model flux \
+  --label baseline \
+  --output-dir "${BENCH_DIR}"
+```
+
+Keep model shape, seed, and GPU topology fixed for every comparison. Save one
+reference image or video before changing code. If the active task requires
+`torch.compile` off, add `--no-torch-compile` here too.
+
+### 2. Capture a representative trace
+
+By default SGLang profiles the denoising stage. The default sampling window is
+5 profiled timesteps after warmup.
+
+```bash
+SGLANG_DIFFUSION_TORCH_PROFILER_DIR="${PROFILE_DIR}/torch" \
+sglang generate \
+  --model-path=black-forest-labs/FLUX.1-dev \
+  --prompt="A futuristic cyberpunk city at night" \
+  --width=1024 --height=1024 --num-inference-steps=50 \
+  --seed=42 --enable-torch-compile --warmup \
+  --profile
+```
+
+Use `--profile-all-stages` only when you really need text encoder, VAE, or
+other non-denoise stages too.
+
+The generated trace path is printed in the console and also lands under
+`SGLANG_DIFFUSION_TORCH_PROFILER_DIR`. The diffusion profiler falls back to
+`SGLANG_TORCH_PROFILER_DIR` and then `./logs` when the diffusion-specific env
+var is unset. Open the trace in Perfetto if you want a timeline view:
+- https://ui.perfetto.dev/
+
+### 3. Rank the hot CUDA kernels
+
+Use this parser for a quick top-k table without opening a browser:
+
+```python
+import collections
+import glob
+import gzip
+import json
+import os
+
+log_dir = (
+    os.environ.get("SGLANG_DIFFUSION_TORCH_PROFILER_DIR")
+    or os.environ.get("SGLANG_TORCH_PROFILER_DIR")
+    or "./logs"
+)
+trace_path = sorted(
+    glob.glob(f"{log_dir}/*.trace.json.gz"),
+    key=os.path.getmtime,
+    reverse=True,
+)[0]
+
+with gzip.open(trace_path, "rb") as f:
+    data = json.loads(f.read())
+
+cuda_ops = collections.defaultdict(lambda: {"total_us": 0, "count": 0})
+for event in data.get("traceEvents", []):
+    if event.get("cat") in ("kernel", "gpu_memcpy") and "dur" in event:
+        cuda_ops[event.get("name", "unknown")]["total_us"] += event["dur"]
+        cuda_ops[event.get("name", "unknown")]["count"] += 1
+
+print(f"{'Kernel':<90} {'Total(ms)':>10} {'Count':>6}")
+for name, stat in sorted(cuda_ops.items(), key=lambda item: -item[1]["total_us"])[:30]:
+    print(f"{name:<90} {stat['total_us'] / 1000:>10.3f} {stat['count']:>6}")
+```
+
+If you need better attribution, add `record_function(...)` scopes around DiT
+attention, norm, modulation, MLP, or communication boundaries and re-run.
+
+### 4. Classify the hotspot with `existing-fast-paths.md`
+
+Do not jump from a hot kernel straight into new code. First classify it against
+the known mainline families.
+
+| What the trace shows | First interpretation |
+| --- | --- |
+| `fused_inplace_qknorm_rope` missing, but separate qk norm plus rope show up | Check whether the fused diffusion `QK norm + RoPE` path should have engaged |
+| `to_q -> to_k -> to_v` on NVFP4 or Nunchaku FLUX-family checkpoints | Treat as a packed-QKV fast-path miss or checkpoint-format mismatch |
+| `fused_norm_tanh_mul_add*` missing on Z-Image | Treat as a missing mainline modulation path, not a new fusion request |
+| LTX-2 split RoPE appears as a long PyTorch elementwise chain | Check the `apply_ltx2_split_rotary_emb` Triton path and its shape guards |
+| masked attention spends time packing/unpacking Q/K/V | Check whether fused varlen USP pack/scatter should have engaged |
+| `all_to_all`, ring attention, or async A2A dominate | Classify against Ulysses, USP, or turbo-layer overlap first |
+| split `fc1 -> gelu -> quant -> fc2.lora_down` on Nunchaku FLUX | Treat as a missing fused GELU MLP path |
+| attention kernels dominate | Confirm backend, topology, and shape guards before proposing a new kernel |
+
+If the hot path is already covered by a mainline optimization family, fix the
+enablement, shape guard, backend choice, or checkpoint mapping first.
+
+### 5. Hand off only real kernel work
+
+Only after the hotspot survives the fast-path checklist:
+
+1. save a baseline perf dump
+2. save a representative `torch.profiler` trace
+3. note the exact model, shape, dtype, and GPU topology
+4. hand the work to the appropriate kernel, Nsight, or framework-specific optimization workflow
+
+This skill intentionally stops here. It tells you whether you are looking at:
+- a missing existing optimization
+- a configuration or backend problem
+- or a real kernel opportunity worth handing off
+
+## Minimal Merge Checklist
+
+- [ ] fixed-shape baseline perf dump saved
+- [ ] fixed-shape new perf dump saved
+- [ ] `compare_perf.py` table generated
+- [ ] one representative `torch.profiler` trace saved
+- [ ] hotspot classified against `existing-fast-paths.md`
+- [ ] reference image or video checked for correctness
+- [ ] any remaining kernel work handed off with perf/profile evidence attached

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

@@ -0,0 +1,261 @@
+# SGLang Diffusion Fast Paths
+
+Use this guide when mapping a diffusion bottleneck to an existing fused path or
+distributed overlap pattern in `sglang.multimodal_gen`. Prefer reuse and
+configuration first before handing the problem to a kernel, Nsight, or
+framework-specific optimization workflow.
+
+**Key Files**
+- `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
+- `python/sglang/multimodal_gen/runtime/layers/elementwise.py`
+- `python/sglang/multimodal_gen/runtime/layers/fused_scale_shift_gate.py`
+- `python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py`
+- `python/sglang/jit_kernel/diffusion/triton/scale_shift.py`
+- `python/sglang/jit_kernel/diffusion/group_norm_silu.py`
+- `python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py`
+- `python/sglang/jit_kernel/diffusion/triton/norm.py`
+- `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py`
+- `python/sglang/jit_kernel/diffusion/triton/rotary.py`
+- `python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py`
+- `python/sglang/jit_kernel/diffusion/residual_gate_add.py`
+- `python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh`
+- `python/sglang/jit_kernel/diffusion/triton/varlen_pack_pad.py`
+- `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
+- `test/registered/jit/diffusion/test_qwen_image_modulation.py`
+- `test/registered/jit/diffusion/test_group_norm_silu.py`
+- `test/registered/jit/diffusion/test_residual_gate_add.py`
+- `test/registered/jit/diffusion/test_varlen_pack_pad.py`
+- `test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py`
+- `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py`
+- `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`
+- `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py`
+- `python/sglang/jit_kernel/norm.py`
+- `python/sglang/multimodal_gen/runtime/platforms/cuda.py`
+- `python/sglang/multimodal_gen/runtime/layers/attention/selector.py`
+- `docs/diffusion/performance/attention_backends.md` (repo root)
+
+**Core Fusion Patterns**
+
+1. Scale/Shift elementwise and gate fusion (AdaLN modulation)
+- Kernels: `fuse_scale_shift_kernel`, `fuse_layernorm_scale_shift_gate_select01_kernel`, `fuse_residual_layernorm_scale_shift_gate_select01_kernel`
+- Locations: `elementwise.py`, `layernorm.py`, `fused_scale_shift_gate.py`, `qwen_image.py`, `triton/scale_shift.py`
+- Use cases: `x * (1 + scale) + shift`, `a * (k + b) + c`, and Qwen-style `(layernorm/residual layernorm) + scale/shift + gate select`.
+- Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`.
+- NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path.
+- Validation: `test/registered/jit/diffusion/test_qwen_image_modulation.py`.
+
+2. Norm + Scale/Shift fusion (CuTe DSL)
+- Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift`
+- Locations: `layernorm.py`, `cutedsl/scale_residual_norm_scale_shift.py`
+- Use cases:
+  - `y = norm(x) * (1 + scale) + shift`
+  - `y = norm(residual + gate * x) * (1 + scale) + shift`
+- Constraints: `D % 256 == 0` and `D <= 8192`. `x/residual/gate/scale/shift` must pass shape and stride validation. Dtypes limited to fp16/bf16/fp32.
+- Behavior: CuTe DSL compilation cached by `(dtype, ndim, D, norm_type)`. `None` tensors replaced by scalar placeholders. If constraints fail, `layernorm.py` warns and falls back to native PyTorch.
+
+3. Z-Image fused tanh/gate modulation
+- Kernels: `fused_norm_tanh_mul_add`, `fused_norm_tanh_mul_add_norm_scale`
+- Locations: `layernorm.py`, `cutedsl/norm_tanh_mul_add_norm_scale.py`, `zimage.py`
+- Use cases:
+  - `y = tanh(gate) * norm(x) + shift`
+  - `y, y2 = tanh(gate) * norm(x) + shift`, then `y2 = norm(y) * (1 + scale)`
+- Constraints: same CuTe DSL envelope as the norm+scale/shift family in practice: contiguous last dim, fp16/bf16/fp32, and `D % 256 == 0`, `D <= 8192`.
+- Validation: `test/registered/jit/diffusion/test_norm_tanh_mul_add_norm_scale.py`
+- Behavior: this is already a mainline fast path, so if Z-Image traces show the unfused chain, treat it as a missing or regressed existing optimization before proposing a new kernel.
+
+4. Triton LayerNorm/RMSNorm fusion
+- Kernels: `rms_norm_fn`, `layer_norm_fn`, `norm_infer`
+- Locations: `triton/norm.py`, `layernorm.py`
+- Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`.
+- Constraints: last dim must be contiguous, and `N * element_size < 64KB`.
+- Validation: `test/registered/jit/test_rmsnorm.py`.
+
+5. Triton one-pass RMSNorm (small hidden size fast path)
+- Kernel: `triton_one_pass_rms_norm`
+- Locations: `triton/rmsnorm_onepass.py`, `layernorm.py`
+- Use case: `hidden_size <= 128` in `RMSNorm.forward_cuda`.
+- `torch.compile` note: keep this path behind the custom-op wrapper in `rmsnorm_onepass.py`; direct `wrap_triton` can recompile on dynamic row counts.
+
+6. Triton RoPE fusion
+- Kernel: `apply_rotary_embedding`
+- Locations: `triton/rotary.py`, `rotary_embedding/utils.py`
+- Use case: GPT-J style RoPE when not Neox.
+- Constraints: `head_size` must be even.
+- NPU fallback: `npu_fallback.apply_rotary_embedding_native`.
+- Validation: `test/registered/jit/test_rope.py`.
+
+7. LTX2 split RoPE fusion
+- Kernel: `apply_ltx2_split_rotary_emb`
+- Locations: `triton/ltx2_rotary.py`, `runtime/models/dits/ltx_2.py`
+- Use case: LTX-2 split rotary embedding over `[B, S, num_heads * head_dim]` with separate `cos` and `sin` tensors.
+- Constraints: `cos` and `sin` shapes must match `[B, H, S, head_dim / 2]`, and `inner_dim == H * head_dim`.
+- Workflow rule: if LTX-2 traces show a large split-RoPE PyTorch chain, check whether the LTX2-specific Triton path was disabled by shape or dtype before proposing a new RoPE kernel.
+
+8. LTX2 residual-gate add fusion
+- Kernel: `diffusion_residual_gate_add`
+- Locations: `diffusion/residual_gate_add.py`, `csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`
+- Use case: `residual + update * gate` in LTX2 self-attention, prompt cross-attention, audio/video cross-attention, and feed-forward residual updates.
+- Constraints: `residual`, `update`, and `gate` must be CUDA tensors on the same device, contiguous, same dtype (`fp16`, `bf16`, or `fp32`), with `update.shape == residual.shape`; `gate` can match `residual` or be row-broadcast with the last dimension matching.
+- Behavior: `_ltx2_residual_gate_add(...)` uses the CUDA custom op while guards pass. On a runtime exception outside `torch.compile`, it logs once, disables the fast path for the process, and falls back to `residual + update * gate`.
+- Validation: `test/registered/jit/diffusion/test_residual_gate_add.py`.
+- Microbench: `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py`.
+- Workflow rule: if LTX2 traces show repeated elementwise `mul` + `add` ladders around attention or MLP residuals, check whether this existing CUDA path was disabled by shape, dtype, contiguity, or a prior runtime failure before proposing another elementwise fusion.
+
+9. HunyuanVideo / LTX upsampler GroupNorm + SiLU fusion
+- Kernel: `triton_group_norm_silu`
+- Locations: `diffusion/group_norm_silu.py`, `triton/group_norm_silu.py`, `runtime/models/vaes/hunyuanvae.py`, `runtime/models/upsampler/latent_upsampler.py`
+- Use case: `activation(group_norm(x))` when the activation is non-inplace `nn.SiLU` and the GroupNorm is affine.
+- Enablement: mainline uses `apply_group_norm_silu(...)` in HunyuanVideo VAE paths and LTX latent upsampler paths by default; there is no env toggle. The wrapper dispatches to Triton only when guards pass.
+- Constraints: CUDA inference path only; no grad, `x.requires_grad == False`, `nn.GroupNorm`, `nn.SiLU(inplace=False)`, affine norm with weight and bias. Unsupported cases fall back to native `activation(norm(x))`.
+- Validation: `test/registered/jit/diffusion/test_group_norm_silu.py`.
+- Microbench: `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`.
+
+**Faster CUDA Kernel Usage Points**
+
+1. sgl-kernel RMSNorm and fused add RMSNorm
+- Location: `layernorm.py`
+- Behavior:
+- Standard `bf16`/`fp16` CUDA paths use `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`.
+- The Z-Image `fp32` `32x2560` path under `torch.compile` avoids `wrap_triton` and uses the native fp32 path.
+- `hidden_size <= 128` uses Triton one-pass.
+- ROCm falls back to native.
+
+2. Attention backend selection (FlashAttention, Sage, SDPA)
+- Locations: `platforms/cuda.py`, `attention/selector.py`, `docs/diffusion/performance/attention_backends.md`
+- Behavior: CUDA prefers FlashAttention (FA3/FA4) when supported, otherwise Torch SDPA. Force via `--attention-backend` or `global_force_attn_backend`.
+
+3. FlashInfer RoPE (Q/K inplace)
+- Location: `rotary_embedding/utils.py`
+- Behavior: `flashinfer.rope.apply_rope_with_cos_sin_cache_inplace` when available, otherwise Triton RoPE fallback.
+
+4. Varlen USP attention pack/scatter
+- Locations: `runtime/layers/attention/layer.py`, `triton/varlen_pack_pad.py`
+- Behavior: masked `USPAttention.forward` can gather dense Q/K/V into packed `[total_valid, H, D]` rows with `fused_pack_qkv`, run varlen attention, then scatter back with `fused_scatter_to_padded`.
+- Validation: `test/registered/jit/diffusion/test_varlen_pack_pad.py` and `test_varlen_uspattn_equivalence.py`.
+- Workflow rule: if a masked attention trace spends time in Python/advanced indexing pack or scatter, first check whether this fused varlen path should have engaged.
+
+**QK Norm Optimization**
+
+- Entry point: `apply_qk_norm` in `layernorm.py`.
+- Fast path: JIT fused inplace QK norm from `python/sglang/jit_kernel/norm.py` via `fused_inplace_qknorm`.
+- Preconditions for fused path:
+  - CUDA only.
+  - `allow_inplace=True` and `q_eps == k_eps`.
+  - `can_use_fused_inplace_qknorm(head_dim, dtype)` returns true.
+  - Supported head dims: `64, 128, 256, 512, 1024`.
+- Behavior: Fused path operates on `q` and `k` in place after reshaping to `[B, -1, head_dim]`. If preconditions fail, fall back to per-tensor RMSNorm.
+- Validation: `test/registered/jit/test_qknorm.py` and `test/registered/jit/test_qknorm_across_heads.py`.
+
+**QK Norm + RoPE Optimization**
+
+- Entry point: `apply_qk_norm_rope` in `layernorm.py`.
+- Fast path: JIT fused inplace QK norm + RoPE from `python/sglang/jit_kernel/diffusion/qknorm_rope.py` via `fused_inplace_qknorm_rope`.
+- Toggle: `SGLANG_ENABLE_FUSED_QKNORM_ROPE=1` keeps the fused path enabled by default.
+- Preconditions for fused path:
+  - CUDA only.
+  - `allow_inplace=True` and `q_eps == k_eps`.
+  - `q` / `k` are contiguous 4D tensors with the same shape.
+  - `q.dtype` is `fp16` or `bf16`, and norm weights match tensor dtype.
+  - `can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, dtype)` returns true.
+  - Supported head dims: `64, 128, 256`.
+- Behavior: `apply_qk_norm_rope` prefers the fused JIT kernel when all guards pass; otherwise it falls back to `apply_qk_norm(...)` plus `apply_flashinfer_rope_qk_inplace(...)`.
+- Validation: `test/registered/jit/diffusion/test_qknorm_rope.py`.
+- Workflow rule: treat LTX2 traces that miss the generic fused path as an enablement/shape-guard issue first, and check the separate LTX2 split-RoPE path before proposing new attention-prep kernels.
+
+**Nunchaku Fused GELU MLP**
+
+- Entry point: `_fused_gelu_mlp` in `runtime/models/dits/flux.py`.
+- Fast path: Nunchaku checkpoints can fuse `fc1 GEMM + GELU + shift + re-quant + fc2.lora_down` before the second GEMM instead of materializing a standalone GELU activation.
+- Scope: this is a model-specific fast path for Nunchaku-quantized FLUX-family checkpoints.
+- Workflow rule: if a Nunchaku trace shows split `fc1 -> gelu -> quant -> fc2.lora_down`, treat it as a missing existing fast path before proposing a new fusion.
+
+**NVFP4 / Nunchaku Packed QKV**
+
+- Entry points: `runtime/models/dits/flux.py`, `runtime/models/dits/flux_2.py`, and the FLUX config remapping in `configs/models/dits/flux.py`.
+- Fast path: quantized FLUX-family checkpoints can store attention projections in packed QKV form, and SGLang intentionally switches to `MergedColumnParallelLinear` paths such as `to_qkv`, `to_added_qkv`, and `to_qkv_mlp_proj` instead of separate `to_q`, `to_k`, `to_v`.
+- FLUX.2 NVFP4 note: `flux_2.py` explicitly enables fused packed QKV when `quant_config` is `ModelOptFp4Config`, because the NVFP4 checkpoint stores image-attention QKV packed on disk.
+- Nunchaku note: raw and converted Nunchaku checkpoint names are remapped onto fused `to_qkv` / `to_added_qkv` names in `configs/models/dits/flux.py`; correctness on NVFP4-style checkpoints also depends on quant metadata such as `wtscale` and attention `wcscales`.
+- Workflow rule: if an NVFP4 or Nunchaku trace shows split `to_q -> to_k -> to_v` where packed QKV is expected, treat it as a missing quantized fast path or checkpoint-format mismatch before proposing a new attention fusion.
+
+**SANA Packed Projection GEMMs**
+
+- Entry point: `runtime/models/dits/sana.py`.
+- Fast path: SANA self-attention uses one `MergedColumnParallelLinear` `to_qkv` GEMM for Q/K/V, and SANA cross-attention uses one `MergedColumnParallelLinear` `to_kv` GEMM for encoder K/V.
+- Scope: this is a mainline SANA model fast path. Query projection in cross-attention remains separate because it uses denoising hidden states, while K/V share step-invariant encoder hidden states.
+- Workflow rule: if a SANA trace shows separate self-attention `to_q`, `to_k`, `to_v` GEMMs, or separate cross-attention `to_k` and `to_v` GEMMs, treat that as a regressed existing packed-projection path before proposing a new GEMM fusion.
+
+**Common Entry Points in Diffusion Models**
+- AdaLN modulation: `LayerNormScaleShift`, `RMSNormScaleShift`, `ScaleResidual*` in `layernorm.py`.
+- Qwen-Image gating: `fuse_layernorm_scale_shift_gate_select01_kernel` and `fuse_residual_layernorm_scale_shift_gate_select01_kernel` through `fused_scale_shift_gate.py` and `qwen_image.py`.
+- Z-Image residual-form modulation: `fused_norm_tanh_mul_add` and `fused_norm_tanh_mul_add_norm_scale` in `zimage.py`.
+- HunyuanVideo VAE and LTX upsampler GroupNorm+SiLU: `apply_group_norm_silu` in `hunyuanvae.py` and `latent_upsampler.py`; default-eligible when wrapper guards pass.
+- QK norm: `apply_qk_norm` used in `flux.py`, `flux_2.py`, `qwen_image.py`, `zimage.py`, `wanvideo.py`, `ltx_2.py`, `hunyuanvideo.py`.
+- QK norm + RoPE: `apply_qk_norm_rope` in `layernorm.py`; use this path when the model wants fused attention prep instead of separate QK norm and RoPE calls.
+- LTX2 split RoPE: `apply_ltx2_split_rotary_emb` in `ltx_2.py`.
+- LTX2 residual-gate add: `_ltx2_residual_gate_add` in `ltx_2.py` wraps the CUDA `diffusion_residual_gate_add` custom op for attention, cross-attention, and MLP residual updates.
+- Varlen USP attention: `fused_pack_qkv` and `fused_scatter_to_padded` in `attention/layer.py`.
+- SANA packed projections: `to_qkv` and `to_kv` in `sana.py`.
+- Nunchaku fused GELU MLP: `_fused_gelu_mlp` in `flux.py` for quantized FLUX-family checkpoints.
+- NVFP4 / packed QKV attention: `to_qkv`, `to_added_qkv`, and `to_qkv_mlp_proj` in FLUX-family quantized paths.
+- RoPE: `_apply_rotary_emb` prefers Triton; Q/K RoPE prefers FlashInfer when present.
+
+**Existing Overlap / Communication Families**
+
+- Ulysses / USP attention: treat `all_to_all`, `ring_attn`, and head / sequence reshards as an existing distributed attention family, not a new overlap idea.
+- Turbo-layer async all-to-all: `all_to_all_single(..., async_op=True)` plus staged waits already form an existing overlap family in `turbo_layer.py`.
+- TorchInductor compute / communication reorder: `torch._inductor.config.reorder_for_compute_comm_overlap = True` can already partially overlap compiled denoise traces.
+- Dual-stream diffusion models: `use_dual_stream = True` in models such as `hunyuan3d.py` is an existing overlap family.
+- Workflow rule: if a hotspot is communication-heavy, rule out these in-repo overlap families before proposing a brand new overlap design.
+
+**Historical PR Watchlist**
+
+These SGLang PRs are useful as upstream direction and prior art, not as
+current-main behavior. Re-check the PR state and the active source tree before
+relying on any file path, flag, or claim about whether the work has merged.
+
+- Norm, modulation, and packed projection fusions:
+  - #24025 LTX2 QK norm fusion.
+  - #24059 Helios fused norm modulation.
+  - #24117 Z-Image packed QKV.
+  - #19488 Wan cross-block elementwise fusion.
+  - #19249 Z-Image `scale residual norm scale shift` plus `add gate norm` fusion.
+  - #18897 dual norm fusion for FLUX-family paths (draft).
+  - #20429 Qwen-Image layernorm and `fuse_scale_shift_gate_select01` work.
+  - #20530 MOVA fused RMSNorm + interleaved RoPE.
+  - #29361 LTX2 residual-gate CUDA fast path for `residual + update * gate`.
+- VAE and decode-side acceleration:
+  - #22531 LTX2 parallel VAE support and #20927 batched tiled VAE decode (draft).
+- Attention, communication, and runtime scheduling:
+  - #22805 FLUX.2 packed QKV for all-to-all.
+  - #21742 hybrid attention schedule.
+  - #24053 USP attention with replicated prefixes.
+  - #18764 dynamic batching v0.
+  - #24200 disaggregated diffusion v2.
+- Cache and CUDA graph:
+  - #21613 TeaCache refactor.
+  - #24227 WanVideo TeaCache skipping fix.
+  - #20447 TeaCache support for GLM-Image, Qwen-Image, and related models.
+  - #19516 Qwen-Image CUDA Graph.
+  - #21912 Z-Image Turbo FP8 full quantization and CUDA Graph.
+
+**Constraints and Fallbacks**
+- `scale_shift` Triton requires CUDA + contiguous `x`. NPU swaps to native.
+- CuTe DSL fused norms require `D % 256 == 0` and `D <= 8192`.
+- Triton norm kernels error on feature size >= 64KB.
+- FlashAttention requires fp16/bf16 and SM80+; otherwise SDPA.
+
+**Integration Checklist for New Models**
+
+1. Reuse `LayerNormScaleShift` or `ScaleResidual*` modules instead of re-implementing fusion logic.
+2. Keep tensors contiguous and satisfy D alignment (`% 256`) and size (`<= 8192`) for CuTe fused paths.
+3. Use `fuse_scale_shift_kernel` for AdaLN modulation and keep a PyTorch fallback.
+4. Use `apply_qk_norm` and ensure head_dim is in the supported list for fused QK norm.
+5. If using FlashInfer RoPE, avoid `pack qkv` and ensure Q/K are contiguous.
+6. For attention, follow `selector.py` priority; override with CLI only if needed.
+
+**When Extending or Modifying Kernels**
+- Add `torch.library.custom_op` and `register_fake` for compile and meta support.
+- Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`.
+- Avoid implicit broadcasts that force hidden `contiguous()` copies.
+- Preserve NPU and ROCm fallback paths.
+- If none of the families above match, package the evidence from the benchmark/profile skill and hand the kernel work to the appropriate kernel, Nsight, or framework-specific optimization workflow.

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import os
+import subprocess
+from pathlib import Path
+
+OUTPUT_DIR_NAMES = {
+    "benchmarks": Path("outputs/diffusion_benchmarks"),
+    "profiles": Path("outputs/diffusion_profiles"),
+}
+
+
+def get_repo_root() -> Path:
+    import sglang
+
+    return Path(sglang.__file__).resolve().parents[2]
+
+
+def get_assets_dir(repo_root: Path | None = None) -> Path:
+    root = repo_root or get_repo_root()
+    return root / "inputs" / "diffusion_benchmark" / "figs"
+
+
+def get_output_dir(name: str, repo_root: Path | None = None) -> Path:
+    if name not in OUTPUT_DIR_NAMES:
+        raise KeyError(f"Unknown output dir name: {name}")
+    root = repo_root or get_repo_root()
+    return root / OUTPUT_DIR_NAMES[name]
+
+
+def ensure_dir(path: Path) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def check_write_access(repo_root: Path | None = None) -> Path:
+    root = repo_root or get_repo_root()
+    probe_dir = ensure_dir(root / ".cache" / "diffusion_skill_write_test")
+    probe_file = probe_dir / "probe.txt"
+    probe_file.write_text("ok", encoding="utf-8")
+    return probe_file
+
+
+def _run_nvidia_smi(query: str) -> list[list[str]]:
+    command = [
+        "nvidia-smi",
+        f"--query-{query}",
+        "--format=csv,noheader,nounits",
+    ]
+    result = subprocess.run(command, check=True, capture_output=True, text=True)
+    rows: list[list[str]] = []
+    for raw_line in result.stdout.splitlines():
+        line = raw_line.strip()
+        if not line:
+            continue
+        rows.append([field.strip() for field in csv.reader([line]).__next__()])
+    return rows
+
+
+def get_gpu_inventory() -> list[dict[str, int | str]]:
+    rows = _run_nvidia_smi("gpu=index,uuid,memory.used,memory.total,utilization.gpu")
+    inventory = []
+    for index, uuid, memory_used, memory_total, utilization_gpu in rows:
+        inventory.append(
+            {
+                "index": int(index),
+                "uuid": uuid,
+                "memory_used_mib": int(memory_used),
+                "memory_total_mib": int(memory_total),
+                "utilization_gpu_pct": int(utilization_gpu),
+            }
+        )
+    return inventory
+
+
+def get_busy_gpu_uuids() -> set[str]:
+    rows = _run_nvidia_smi("compute-apps=gpu_uuid,pid,process_name,used_gpu_memory")
+    return {gpu_uuid for gpu_uuid, *_ in rows}
+
+
+def pick_idle_gpus(
+    required_gpus: int,
+    max_memory_used_mib: int = 32,
+    max_utilization_gpu_pct: int = 5,
+) -> list[int]:
+    inventory = get_gpu_inventory()
+    busy_uuids = get_busy_gpu_uuids()
+
+    idle = [
+        int(gpu["index"])
+        for gpu in inventory
+        if gpu["uuid"] not in busy_uuids
+        and int(gpu["memory_used_mib"]) <= max_memory_used_mib
+        and int(gpu["utilization_gpu_pct"]) <= max_utilization_gpu_pct
+    ]
+    if len(idle) < required_gpus:
+        raise RuntimeError(
+            "Not enough idle GPUs. "
+            f"required={required_gpus}, idle={idle}, inventory={inventory}, busy={sorted(busy_uuids)}"
+        )
+    return idle[:required_gpus]
+
+
+def configure_runtime_env(required_gpus: int = 1) -> str | None:
+    os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+    if os.environ.get("CUDA_VISIBLE_DEVICES"):
+        return None
+    selected = ",".join(str(index) for index in pick_idle_gpus(required_gpus))
+    os.environ["CUDA_VISIBLE_DEVICES"] = selected
+    return selected
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Resolve SGLang diffusion skill paths and idle GPUs."
+    )
+    parser.add_argument(
+        "command",
+        choices=[
+            "print-root",
+            "print-assets-dir",
+            "print-output-dir",
+            "print-idle-gpus",
+            "check-write-access",
+        ],
+    )
+    parser.add_argument(
+        "--kind",
+        choices=sorted(OUTPUT_DIR_NAMES),
+        help="Output directory kind for print-output-dir.",
+    )
+    parser.add_argument(
+        "--count",
+        type=int,
+        default=1,
+        help="Number of idle GPUs to print.",
+    )
+    parser.add_argument(
+        "--mkdir",
+        action="store_true",
+        help="Create the requested directory before printing it.",
+    )
+    args = parser.parse_args()
+
+    if args.command == "print-root":
+        print(get_repo_root())
+        return
+    if args.command == "print-assets-dir":
+        path = get_assets_dir()
+        if args.mkdir:
+            ensure_dir(path)
+        print(path)
+        return
+    if args.command == "print-output-dir":
+        if not args.kind:
+            raise SystemExit("--kind is required for print-output-dir")
+        path = get_output_dir(args.kind)
+        if args.mkdir:
+            ensure_dir(path)
+        print(path)
+        return
+    if args.command == "print-idle-gpus":
+        print(",".join(str(index) for index in pick_idle_gpus(args.count)))
+        return
+    if args.command == "check-write-access":
+        print(check_write_access())
+        return
+    raise SystemExit(f"Unhandled command: {args.command}")
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md

@@ -0,0 +1,407 @@
+---
+name: sglang-diffusion-modelopt-quant
+description: Use when quantizing a diffusion DiT with NVIDIA ModelOpt and making the resulting FP8 or NVFP4 checkpoint loadable, verifiable, and benchmarkable in SGLang Diffusion.
+---
+
+# SGLang Diffusion ModelOpt Quant
+
+## Overview
+
+Use this skill when the task is to take a diffusion transformer through the full ModelOpt workflow:
+
+- quantize it with NVIDIA ModelOpt
+- adapt the exported checkpoint to SGLang Diffusion
+- verify that quality holds up
+- benchmark whether the quantized checkpoint is actually faster
+
+This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning skill.
+
+## Core Rules
+
+- Use ModelOpt's official `quantize.py` as the PTQ source of truth.
+- Keep the workflow generic. Put model-specific fallback logic in small isolated branches, not in the main conversion path.
+- Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested.
+- For diffusion FP8, keep `dit_cpu_offload=false`. `dit_layerwise_offload=true` is valid on the fixed path when you want lower DiT residency.
+- For multi-transformer pipelines, use per-component overrides when different components need different checkpoints.
+- For B200 NVFP4 validation, keep backend-sensitive environment variables explicit. Wan2.2 NVFP4 is commonly validated with `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=cudnn`; benchmark the default CUTLASS path separately if that is what you are evaluating.
+- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`, `python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`, and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere.
+- After validating a new ModelOpt quant path, update the ModelOpt support matrix in `docs/diffusion/quantization.md` before closing the task.
+
+## Read First
+
+Read these sources before changing code:
+
+- NVIDIA ModelOpt diffusers guide: `examples/diffusers/README.md`
+- ModelOpt quantization entrypoint: `examples/diffusers/quantization/quantize.py`
+- ModelOpt diffusers quant presets: `examples/diffusers/quantization/config.py`
+- SGLang diffusion quant runtime:
+  - `python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py`
+  - `python/sglang/multimodal_gen/runtime/utils/quantization_utils.py`
+  - `python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py`
+- Helper tools in this repo:
+  - [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py)
+  - [`python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`](../../../tools/build_modelopt_nvfp4_transformer.py)
+  - [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py)
+
+If you are working on a new model family, inspect the transformer's config and tensor naming before changing the generic converter.
+
+## What SGLang Supports Here
+
+This repo now contains:
+
+- flat `quant_method=modelopt` plus `quant_algo=FP8/NVFP4` resolution
+- diffusion-side ModelOpt FP8 linear loading
+- diffusion-side NVFP4 loading from ModelOpt exports
+- FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports
+- automatic protection against incompatible FP8 CPU offload while keeping layerwise DiT offload available
+- separate online diffusion quantization paths such as `--quantization fp8` / `mxfp4`; keep those out of this ModelOpt PTQ/export workflow unless the user explicitly asks for runtime quantization
+- FP8 transformer build:
+  [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py)
+- NVFP4 mixed transformer build:
+  [`python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`](../../../tools/build_modelopt_nvfp4_transformer.py)
+- trajectory similarity validation:
+  [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py)
+
+Validated documentation and CI coverage currently center on these ModelOpt diffusion transformer override families:
+
+- FP8: FLUX.1-dev, FLUX.2-dev, Wan2.2, HunyuanVideo, Qwen Image, Qwen Image Edit
+- NVFP4: FLUX.1-dev, FLUX.2-dev, Wan2.2
+
+Treat a new family, a new precision, or a new checkpoint layout as unsupported until it has a documented matrix row and a matching validation story.
+Current B200 CI also contains an Ideogram4 NVFP4 native load case
+(`ideogram4_nvfp4_t2i` via `Comfy-Org/Ideogram-4`). Treat that as source
+evidence for an existing NVFP4 path, but do not expand the ModelOpt support
+matrix to Ideogram4 unless `docs/diffusion/quantization.md` is updated with the
+exact checkpoint, loader path, quality check, and benchmark scope.
+Before writing CLI examples, re-read the active branch's `docs/diffusion/quantization.md`: FLUX.2 NVFP4 is an official `black-forest-labs/*` repo rather than a `lmsys/*` converted repo, and its preferred flag depends on the current documented loader flow. Use `--transformer-path` for a component override directory with `config.json`; use `--transformer-weights-path` when the repo or path should be probed as raw weights.
+
+B200 CI coverage can include loose BF16-vs-quantized quality checks. Inspect the active branch's `run_suite.py` before assuming they are part of the suite; mainline and feature branches may differ. Those checks are intended to catch blank, corrupted, or obviously divergent images, not exact image parity.
+
+Mainline documentation now uses `lmsys/*` for the eight converted ModelOpt
+checkpoint repos; the FLUX.2 NVFP4 raw export remains
+`black-forest-labs/FLUX.2-dev-NVFP4`. Do not use older `BBuf/*` examples unless
+you are explicitly testing a historical branch.
+
+## Related PR Watchlist
+
+These related SGLang PRs are useful as ModelOpt diffusion support history.
+Re-check the PR state and the active source tree before treating any item as
+current behavior, and keep the docs/CI matrix as the support boundary.
+
+- #23155 added Qwen Image ModelOpt FP8 support.
+- #23199 adds HunyuanVideo ModelOpt FP8 support.
+- #23373 adds a runtime quantization flag; keep PTQ/export workflows separate from runtime quant examples until the CLI behavior is merged.
+- #24024 adds transformer FP8-cast compatibility mode.
+- #24186 re-enables B200 multimodal CI with NVFP4 fixes for FLUX.2 and Wan2.2.
+
+Do not expand the validated matrix beyond the documented rows solely because a
+related PR exists. Add a row only after the exact checkpoint, loader path,
+accuracy check, and benchmark scope are validated on the active branch.
+
+## Documentation Maintenance
+
+- Keep the validated ModelOpt support matrix in `docs/diffusion/quantization.md`.
+- Each row should record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats.
+- If the quantized DiT weights are not published yet, write `unpublished` explicitly instead of leaving the field blank.
+
+## FP8 Vs NVFP4
+
+FP8 and NVFP4 are not wired into SGLang in exactly the same way.
+
+FP8:
+
+- the validated ModelOpt diffusers FP8 export still needs an extra SGLang-side conversion step
+- SGLang expects explicit `weight_scale` and `input_scale`
+- the validated path also materializes SGLang-native `float8_e4m3fn` weights from `backbone.pt`
+
+NVFP4:
+
+- the official diffusers export often already contains packed FP4 weights, scale tensors, and enough safetensors metadata for SGLang to rebuild the quant config
+- in that case SGLang mainly needs to detect the checkpoint family and rearrange tensors into the runtime layout
+- this is why NVFP4 often does not need an extra offline conversion pass like FP8 does
+- backend choice matters on B200; record whether the run used the default CUTLASS path or a cuDNN-backed FlashInfer FP4 GEMM path
+
+Important caveat:
+
+- "often" does not mean "always"
+- the exact load path still depends on the checkpoint family, especially whether a model uses a packed-QKV layout
+
+## Generic Workflow
+
+### 1. Verify The BF16 Baseline First
+
+Before quantizing anything:
+
+- run the original BF16 model in SGLang
+- fix the prompt, seed, size, step count, and GPU topology
+- save the output and `perf.json`
+
+Do not start quantization work until the BF16 path is already healthy.
+
+### 2. Quantize With Official ModelOpt
+
+Use ModelOpt's official script. Generic template:
+
+```bash
+python quantize.py \
+  --model <model-name> \
+  --override-model-path <hf-repo-or-local-model> \
+  --model-dtype <Half|BFloat16> \
+  --format <fp8|fp4> \
+  --batch-size 1 \
+  --calib-size <calib-size> \
+  --n-steps <calib-steps> \
+  --quantize-mha \
+  --prompts-file <prompt-file> \
+  --quantized-torch-ckpt-save-path <out>/ckpt \
+  --hf-ckpt-dir <out>/hf
+```
+
+For current ModelOpt diffusion examples, use `--format fp4` for NVFP4 exports.
+Do not assume the checked-out ModelOpt version accepts a literal `nvfp4` format string unless you verified it locally.
+
+For multi-transformer models:
+
+- quantize each backbone deliberately
+- keep each output directory separate
+- save both `backbone.pt` and the matching `hf/<component>` export
+
+### 3. Convert FP8 Exports For SGLang
+
+FP8 requires an extra conversion step:
+
+```bash
+PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \
+  --modelopt-hf-dir <out>/hf \
+  --modelopt-backbone-ckpt <out>/ckpt/backbone.pt \
+  --base-transformer-dir <base-model-transformer-dir> \
+  --output-dir <out>/sglang_transformer \
+  --overwrite
+```
+
+What the converter does:
+
+- reads `weight_quantizer._amax` and `input_quantizer._amax` from `backbone.pt`
+- writes `weight_scale` and `input_scale`
+- materializes eligible FP8 weights as `float8_e4m3fn`
+- preserves ModelOpt `ignore` layers as BF16
+- strips stale `_quantizer.*` tensors and fallback-layer scales that should not survive into the SGLang-native checkpoint
+
+For `FLUX.1-dev`, the validated fallback set currently keeps these modules in BF16:
+
+- `transformer_blocks.*.norm1.linear`
+- `transformer_blocks.*.norm1_context.linear`
+- `transformer_blocks.*.ff.net.0.proj`
+- `transformer_blocks.*.ff.net.2`
+- `transformer_blocks.*.ff_context.net.0.proj`
+- `transformer_blocks.*.ff_context.net.2`
+- `single_transformer_blocks.*.norm.linear`
+- `single_transformer_blocks.*.proj_mlp`
+
+Use `--model-type flux1` to force that profile, or rely on `--model-type auto` when the export config identifies `FluxTransformer2DModel`.
+
+HunyuanVideo uses `HunyuanVideoTransformer3DModel`, so the validated
+HunyuanVideo FP8 fallback preset keeps these modules in BF16:
+
+- `context_embedder.*`
+- `x_embedder.proj`
+- `time_text_embed.(timestep_embedder|guidance_embedder|text_embedder).linear_[12]`
+- `norm_out.linear`
+- `proj_out`
+- `transformer_blocks.*.norm1.linear`
+- `transformer_blocks.*.norm1_context.linear`
+- `single_transformer_blocks.*.norm.linear`
+
+Use `--model-type hunyuan-video` to force that profile, or rely on
+`--model-type auto` when the export config identifies
+`HunyuanVideoTransformer3DModel`.
+
+HunyuanVideo ModelOpt exports use diffusers module names that differ from
+SGLang runtime names for fused QKV and fused QKV+MLP layers. Keep the
+diffusers-to-runtime mapping in `build_modelopt_fp8_transformer.py` in sync
+with `runtime/models/dits/hunyuanvideo.py` before trusting converted scale
+tensors.
+
+Qwen Image and Qwen Image Edit share `QwenImageTransformer2DModel`, so one
+ModelOpt FP8 fallback preset covers both. The validated Qwen Image fallback set
+keeps these modules in BF16:
+
+- `img_in`
+- `txt_in`
+- `time_text_embed.timestep_embedder.linear_1`
+- `time_text_embed.timestep_embedder.linear_2`
+- `norm_out.linear`
+- `proj_out`
+- `transformer_blocks.*.img_mlp.net.2`
+- `transformer_blocks.*.img_mod`
+- `transformer_blocks.*.txt_mod`
+
+Use `--model-type qwen-image` to force that profile, or rely on
+`--model-type auto` when the export config identifies
+`QwenImageTransformer2DModel`.
+
+Qwen modulation weights can appear in safetensors as `.img_mod.1.weight` and
+`.txt_mod.1.weight`. Canonicalize those module names to `.img_mod` and
+`.txt_mod` before fallback matching.
+
+For Qwen Image FP8, explicit BF16 fallback tensors must be written before
+honoring ModelOpt ignored weights. Otherwise converter stats can report a
+fallback while the output checkpoint still retains the source FP8 tensor, which
+causes severe image-quality regressions.
+
+For FLUX.1-dev NVFP4 model families that need a mixed BF16+NVFP4 checkpoint, build the merged transformer explicitly:
+
+```bash
+PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer \
+  --base-transformer-dir <base-model-transformer-dir> \
+  --modelopt-hf-dir <out>/hf/transformer \
+  --output-dir <out>/transformer-mixed \
+  --pattern-preset flux1-nvfp4
+```
+
+The validated FLUX.1-dev mixed builder also needs to preserve:
+
+- `quant_type: NVFP4` in `config.json`
+- `swap_weight_nibbles: false` for the validated diffusers export
+
+### 4. Load The Quantized Checkpoint In SGLang
+
+Single-transformer example:
+
+```bash
+sglang generate \
+  --model-path <base-model> \
+  --transformer-path <quantized-transformer> \
+  --prompt "<prompt>" \
+  --seed <seed> \
+  --save-output
+```
+
+Multi-transformer example:
+
+```bash
+sglang generate \
+  --model-path <base-model> \
+  --transformer-path <quantized-transformer> \
+  --transformer-2-path <another-transformer-or-bf16-override> \
+  --prompt "<prompt>" \
+  --seed <seed> \
+  --save-output
+```
+
+Guideline:
+
+- use the global `--transformer-path` only when the model effectively has one transformer override to apply
+- use per-component overrides when different backbones need different checkpoints
+- the preferred CLI form is `--<component>-path`
+- config-expanded forms such as `--component_paths.transformer_2=...` also resolve to the same internal override map
+
+### 5. Validate Accuracy
+
+Use two levels of validation.
+
+Reduced deterministic validation:
+
+- keep prompt, seed, resolution, and step count fixed
+- compare BF16 and quantized runs
+- capture denoising trajectories
+- inspect per-step latent cosine similarity plus MAE or RMSE
+- compare final frames with image metrics such as PSNR or MAE
+
+Tool:
+
+```bash
+PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.compare_diffusion_trajectory_similarity \
+  --model-path <base-model> \
+  --model-id <optional-native-model-id> \
+  --prompt "<prompt>" \
+  --width <w> \
+  --height <h> \
+  --num-inference-steps <steps> \
+  --guidance-scale <cfg> \
+  --seed <seed> \
+  --candidate-transformer-path <quantized-transformer> \
+  --output-json <report.json>
+```
+
+Use `--model-id FLUX.1-dev` when `--model-path` points to a local directory but the runtime still needs the native FLUX.1 model registration.
+
+Full-output validation:
+
+- run the same user-facing generation config in BF16 and quantized mode
+- inspect the output visually
+- only claim "quality preserved" for the exact scope you actually checked
+
+### 6. Benchmark Correctly
+
+Benchmark only when these match between BF16 and quantized:
+
+- prompt
+- seed
+- width and height
+- frame count
+- inference step count
+- GPU count and topology
+- offload flags
+- compile settings
+- profiler settings
+
+Only the quantized checkpoint path should differ.
+
+Interpretation rule:
+
+- the primary expected gain is in denoising
+- text-encoding and VAE differences are secondary and should not be over-attributed unless they were quantized too
+
+### 7. Add Model-Specific Fallbacks Only When Needed
+
+If the generic FP8 path fails on a new model family:
+
+- inspect which modules are numerically sensitive or loader-incompatible
+- keep fallback patterns small and explicit
+- isolate them in the converter instead of scattering ad-hoc exceptions
+- re-run deterministic trajectory checks after every fallback change
+
+Do not turn one validated model quirk into a generic rule unless another family also needs it.
+
+## FP8 Offload Constraint
+
+Current diffusion ModelOpt FP8 support requires:
+
+- `dit_cpu_offload=false`
+- `dit_layerwise_offload` may be enabled when you want lower DiT residency
+
+Reason:
+
+- the FP8 linear path depends on a CUTLASS-compatible weight layout after loading
+- `dit_cpu_offload` is still treated conservatively
+- the fixed layerwise offload path now preserves non-contiguous tensor strides instead of flattening and rebuilding FP8 weights into a contiguous layout
+
+Runtime behavior:
+
+- SGLang still force-disables `dit_cpu_offload` when it detects `modelopt_fp8`
+- benchmark commands should still pin offload flags explicitly so the command line itself makes the comparison rule obvious
+
+## Claim Discipline
+
+When documenting results:
+
+- claim only scopes that were actually validated end to end
+- do not collapse "single-transformer FP8 override" into "full-model FP8"
+- do not call a practical deployment comparison a benchmark if BF16 and quantized commands used different offload behavior
+
+## Current Code Areas
+
+| File | Role |
+| --- | --- |
+| `runtime/layers/quantization/__init__.py` | registers diffusion quant methods |
+| `runtime/layers/quantization/modelopt_quant.py` | ModelOpt FP8 and NVFP4 runtime loading |
+| `runtime/utils/quantization_utils.py` | resolves flat ModelOpt configs and reconstructs NVFP4 config from metadata |
+| `runtime/loader/transformer_load_utils.py` | guards incompatible FP8 offload modes |
+| `runtime/models/dits/flux_2.py` | packed-QKV handling for the packed FLUX.2 NVFP4 family |
+| `tools/build_modelopt_fp8_transformer.py` | Build an SGLang-loadable FP8 transformer from a ModelOpt export |
+| `tools/build_modelopt_nvfp4_transformer.py` | Build mixed BF16+NVFP4 transformer directories when a family needs preserved BF16 layers |
+| `tools/compare_diffusion_trajectory_similarity.py` | reduced deterministic BF16-vs-quantized validation |
+| `docs/diffusion/quantization.md` | public ModelOpt support matrix and CLI examples |
+| `test/server/testcase_configs.py` | reusable ModelOpt testcase constants, thresholds, and helpers |
+| `test/server/gpu_cases.py` | concrete GPU and B200 ModelOpt CI case lists |

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md

@@ -0,0 +1,290 @@
+---
+name: sglang-diffusion-performance
+description: Use when choosing the fastest SGLang Diffusion flags for a model, GPU, and VRAM budget.
+---
+
+# SGLang Diffusion Performance Tuning
+
+Use this skill when the user wants the fastest command line, lower VRAM, or the right performance flags for a specific model and GPU setup.
+
+Before running any `sglang generate` command below inside the diffusion container:
+- use `python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py` to derive the repo root, verify write access, and choose idle GPU(s)
+- export `HF_TOKEN` first when the selected model lives in a gated Hugging Face repo such as `black-forest-labs/FLUX.*`
+- export `FLASHINFER_DISABLE_VERSION_CHECK=1`
+- `cd` to the repo root resolved from `sglang.__file__`
+
+## Native Backend Gate
+
+Performance numbers are useful only when the intended backend actually ran.
+
+- Treat any log containing `Falling back to diffusers backend`, `Using diffusers backend`, or `Loaded diffusers pipeline` as invalid for native SGLang performance tuning.
+- Use `--backend diffusers` only for an explicit diffusers baseline. For native recipes, leave the default backend or pin `--backend sglang`.
+- If a fallback happened, fix pipeline registration/model-path/config issues first, then rerun. Do not compare perf dumps collected from a fallback run.
+- When the runtime auto-selects parallel settings because the user omitted them, keep the result as an auto-tuned baseline. For reproducible tuning, pin `--num-gpus`, `--ulysses-degree`, `--ring-degree`, and `--enable-cfg-parallel` explicitly.
+
+Reference: [SGLang-Diffusion Advanced Optimizations Blog](https://lmsys.org/blog/2026-02-16-sglang-diffusion-advanced-optimizations/)
+
+---
+
+## Section 1: Lossless Optimizations
+
+These options are intended to preserve output quality. In practice, some paths (most notably `torch.compile`) can still introduce small floating-point drift, so validate on your target model when numerical parity matters.
+
+| Option | CLI Flag / Env Var | What It Does | Speedup | Limitations / Notes |
+|---|---|---|---|---|
+| **torch.compile** | `--enable-torch-compile` | Applies `torch.compile` to the DiT forward pass, fusing ops and reducing kernel launch overhead. | ~1.2–1.5x on denoising | First request is slow (compilation). May cause minor precision drifts due to [PyTorch issue #145213](https://github.com/pytorch/pytorch/issues/145213). Pair with `--warmup` for best results. |
+| **Warmup** | `--warmup` | Runs dummy forward passes to warm up CUDA caches, JIT, and `torch.compile`. Eliminates cold-start penalty. | Removes first-request latency spike | Adds startup time. Without `--warmup-resolutions`, warmup happens on first request. |
+| **Warmup Resolutions** | `--warmup-resolutions 256x256 720x720` | Pre-compiles and warms up specific resolutions at server startup (instead of lazily on first request). | Faster first request per resolution | Each resolution adds to startup time. Serving mode only; useful when you know your target resolutions in advance. |
+| **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. For Wan2.2 video, start by benchmarking pure Ulysses before assuming a mixed Ulysses/Ring layout is fastest. |
+| **CFG Parallel** | `--enable-cfg-parallel` | Runs conditional and unconditional CFG branches in parallel across GPUs. For CFG models on multi-GPU, benchmark this against pure Ulysses on your topology instead of assuming one always wins. | Often faster than pure SP for CFG models | Requires `num_gpus >= 2`. Halves the Ulysses group size (e.g. 8 GPU → two 4-GPU groups). Only for models that use CFG. Nightly coverage configs may intentionally use smaller Ulysses groups to keep ring behavior exercised; that does not automatically make them the lowest-latency choice. |
+| **Layerwise Offload** | `--dit-layerwise-offload` | Async layer-by-layer H2D prefetch with compute overlap. Only ~2 DiT layers reside on GPU at a time, dramatically reducing VRAM. For some video models the copy stream can be almost fully hidden behind compute. | Saves VRAM (40 GB → ~11 GB for Wan A14B); can be near-zero speed cost on the right workload | Enabled by default for Wan/MOVA video models. Incompatible with Cache-DiT. For **image models** or highly parallelized setups (many GPUs, small per-GPU compute), the copy stream may not be fully hidden and can cause slowdown. |
+| **Offload Prefetch Size** | `--dit-offload-prefetch-size F` | Fine-grained control over layerwise offload: how many layers to prefetch ahead. `0.0` = 1 layer (min VRAM), `0.1` = 10% of layers, `≥1` = absolute layer count. | Tune for cases where default offload has copy stream interference (e.g. image models). 0.05–0.1 is a good starting point. | Values ≥ 0.5 approach no-offload VRAM with worse performance. Use lower values when copy overlap is weak; disable offload when memory allows and latency dominates. |
+| **FSDP Inference** | `--use-fsdp-inference` | Uses PyTorch FSDP to shard model weights across GPUs with prefetch. Low latency, low VRAM. | Reduces per-GPU VRAM | Mutually exclusive with `--dit-layerwise-offload`. More overhead than SP on high-bandwidth interconnects. |
+| **CPU Offload (components)** | `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, `--dit-cpu-offload` | Offloads specific pipeline components to CPU when not in use. | Reduces peak VRAM | Adds H2D transfer latency when the component is needed. Auto-enabled for low-VRAM GPUs (<30 GB). **Tip:** after the first request completes, the console prints a peak VRAM analysis with suggestions on which offload flags can be safely disabled — look for the `"Components that could stay resident"` log line. |
+| **Pin CPU Memory** | `--pin-cpu-memory` | Uses pinned (page-locked) memory for CPU offload transfers. | Faster H2D transfers | Slightly higher host memory usage. Enabled by default; disable only as workaround for CUDA errors. |
+| **Attention Backend (lossless)** | `--attention-backend fa` | Selects a lossless attention kernel for SGLang-native pipelines: `fa` (FlashAttention 2/3/4 alias) or `torch_sdpa`. | FA is usually faster than SDPA on long sequences | FA requires compatible GPU (Ampere+). For `--backend diffusers`, valid backend names differ; use the names documented in `docs/diffusion/performance/attention_backends.md`. |
+| **Parallel Folding** | *(automatic when SP > 1)* | Reuses the SP process group as TP for the T5 text encoder, so text encoding is parallelized "for free". | Faster text encoding on multi-GPU | Automatic; no user action needed. Only applies to T5-based pipelines. |
+
+---
+
+## Section 2: Lossy Optimizations
+
+These options **trade output quality** for speed or VRAM savings. Results will differ from the baseline.
+
+| Option | CLI Flag / Env Var | What It Does | Speedup | Quality Impact / Limitations |
+|---|---|---|---|---|
+| **Approximate Attention** | `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn` | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific and may require config files (e.g. `--mask-strategy-file-path` for STA). Requires corresponding packages installed. |
+| **Cache-DiT** | Native: `SGLANG_CACHE_DIT_ENABLED=true` plus `SGLANG_CACHE_DIT_*` env vars. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Incompatible with `--dit-layerwise-offload`. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
+| **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path <path>` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. |
+| **Pre-quantized Transformer Override** | `--transformer-path <dir-or-repo>` / `--transformer-weights-path <path>` | Load a quantized transformer component or raw transformer weights. For converted ModelOpt FP8/NVFP4 directories, prefer `--transformer-path`; use `--transformer-weights-path` for weight-only artifacts the model loader expects. | ~1.3–1.5x compute (dtype dependent) | Requires a validated quantized transformer override, such as one produced by the ModelOpt helper tools. Quality is usually slightly worse than BF16 and depends on the format, fallback layers, and calibration scope. |
+| **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. |
+| **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. |
+
+---
+
+## Quick Recipes
+
+### Maximum speed, video model, multi-GPU, lossless (Wan A14B, 8 GPUs)
+
+```bash
+sglang generate --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --num-gpus 8 --enable-cfg-parallel --ulysses-degree 4 \
+  --enable-torch-compile --warmup \
+  --text-encoder-cpu-offload true \
+  --prompt "..." --save-output
+```
+
+Note: `--dit-layerwise-offload` is enabled by default for Wan/MOVA video models and is often a good default, but still benchmark it on your exact workload if latency matters.
+
+For Wan2.2 specifically:
+- the nightly-aligned 4-GPU benchmark may use `--enable-cfg-parallel --ulysses-degree=2` to keep CFG and ring behavior covered
+- that is a **coverage** choice, not a guaranteed best-performance choice
+- for pure latency tuning, benchmark pure Ulysses too, for example `--ulysses-degree=4 --ring-degree=1` on 4 GPUs
+- on 8 GPUs, compare pure `--ulysses-degree=8` against `--enable-cfg-parallel --ulysses-degree=4`
+
+### Nightly-aligned model, 2 GPUs: LTX-2 two-stage
+
+```bash
+sglang generate --model-path Lightricks/LTX-2 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --prompt "A cat and a dog baking a cake together in a kitchen." \
+  --width 768 --height 512 \
+  --num-frames 121 \
+  --seed 42 --num-gpus 2 --enable-cfg-parallel \
+  --enable-torch-compile --warmup --save-output
+```
+
+Note: this generate recipe is aligned with the nightly comparison case `ltx2_twostage_t2v`. The nightly config omits explicit steps and guidance, so this command omits them too and uses runtime defaults. `LTX2TwoStagePipeline` is a native path and auto-resolves the spatial upsampler plus distilled LoRA from the same model snapshot unless you override them.
+
+### Nightly-aligned model, 2 GPUs: LTX-2.3 TI2V two-stage
+
+```bash
+sglang generate --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --prompt "The cat starts walking slowly towards the camera." \
+  --image-path "${ASSET_DIR}/cat.png" \
+  --width 768 --height 512 \
+  --num-frames 121 \
+  --seed 42 --num-gpus 2 --cfg-parallel-size 2 \
+  --enable-torch-compile --warmup --save-output
+```
+
+Note: this matches the nightly comparison case `ltx2.3_twostage_ti2v_2gpus`. The nightly config omits explicit steps and guidance, so this command omits them too and uses runtime defaults. Download `${ASSET_DIR}/cat.png` with the benchmark/profile skill before running it.
+
+### Native baseline, 2 GPUs: LTX-2.3 one-stage
+
+```bash
+sglang generate --model-path Lightricks/LTX-2.3 \
+  --prompt "A beautiful sunset over the ocean" \
+  --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
+  --width 768 --height 512 \
+  --num-frames 121 --fps 24 \
+  --num-inference-steps 30 --guidance-scale 3.0 \
+  --seed 1234 --num-gpus 2 \
+  --enable-torch-compile --warmup --save-output
+```
+
+Note: use this as the native `LTX2Pipeline` baseline for `LTX-2.3`. It keeps the validated one-stage resolution and explicit `LTX-2.3` sampling defaults, and matches the `ltx23-one-stage` benchmark preset in `sglang-diffusion-benchmark-profile`.
+
+### Skill-only stress target, 2 GPUs: LTX-2.3 two-stage high resolution
+
+```bash
+sglang generate --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --prompt "A beautiful sunset over the ocean" \
+  --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \
+  --width 1536 --height 1024 \
+  --num-frames 121 --fps 24 \
+  --num-inference-steps 30 --guidance-scale 3.0 \
+  --seed 1234 --num-gpus 2 \
+  --enable-torch-compile --warmup --save-output
+```
+
+Note: this is a high-resolution stress target for the native `LTX-2.3` two-stage path. It matches the skill-only `ltx23-two-stage` benchmark preset, not a nightly comparison case.
+
+### Maximum speed, image model, single GPU, lossless
+
+```bash
+sglang generate --model-path <IMAGE_MODEL> \
+  --enable-torch-compile --warmup \
+  --dit-layerwise-offload false \
+  --dit-cpu-offload false \
+  --prompt "..." --save-output
+```
+
+Note: for image models, per-layer compute is smaller, so layerwise offload may not fully hide H2D transfer. Disable DiT layerwise and CPU offload if VRAM allows; otherwise a large image DiT can stay resident on CPU and make the denoise loop H2D-bound.
+
+### Image-edit baselines: JoyAI and FireRed
+
+```bash
+sglang generate --backend=sglang \
+  --model-path jdopensource/JoyAI-Image-Edit-Diffusers \
+  --prompt "Make the cat wear a red hat" \
+  --image-path "${ASSET_DIR}/cat.png" \
+  --width 1024 --height 1024 \
+  --num-inference-steps 40 --guidance-scale 4.0 \
+  --num-gpus 2 --enable-cfg-parallel --ulysses-degree 1 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --enable-torch-compile --warmup --save-output
+```
+
+```bash
+sglang generate --backend=sglang \
+  --model-path FireRedTeam/FireRed-Image-Edit-1.1 \
+  --prompt "Make the cat wear a red hat" \
+  --image-path "${ASSET_DIR}/cat.png" \
+  --width 1024 --height 1024 \
+  --num-inference-steps 40 --guidance-scale 4.0 \
+  --num-gpus 2 --enable-cfg-parallel --ulysses-degree 1 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --enable-torch-compile --warmup --save-output
+```
+
+Use `FireRedTeam/FireRed-Image-Edit-1.0` in the same command when comparing
+FireRed 1.0. These are native image-edit paths; keep the reference image, prompt,
+seed, and output size fixed when comparing denoise numbers. On H100, 2-GPU CFG
+parallel was faster than the otherwise matching 2-GPU Ulysses command: FireRed
+1.0 improved from 13419.15 ms to 10955.90 ms, and FireRed 1.1 improved from
+13414.72 ms to 10934.21 ms.
+
+### Hunyuan3D shape baseline
+
+```bash
+OUTPUT_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir)
+CONFIG_DIR="${OUTPUT_DIR}/generated_configs"
+mkdir -p "${CONFIG_DIR}"
+printf '{"paint_enable": false}\n' > "${CONFIG_DIR}/hunyuan3d-shape.json"
+
+sglang generate --backend=sglang \
+  --model-path tencent/Hunyuan3D-2 \
+  --prompt "generate 3d mesh" \
+  --image-path "${ASSET_DIR}/cat.png" \
+  --config "${CONFIG_DIR}/hunyuan3d-shape.json" \
+  --num-inference-steps 50 --guidance-scale 5.0 \
+  --dit-layerwise-offload false --dit-cpu-offload false \
+  --enable-torch-compile --warmup --save-output
+```
+
+For Hunyuan3D, treat `Hunyuan3DShapeDenoisingStage` as the primary latency
+metric. Mesh export and paint stages are useful end-to-end checks but should not
+drive DiT optimization decisions.
+
+### Low VRAM, decent speed (single GPU)
+
+```bash
+sglang generate --model-path <MODEL> \
+  --enable-torch-compile --warmup \
+  --dit-layerwise-offload --dit-offload-prefetch-size 0.1 \
+  --text-encoder-cpu-offload true --vae-cpu-offload true \
+  --prompt "..." --save-output
+```
+
+### Maximum speed, lossy native path (SageAttention + Cache-DiT)
+
+```bash
+SGLANG_CACHE_DIT_ENABLED=true sglang generate --model-path <MODEL> \
+  --attention-backend sage_attn \
+  --dit-layerwise-offload false \
+  --enable-torch-compile --warmup \
+  --prompt "..." --save-output
+```
+
+Add native Cache-DiT knobs such as `SGLANG_CACHE_DIT_SCM_PRESET=medium`,
+`SGLANG_CACHE_DIT_RDT=0.24`, or `SGLANG_CACHE_DIT_TAYLORSEER=true` only after
+you have a BF16 baseline output to compare against.
+
+For a diffusers-backend Cache-DiT YAML/JSON config baseline, make the fallback
+explicit:
+
+```bash
+sglang generate --backend diffusers --model-path <MODEL> \
+  --cache-dit-config <config.yaml> \
+  --dit-layerwise-offload false \
+  --prompt "..." --save-output
+```
+
+---
+
+## Model-Specific Starting Points
+
+Use these as first commands to benchmark, not as universal winners.
+
+| Model family | First performance shape | Starting flags | Notes |
+|---|---|---|---|
+| FLUX.1 / FLUX.2 image | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup --dit-layerwise-offload false` | `black-forest-labs/FLUX.*` repos are gated; for FP8/NVFP4 use validated `--transformer-path` or `--transformer-weights-path` flows from the quant skill. |
+| FLUX.2 Klein / Klein Base | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup --dit-layerwise-offload false` | Current registry has `black-forest-labs/FLUX.2-klein-4B`, `FLUX.2-klein-9B`, and base variants. Klein is step-distilled; Klein Base is not. |
+| Qwen-Image / Qwen-Image-Edit | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup`; optionally native `SGLANG_CACHE_DIT_ENABLED=true` | Cache-DiT is lossy. For edit tasks, keep reference image, seed, and output size fixed. |
+| Z-Image / Z-Image-Turbo | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup` | Keep base Z-Image separate from Turbo: base uses 50-step CFG defaults, Turbo uses 9-step zero-CFG defaults. Mainline has Z-Image tanh/gate norm fusions. |
+| Wan2.2 A14B T2V/I2V | 1280x720, 81 frames | Nightly: `--num-gpus 4 --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory` | For lowest latency, also benchmark pure Ulysses on the same GPUs. |
+| Wan2.2 TI2V 5B | 1280x720, 81 frames, 1 GPU | `--enable-torch-compile --warmup` | Keep the input image and motion prompt fixed when comparing sparse attention or Cache-DiT. |
+| Wan2.1 / FastWan / TurboWan variants | 480p or 720p video, family defaults | `--enable-torch-compile --warmup`; add `--ulysses-degree` / CFG parallel only after measuring | Current registry includes Wan2.1, FastWan2.1, FastWan2.2 TI2V, TurboWan2.1, TurboWan2.2 I2V, and Wan2.1-Fun InP. Use the compatibility matrix and benchmark presets before choosing topology. |
+| Cosmos3 Nano / Super | T2I: 1024x1024 with `--num-frames 1`; T2V/I2V: 480p/720p video | `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` for benchmark isolation; `--enable-torch-compile --warmup` | One checkpoint serves T2I/T2V/I2V. Mode is request-driven: `num_frames == 1` means T2I, `--image-path` means I2V. |
+| Ideogram 4 FP8/NVFP4 | 1024x1024, native preset defaults | `--enable-torch-compile --warmup` | Do not set `--num-inference-steps` or `--guidance-scale` directly unless you also update the Ideogram preset; sampling params derive them from `preset`. |
+| ERNIE-Image / GLM-Image / SANA / SD3 | 1024-class image, family defaults | `--enable-torch-compile --warmup`; disable offload only after checking VRAM | Treat these as current native image families. Start with benchmark/profile presets for ERNIE, GLM, and SANA; use registry/config defaults for SD3 unless you add a new preset. |
+| LTX-2 / LTX-2.3 | 768x512 or HQ 1920x1088, 121 frames | `--pipeline-class-name LTX2TwoStagePipeline --enable-torch-compile --warmup`; HQ uses `LTX2TwoStageHQPipeline` | Use benchmark/profile presets for nightly alignment, one-stage, high-resolution stress, and HQ. Device mode choices are `original` and `resident`; `resident` is fastest but uses more VRAM. `snapshot` is a deprecated alias for `original`, so do not use it in new commands. |
+| HunyuanVideo | 848x480 or 720p class video | `--text-encoder-cpu-offload --pin-cpu-memory --enable-torch-compile --warmup` | Check VAE decode separately. GroupNorm+SiLU is default-eligible in mainline when wrapper guards pass; use `bench_group_norm_silu.py` when VAE residual blocks are hot. |
+| JoyAI-Image-Edit | 1024-class TI2I, 40 steps, guidance 4.0 | `--backend=sglang --num-gpus 2 --enable-cfg-parallel --ulysses-degree 1 --enable-torch-compile --warmup --dit-layerwise-offload false --dit-cpu-offload false` | Newly supported image-edit path. Keep the input image, prompt, seed, and output size fixed; 2-GPU CFG parallel is the validated H100 starting point. |
+| FireRed-Image-Edit 1.0 / 1.1 | 1024x1024 image edit, 40 steps, guidance 4.0 | `--backend=sglang --num-gpus 2 --enable-cfg-parallel --ulysses-degree 1 --enable-torch-compile --warmup --dit-layerwise-offload false --dit-cpu-offload false` | Uses the native `QwenImageEditPlusPipeline` path. 2-GPU CFG parallel is the validated H100 starting point; benchmark 1.0 and 1.1 separately because checkpoint differences can change denoise latency. |
+| Hunyuan3D-2 shape | Shape generation, 50 steps, guidance 5.0 | `--backend=sglang --enable-torch-compile --warmup --dit-layerwise-offload false --dit-cpu-offload false` | Focus on `Hunyuan3DShapeDenoisingStage`; keep mesh export/paint timings separate from denoise. |
+| MOVA / Helios / LingBot World | Use the benchmark/profile presets or server test cases first | `--enable-torch-compile --warmup`; pin offload and topology flags explicitly | These video/realtime families have model-specific stages and condition handling. Keep prompt/image/action inputs fixed and prefer perf dumps over wall time alone. |
+
+## Historical PR Watchlist
+
+Treat these performance PRs as direction and prior art only. Re-check the PR
+state and the active source tree before relying on any path, flag, or claim
+about whether the work has merged:
+
+- Fusion/kernel: #24025 LTX2 QK norm, #24059 Helios norm modulation, #24117 Z-Image packed QKV, #19488 Wan elementwise cross-block fusion, #19249 Z-Image gate/norm fusion, #20429 Qwen-Image layernorm/modulation, #20530 MOVA RMSNorm+RoPE.
+- VAE/decode: #22531 LTX2 parallel VAE, #20927 batched tiled VAE decode.
+- Runtime/parallel/cache: #22805 FLUX.2 packed QKV for A2A, #21742 hybrid attention schedule, #24053 USP replicated-prefix fix, #21613 TeaCache refactor, #24227 WanVideo TeaCache fix, #18764 dynamic batching, #24200 disaggregated diffusion.
+
+## Tips
+
+- **Benchmarking**: always use `--warmup` and look for the line ending with `(with warmup excluded)` for accurate timing.
+- **Perf dump**: use `--perf-dump-path result.json` to save structured metrics, then compare with `python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json result.json`.
+- **Offload tuning**: after the first request, the runtime logs peak GPU memory and which components could stay resident. Use this to decide which `--*-cpu-offload` flags to disable.
+- **Backend selection**: `--backend sglang` (default, auto-detected) enables native optimizations (fused kernels, SP, native Cache-DiT env knobs, etc.). `--backend diffusers` falls back to Diffusers pipelines and is the path that accepts `--cache-dit-config` plus diffusers attention backend names.
+- **Wan2.2-I2V sizing**: explicit `--width/--height` on `Wan2.2-I2V-A14B` control the target area while preserving the condition-image aspect ratio.
+- **Mainline diffusion fast paths**: before proposing a new kernel or overlap scheme, check `sglang-diffusion-benchmark-profile/existing-fast-paths.md`. It covers GroupNorm+SiLU, Z-Image residual-form modulation, fused diffusion `QK norm + RoPE`, LTX2 split RoPE, LTX2 residual-gate add, varlen USP pack/scatter, packed QKV/NVFP4 expectations, and existing multi-GPU overlap families such as Ulysses / USP and turbo-layer async all-to-all.
+- **NVFP4 trace interpretation**: on FLUX.2 NVFP4 and Nunchaku-style checkpoints, packed QKV is expected. SGLang intentionally uses fused projection modules such as `to_qkv` / `to_added_qkv` instead of separate `to_q` / `to_k` / `to_v`, so a split-QKV trace usually means the quantized path did not engage rather than a brand new fusion opportunity.
+- **Hotspot workflow split**: use `sglang-diffusion-benchmark-profile` to prove and classify a slowdown with perf dumps plus `torch.profiler`; hand concrete kernel work off with the perf/profile evidence attached instead of expanding the benchmark skill.

python/sglang/multimodal_gen/README.md

@@ -0,0 +1,105 @@
+<div align="center"  style="display:block; margin:auto;">
+<img src=https://github.com/lm-sys/lm-sys.github.io/releases/download/test/sgl-diffusion-logo.png width="80%"/>
+</div>
+
+**SGLang diffusion is an inference framework for accelerated image/video generation.**
+
+SGLang diffusion features an end-to-end unified pipeline for accelerating diffusion models. It is designed to be modular and extensible, allowing users to easily add new models and optimizations.
+
+## Key Features
+
+SGLang Diffusion has the following features:
+  - Broad model support: Wan, FastWan, FLUX, Qwen-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3, LingBot World, SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
+  - Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations
+  - Ease of use: OpenAI-compatible api, CLI, and python sdk support
+  - Multi-platform support:
+    - NVIDIA GPUs (H100, H200, A100, B200, 4090, 5090)
+    - AMD GPUs (MI300X, MI325X, MI355X)
+    - Intel XPUs
+    - Ascend NPU (A2, A3)
+    - Apple Silicon (M-series via MPS)
+    - Moore Threads GPUs (MTT S5000)
+
+### AMD/ROCm Support
+
+SGLang Diffusion supports AMD Instinct GPUs through ROCm. On AMD platforms, we use the Triton attention backend and leverage AITER kernels for optimized layernorm and other operations. See the [installation guide](https://docs.sglang.io/docs/sglang-diffusion/installation) for setup instructions.
+
+### Moore Threads/MUSA Support
+
+SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software stack. On MUSA platforms, we use FlashAttention (FA3) when available; also supports Sage Attention when installed; otherwise falls back to the Torch SDPA backend. See the [installation guide](https://docs.sglang.io/docs/sglang-diffusion/installation) for setup instructions.
+
+### Apple MPS Support
+
+SGLang Diffusion supports Apple Silicon (M-series) via the MPS backend. Since Triton is Linux-only, all Triton kernels are replaced with PyTorch-native fallbacks on MPS. Norm operations can be optionally accelerated with MLX fused Metal kernels (`SGLANG_USE_MLX=1`). See the [installation guide](https://docs.sglang.io/docs/sglang-diffusion/installation) for setup instructions.
+
+## Getting Started
+
+```bash
+uv pip install 'sglang[diffusion]' --prerelease=allow
+```
+
+For more installation methods (e.g. pypi, uv, docker, ROCm/AMD, MUSA/Moore Threads), check the [installation guide](https://docs.sglang.io/docs/sglang-diffusion/installation).
+
+## Inference
+
+Here's a minimal example to generate a video using the default settings:
+
+```python
+from sglang.multimodal_gen import DiffGenerator
+
+def main():
+    # Create a diff generator from a pre-trained model
+    generator = DiffGenerator.from_pretrained(
+        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Generate the video
+    video = generator.generate(
+        sampling_params_kwargs=dict(
+            prompt="A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest.",
+            return_frames=True,  # Also return frames from this call (defaults to False)
+            output_path="my_videos/",  # Controls where videos are saved
+            save_output=True
+        )
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Or, more simply, with the CLI:
+
+```bash
+sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --text-encoder-cpu-offload --pin-cpu-memory \
+    --prompt "A curious raccoon" \
+    --save-output
+```
+
+### LoRA support
+
+Apply LoRA adapters via `--lora-path`:
+
+```bash
+sglang generate \
+  --model-path Qwen/Qwen-Image-Edit-2511 \
+  --lora-path prithivMLmods/Qwen-Image-Edit-2511-Anime \
+  --prompt "Transform into anime." \
+  --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" \
+  --save-output
+```
+
+For more usage examples (e.g. OpenAI compatible API, server mode), check the [CLI reference](https://docs.sglang.io/docs/sglang-diffusion/api/cli).
+
+## Contributing
+
+All contributions are welcome. The contribution guide is available [here](https://docs.sglang.io/docs/sglang-diffusion/contributing).
+
+## Acknowledgement
+
+We learnt and reused code from the following projects:
+
+- [FastVideo](https://github.com/hao-ai-lab/FastVideo.git). The major components of this repo are based on a fork of FastVideo on Sept. 24, 2025.
+- [xDiT](https://github.com/xdit-project/xDiT). We used the parallelism library from it.
+- [diffusers](https://github.com/huggingface/diffusers) We used the pipeline design from it.

python/sglang/multimodal_gen/__init__.py

@@ -0,0 +1,8 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+from sglang.multimodal_gen.configs.pipeline_configs import PipelineConfig
+from sglang.multimodal_gen.configs.sample import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+
+__all__ = ["DiffGenerator", "PipelineConfig", "SamplingParams"]
+
+# Trigger multimodal CI tests

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/README.md

@@ -0,0 +1,58 @@
+# ComfyUI SGLDiffusion Plugin
+
+A ComfyUI plugin for integrating with SGLang Diffusion server, supporting image and video generation capabilities.
+
+## Installation
+
+1. **Install SGLang**: Follow the [Installation Guide](../../../../../docs/diffusion/installation.md) to install `sglang[diffusion]`.
+2. **Install Plugin**: Copy this entire directory (`ComfyUI_SGLDiffusion`) to your ComfyUI `custom_nodes/` folder.
+3. **Restart ComfyUI**: Restart ComfyUI to load the plugin.
+
+## Usage
+
+The plugin supports two modes of operation: **Server Mode** (via HTTP API) and **Integrated Mode** (tight integration with ComfyUI).
+
+### Supported Models
+- **Z-Image**: High-speed image generation models (e.g., `Z-Image-Turbo`)
+- **FLUX**: State-of-the-art text-to-image models (e.g., `FLUX.1-dev`)
+- **Qwen-Image**: Multi-modal image generation models (e.g., `Qwen-Image`,`Qwen-Image-2512`). *Note: Image editing support is currently experimental and may have some issues.*
+
+### Mode 1: Server Mode (HTTP API)
+Connect to a standalone SGLang Diffusion server.
+
+1. **Start SGLang Diffusion Server**: Ensure the server is running and accessible.
+2. **Connect to Server**: Use the `SGLDiffusion Server Model` node to connect (default: `http://localhost:3000/v1`).
+3. **Generate Content**:
+   - `SGLDiffusion Generate Image`: For text-to-image and image editing.
+   - `SGLDiffusion Generate Video`: For text-to-video and image-to-video.
+4. **LoRA Support**: Use `SGLDiffusion Server Set LoRA` and `SGLDiffusion Server Unset LoRA`.
+
+### Mode 2: Integrated Mode (Tight Integration)
+Leverage SGLang's high-performance sampling directly within ComfyUI while using ComfyUI's front-end nodes (CLIP, VAE, etc.).
+
+1. **Load Model**: Use the `SGLDiffusion UNET Loader` node to load your diffusion model.
+2. **Configure Options**: Use the `SGLDiffusion Options` node to set runtime parameters like `num_gpus`, `tp_size`, `model_type`, or `enable_torch_compile`.
+3. **Sample**: Connect the loaded model to standard ComfyUI samplers. SGLang will handle the sampling process efficiently.
+4. **LoRA Support**: Use the `SGLDiffusion LoRA Loader` for native LoRA integration.
+
+## Example Workflows
+
+Reference workflow files are provided in the `workflows/` directory:
+
+- **`flux_sgld_sp.json`**: Multi-GPU (Sequence Parallelism) workflow for FLUX models. High-performance inference across multiple cards.
+- **`qwen_image_sgld.json`**: Qwen-Image generation with LoRA support. Optimized for multi-modal image tasks.
+- **`z-image_sgld.json`**: High-speed image generation using Z-Image.
+- **`sgld_text2img.json`**: Server-mode text-to-image generation with LoRA support.
+- **`sgld_image2video.json`**: Server-mode image-to-video generation.
+
+For other workflows supporting the models, you can easily use SGLang by replacing the official `UNET Loader` node with the `SGLDUNETLoader` node. Similarly, for LoRA support, replace the official LoRA loader with the `SGLDiffusion LoRA Loader`.
+
+To use these workflows:
+1. Open ComfyUI.
+2. Load the workflow JSON file from the `workflows/` directory.
+3. Adjust the parameters and model paths as needed.
+4. Run the workflow.
+
+## Current Implementation
+
+This plugin provides a high-performance backend for diffusion models in ComfyUI. By leveraging SGLang's optimized kernels and parallelization techniques (Tensor Parallelism, TeaCache, etc.), it significantly accelerates the sampling process, especially for large models like FLUX.

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/__init__.py

@@ -0,0 +1,13 @@
+"""
+ComfyUI SGLang Diffusion nodes package.
+"""
+
+try:
+    from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+except ImportError:
+    # ComfyUI dependencies not available (e.g., in test environment)
+    NODE_CLASS_MAPPINGS = {}
+    NODE_DISPLAY_NAME_MAPPINGS = {}
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/__init__.py

@@ -0,0 +1,14 @@
+"""
+Core components for SGLang Diffusion ComfyUI integration.
+Provides generator, model patcher, and server API client.
+"""
+
+from .generator import SGLDiffusionGenerator
+from .model_patcher import SGLDModelPatcher
+from .server_api import SGLDiffusionServerAPI
+
+__all__ = [
+    "SGLDiffusionGenerator",
+    "SGLDModelPatcher",
+    "SGLDiffusionServerAPI",
+]

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py

@@ -0,0 +1,231 @@
+"""
+Generator for SGLang Diffusion ComfyUI integration.
+"""
+
+import logging
+import os
+
+import psutil
+from comfy import model_detection, model_management
+from comfy.utils import (
+    calculate_parameters,
+    load_torch_file,
+    state_dict_prefix_replace,
+    unet_to_diffusers,
+)
+
+logger = logging.getLogger(__name__)
+
+try:
+    from sglang.multimodal_gen import DiffGenerator
+except ImportError:
+    logger.error(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+from ..executors import (
+    FluxExecutor,
+    QwenImageEditExecutor,
+    QwenImageExecutor,
+    ZImageExecutor,
+)
+from .model_patcher import SGLDModelPatcher
+
+
+class SGLDiffusionGenerator:
+    """Generator for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(self):
+        self.model_path = None
+        self.generator = None
+        self.executor = None
+        self.last_options = None
+
+        self.pipeline_class_dict = {
+            "flux": "ComfyUIFluxPipeline",
+            "lumina2": "ComfyUIZImagePipeline",  # zimage
+            "qwen_image": "ComfyUIQwenImagePipeline",
+            "qwen_image_edit": "ComfyUIQwenImageEditPipeline",
+        }
+        self.executor_class_dict = {
+            "flux": FluxExecutor,
+            "lumina2": ZImageExecutor,
+            "qwen_image": QwenImageExecutor,
+            "qwen_image_edit": QwenImageEditExecutor,
+        }
+
+    def __del__(self):
+        self.close_generator()
+
+    def init_generator(
+        self, model_path: str, pipeline_class_name: str, kwargs: dict = None
+    ):
+        """Initialize the diffusion generator."""
+        if self.generator is not None:
+            return self.generator
+        if kwargs is None:
+            kwargs = {}
+        # Set comfyui_mode for ComfyUI integration
+        kwargs["comfyui_mode"] = True
+        self.generator = DiffGenerator.from_pretrained(
+            model_path=model_path,
+            pipeline_class_name=pipeline_class_name,
+            **kwargs,
+        )
+        return self.generator
+
+    def kill_generator(self):
+        """Kill worker processes manually because generator shutdown cannot terminate them."""
+        current_pid = os.getpid()
+        worker_processes = []
+        for proc in psutil.process_iter(["pid", "name", "cmdline"]):
+            try:
+                # Look for sglang-diffusionWorker processes
+                if proc.info["cmdline"]:
+                    cmdline = " ".join(proc.info["cmdline"])
+                    if "sgl_diffusion::" in cmdline:
+                        if proc.info["pid"] != current_pid:
+                            worker_processes.append(proc)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+
+        if worker_processes:
+            logger.info(
+                f"Found {len(worker_processes)} worker processes to terminate..."
+            )
+            for proc in worker_processes:
+                try:
+                    logger.info(
+                        f"Terminating worker process {proc.info['pid']}: {proc.info['name']}"
+                    )
+                    proc.terminate()
+                    proc.wait(timeout=5)
+                except psutil.TimeoutExpired:
+                    logger.warning(
+                        f"Process {proc.info['pid']} did not terminate, forcing kill..."
+                    )
+                    try:
+                        proc.kill()
+                        proc.wait(timeout=2)
+                    except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+                        pass
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+    def close_generator(self):
+        """Close and cleanup the generator and all associated resources."""
+        if self.generator is not None:
+            self.generator.shutdown()
+            self.kill_generator()
+            # Clear other references
+            self.last_options = None
+            self.model_path = None
+            self.generator = None
+            self.executor = None
+
+    def get_comfyui_model(self, model_path: str, model_options: dict = None):
+        """Get ComfyUI model from model path."""
+        if model_options is None:
+            model_options = {}
+        dtype = model_options.get("dtype", None)
+        # Allow loading unets from checkpoint files
+        sd = load_torch_file(model_path)
+        diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
+        temp_sd = state_dict_prefix_replace(
+            sd, {diffusion_model_prefix: ""}, filter_keys=True
+        )
+        if len(temp_sd) > 0:
+            sd = temp_sd
+
+        parameters = calculate_parameters(sd)
+        load_device = model_management.get_torch_device()
+
+        model_detect_config = model_detection.detect_unet_config(sd, "")
+        model_type = model_detect_config.get("image_model", None)
+        if model_type is None or model_type not in self.pipeline_class_dict:
+            raise ValueError(f"Unsupported model type: {model_type}")
+        model_config = model_detection.model_config_from_unet(sd, "")
+
+        if model_config is not None:
+            new_sd = sd
+        else:
+            new_sd = model_detection.convert_diffusers_mmdit(sd, "")
+            if new_sd is not None:  # diffusers mmdit
+                model_config = model_detection.model_config_from_unet(new_sd, "")
+                if model_config is None:
+                    return None
+            else:  # diffusers unet
+                model_config = model_detection.model_config_from_diffusers_unet(sd)
+                if model_config is None:
+                    return None
+
+                diffusers_keys = unet_to_diffusers(model_config.unet_config)
+                new_sd = {}
+                for k in diffusers_keys:
+                    if k in sd:
+                        new_sd[diffusers_keys[k]] = sd.pop(k)
+        if dtype is None:
+            unet_dtype = model_management.unet_dtype(
+                model_params=parameters,
+                supported_dtypes=model_config.supported_inference_dtypes,
+            )
+        else:
+            unet_dtype = dtype
+
+        manual_cast_dtype = model_management.unet_manual_cast(
+            unet_dtype, load_device, model_config.supported_inference_dtypes
+        )
+        model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+        model_config.custom_operations = model_options.get("custom_operations", None)
+        model_config.unet_config["disable_unet_model_creation"] = True
+        comfyui_model = model_config.get_model({})
+        return comfyui_model, model_config, model_type
+
+    def load_model(
+        self, model_path: str, model_options: dict = None, sgld_options: dict = None
+    ):
+        """Load model and return model patcher."""
+        gather_options = {
+            "model_path": model_path,
+            "model_options": model_options,
+            "sgld_options": sgld_options,
+        }
+        if (
+            self.last_options is not None
+            and self.last_options == gather_options
+            and self.generator is not None
+        ):
+            return self.generator
+        else:
+            self.close_generator()
+
+        self.last_options = gather_options
+        self.model_path = model_path
+
+        comfyui_model, model_config, model_type = self.get_comfyui_model(
+            model_path, model_options
+        )
+        if model_type is None or model_type not in self.pipeline_class_dict:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+        set_model_type = sgld_options.pop("model_type", None) if sgld_options else None
+        if set_model_type is not None and set_model_type in self.pipeline_class_dict:
+            model_type = set_model_type
+
+        pipeline_class_name = self.pipeline_class_dict[model_type]
+        self.generator = self.init_generator(
+            model_path, pipeline_class_name, sgld_options
+        )
+
+        executor_class = self.executor_class_dict[model_type]
+        self.executor = executor_class(
+            self.generator, model_path, comfyui_model, model_config
+        )
+        comfyui_model.diffusion_model = self.executor
+
+        load_device = model_management.get_torch_device()
+        offload_device = model_management.unet_offload_device()
+
+        return SGLDModelPatcher(
+            comfyui_model, load_device, offload_device, model_type=model_type
+        )

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/model_patcher.py

@@ -0,0 +1,82 @@
+"""
+Model patcher for SGLang Diffusion ComfyUI integration.
+"""
+
+import copy
+
+from comfy.model_patcher import ModelPatcher
+
+
+class SGLDModelPatcher(ModelPatcher):
+    """Model patcher for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(
+        self,
+        model,
+        load_device,
+        offload_device,
+        size=0,
+        weight_inplace_update=False,
+        model_type=None,
+    ):
+        super().__init__(
+            model, load_device, offload_device, size, weight_inplace_update
+        )
+        self.lora_cache = {}
+        self.model_type = model_type
+        self.model_size_dict = {
+            "flux": 27 * 1024 * 1024 * 1024,
+            "lumina2": 8 * 1024 * 1024 * 1024,
+        }
+
+    def clone(self):
+        """Clone the model patcher."""
+        n = SGLDModelPatcher(
+            self.model,
+            self.load_device,
+            self.offload_device,
+            self.size,
+            weight_inplace_update=self.weight_inplace_update,
+        )
+        n.patches = {}
+        for k in self.patches:
+            n.patches[k] = self.patches[k][:]
+        n.patches_uuid = self.patches_uuid
+
+        n.object_patches = self.object_patches.copy()
+        n.model_options = copy.deepcopy(self.model_options)
+        n.backup = self.backup
+        n.object_patches_backup = self.object_patches_backup
+        n.lora_cache = copy.copy(self.lora_cache)
+        return n
+
+    def model_size(self):
+        """Get the model size in bytes."""
+        if self.model_type in self.model_size_dict:
+            return self.model_size_dict[self.model_type]
+        else:
+            return 0
+
+    def load(
+        self,
+        device_to=None,
+        lowvram_model_memory=0,
+        force_patch_weights=False,
+        full_load=False,
+    ):
+        """Load model (no-op for SGLang Diffusion)."""
+        pass
+
+    def patch_model(
+        self,
+        device_to=None,
+        lowvram_model_memory=0,
+        load_weights=True,
+        force_patch_weights=False,
+    ):
+        """Patch model (no-op for SGLang Diffusion)."""
+        pass
+
+    def unpatch_model(self, device_to=None, unpatch_weights=True):
+        """Unpatch model (no-op for SGLang Diffusion)."""
+        pass

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/server_api.py

@@ -0,0 +1,539 @@
+"""
+SGLang Diffusion Server API client.
+Provides a low-level interface for interacting with SGLang Diffusion HTTP server.
+"""
+
+import base64
+import io
+import os
+import time
+from typing import Any, Dict, Optional
+
+import requests
+from PIL import Image
+
+
+class SGLDiffusionServerAPI:
+    """Client for SGLang Diffusion HTTP server API."""
+
+    def __init__(self, base_url: str, api_key: str = "sk-proj-1234567890"):
+        """
+        Initialize the API client.
+
+        Args:
+            base_url: Base URL of the SGLang Diffusion server (e.g., "http://localhost:30010/v1")
+            api_key: API key for authentication (default: "sk-proj-1234567890")
+        """
+        # Ensure base_url doesn't end with /v1 if it's already there
+        if base_url.endswith("/v1"):
+            self.base_url = base_url
+        elif base_url.endswith("/v1/"):
+            self.base_url = base_url.rstrip("/")
+        else:
+            self.base_url = f"{base_url.rstrip('/')}/v1"
+
+        self.api_key = api_key
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model served by this server.
+
+        Returns:
+            Dictionary containing model information including:
+            - model_path: Path to the model
+            - task_type: Type of task (e.g., "T2V", "I2I")
+            - pipeline_name: Name of the pipeline
+            - num_gpus: Number of GPUs
+            - dit_precision: DiT model precision
+            - vae_precision: VAE model precision
+        """
+        try:
+            # Remove /v1 from base_url for /models endpoint
+            models_url = self.base_url.removesuffix("/v1") + "/models"
+            response = requests.get(models_url, headers=self.headers, timeout=30)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to get model info: {str(e)}")
+
+    def generate_image(
+        self,
+        prompt: str,
+        image_path: Optional[str] = None,
+        mask_path: Optional[str] = None,
+        size: Optional[str] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        n: int = 1,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        response_format: str = "b64_json",
+        quality: Optional[str] = "auto",
+        style: Optional[str] = "vivid",
+        background: Optional[str] = "auto",
+        output_format: Optional[str] = None,
+        generator_device: Optional[str] = "cuda",
+    ) -> Dict[str, Any]:
+        """
+        Generate or edit an image using SGLang Diffusion API.
+        If image_path is provided, calls the edit endpoint; otherwise calls the generation endpoint.
+
+        Args:
+            prompt: Text prompt for image generation/editing
+            image_path: Optional path to input image file for editing. If provided, uses edit API.
+            mask_path: Optional path to mask image file (only used when image_path is provided)
+            size: Image size in format "WIDTHxHEIGHT" (e.g., "1024x1024")
+            width: Image width (used if size is not provided)
+            height: Image height (used if size is not provided)
+            n: Number of images to generate (1-10)
+            negative_prompt: Negative prompt to avoid certain elements
+            guidance_scale: Classifier-free guidance scale
+            num_inference_steps: Number of denoising steps
+            seed: Random seed for reproducible generation
+            enable_teacache: Enable TEA cache acceleration
+            response_format: Response format ("b64_json" or "url")
+            quality: Image quality ("auto", "standard", "hd") - only for generation
+            style: Image style ("vivid" or "natural") - only for generation
+            background: Background type ("auto", "transparent", "opaque")
+            output_format: Output format ("png", "jpeg", "webp")
+            generator_device: Device for random generator ("cuda" or "cpu")
+
+        Returns:
+            Dictionary containing the API response with generated/edited image data
+        """
+        if not prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        # Determine size
+        if size is None:
+            if width is not None and height is not None:
+                size = f"{width}x{height}"
+            else:
+                size = "1024x1024"
+
+        # Build common parameters
+        common_params = self._build_image_common_params(
+            prompt=prompt,
+            size=size,
+            n=n,
+            response_format=response_format,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            seed=seed,
+            enable_teacache=enable_teacache,
+            background=background,
+            output_format=output_format,
+            generator_device=generator_device,
+        )
+
+        # If image_path is provided, use edit endpoint
+        if image_path:
+            if not os.path.exists(image_path):
+                raise FileNotFoundError(f"Image file not found: {image_path}")
+
+            # Prepare multipart form data for edit
+            files: Dict[str, Any] = {}
+            data = common_params.copy()
+
+            # Add image file
+            files["image"] = (
+                os.path.basename(image_path),
+                open(image_path, "rb"),
+                self._get_content_type(image_path),
+            )
+
+            # Add mask file if provided
+            if mask_path:
+                if not os.path.exists(mask_path):
+                    raise FileNotFoundError(f"Mask file not found: {mask_path}")
+                files["mask"] = (
+                    os.path.basename(mask_path),
+                    open(mask_path, "rb"),
+                    self._get_content_type(mask_path),
+                )
+
+            # Prepare headers for multipart form data
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+            }
+
+            try:
+                response = requests.post(
+                    f"{self.base_url}/images/edits",
+                    files=files,
+                    data=data,
+                    headers=headers,
+                    timeout=300,  # 5 minutes timeout for generation
+                )
+                response.raise_for_status()
+                return response.json()
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError(f"Failed to edit image: {str(e)}")
+            finally:
+                # Close file handles
+                for file_tuple in files.values():
+                    if isinstance(file_tuple, tuple) and len(file_tuple) > 1:
+                        file_tuple[1].close()
+        else:
+            # Use generation endpoint - add generation-specific parameters
+            payload = common_params.copy()
+            if quality:
+                payload["quality"] = quality
+            if style:
+                payload["style"] = style
+
+            try:
+                response = requests.post(
+                    f"{self.base_url}/images/generations",
+                    json=payload,
+                    headers=self.headers,
+                    timeout=300,  # 5 minutes timeout for generation
+                )
+                response.raise_for_status()
+                return response.json()
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError(f"Failed to generate image: {str(e)}")
+
+    def generate_video(
+        self,
+        prompt: str,
+        size: Optional[str] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        seconds: Optional[int] = 4,
+        fps: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        generator_device: Optional[str] = "cuda",
+        input_reference: Optional[str] = None,
+        output_path: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Generate a video using SGLang Diffusion API and wait for completion.
+
+        Args:
+            prompt: Text prompt for video generation
+            size: Video size in format "WIDTHxHEIGHT" (e.g., "1280x720")
+            width: Video width (used if size is not provided)
+            height: Video height (used if size is not provided)
+            seconds: Duration of the video in seconds
+            fps: Frames per second
+            num_frames: Number of frames (overrides seconds * fps if provided)
+            negative_prompt: Negative prompt to avoid certain elements
+            guidance_scale: Classifier-free guidance scale
+            num_inference_steps: Number of denoising steps
+            seed: Random seed for reproducible generation
+            enable_teacache: Enable TEA cache acceleration
+            generator_device: Device for random generator ("cuda" or "cpu")
+            input_reference: Path to input reference image for image-to-video
+
+        Returns:
+            Dictionary containing completed video job information with file_path
+        """
+        if not prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        # Determine size
+        if size is None:
+            if width is not None and height is not None:
+                size = f"{width}x{height}"
+            else:
+                size = "720x1280"
+
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "prompt": prompt,
+            "size": size,
+        }
+
+        # Add optional parameters
+        if seconds is not None:
+            payload["seconds"] = seconds
+        if fps is not None:
+            payload["fps"] = fps
+        if num_frames is not None:
+            payload["num_frames"] = num_frames
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if guidance_scale is not None:
+            payload["guidance_scale"] = guidance_scale
+        if num_inference_steps is not None:
+            payload["num_inference_steps"] = num_inference_steps
+        if seed is not None and seed >= 0:
+            payload["seed"] = seed
+        if enable_teacache:
+            payload["enable_teacache"] = True
+        if generator_device:
+            payload["generator_device"] = generator_device
+        if input_reference:
+            payload["input_reference"] = input_reference
+        if output_path:
+            payload["output_path"] = output_path
+
+        try:
+            # Create video generation job
+            response = requests.post(
+                f"{self.base_url}/videos",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            video_job = response.json()
+            video_id = video_job.get("id")
+
+            # Wait for completion with fixed polling
+            poll_interval = 5  # 5 seconds
+            max_wait_time = 3600  # 1 hour
+            max_consecutive_errors = 5
+            consecutive_errors = 0
+            start_time = time.time()
+
+            while time.time() - start_time < max_wait_time:
+                try:
+                    status_response = requests.get(
+                        f"{self.base_url}/videos/{video_id}",
+                        headers=self.headers,
+                        timeout=30,
+                    )
+                    status_response.raise_for_status()
+                    status = status_response.json()
+
+                    # Reset error counter on successful request
+                    consecutive_errors = 0
+
+                    if status.get("status") == "completed":
+                        return status
+                    elif status.get("status") == "failed":
+                        error = status.get("error", {})
+                        error_msg = (
+                            error.get("message", "Unknown error")
+                            if error
+                            else "Unknown error"
+                        )
+                        raise RuntimeError(f"Video generation failed: {error_msg}")
+                except requests.exceptions.ConnectionError as e:
+                    # Connection errors - likely server is down
+                    consecutive_errors += 1
+                    if consecutive_errors >= max_consecutive_errors:
+                        raise RuntimeError(
+                            f"Lost connection to server after {consecutive_errors} consecutive errors. "
+                            f"Server may be unavailable: {str(e)}"
+                        )
+                except requests.exceptions.RequestException as e:
+                    # Other network errors - continue polling but track errors
+                    consecutive_errors += 1
+                    if consecutive_errors >= max_consecutive_errors:
+                        raise RuntimeError(
+                            f"Network error after {consecutive_errors} consecutive failures: {str(e)}"
+                        )
+
+                time.sleep(poll_interval)
+
+            raise TimeoutError(
+                f"Video generation timed out after {max_wait_time} seconds"
+            )
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to generate video: {str(e)}")
+
+    def _build_image_common_params(
+        self,
+        prompt: str,
+        size: str,
+        n: int,
+        response_format: str,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        background: Optional[str] = None,
+        output_format: Optional[str] = None,
+        generator_device: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Build common parameters for both image generation and editing.
+
+        Returns:
+            Dictionary containing common parameters
+        """
+        params: Dict[str, Any] = {
+            "prompt": prompt,
+            "size": size,
+            "n": max(1, min(n, 10)),
+            "response_format": response_format,
+        }
+
+        # Add optional parameters
+        if negative_prompt:
+            params["negative_prompt"] = negative_prompt
+        if guidance_scale is not None:
+            params["guidance_scale"] = guidance_scale
+        if num_inference_steps is not None:
+            params["num_inference_steps"] = num_inference_steps
+        if seed is not None and seed >= 0:
+            params["seed"] = seed
+        if enable_teacache:
+            params["enable_teacache"] = True
+        if background:
+            params["background"] = background
+        if output_format:
+            params["output_format"] = output_format
+        if generator_device:
+            params["generator_device"] = generator_device
+
+        return params
+
+    def _get_content_type(self, file_path: str) -> str:
+        """Get content type based on file extension."""
+        ext = os.path.splitext(file_path)[1].lower()
+        content_types = {
+            ".png": "image/png",
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".webp": "image/webp",
+        }
+        return content_types.get(ext, "image/png")
+
+    def decode_image_from_response(
+        self, response_data: Dict[str, Any], index: int = 0
+    ) -> Image.Image:
+        """
+        Decode base64 image from API response.
+
+        Args:
+            response_data: API response dictionary
+            index: Index of the image in the response (default: 0)
+
+        Returns:
+            PIL Image object
+        """
+        if "data" not in response_data or not response_data["data"]:
+            raise ValueError("No image data in response")
+
+        if index >= len(response_data["data"]):
+            raise IndexError(f"Image index {index} out of range")
+
+        image_data = response_data["data"][index]
+        if "b64_json" not in image_data or not image_data["b64_json"]:
+            raise ValueError("No base64 image data found")
+
+        image_bytes = base64.b64decode(image_data["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes))
+
+        # Convert to RGB if needed
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        return image
+
+    def set_lora(
+        self,
+        lora_nickname: str,
+        lora_path: Optional[str] = None,
+        target: str = "all",
+    ) -> Dict[str, Any]:
+        """
+        Set a LoRA adapter for the specified transformer(s).
+
+        Args:
+            lora_nickname: The nickname of the adapter (required).
+            lora_path: Path to the LoRA adapter (local path or HF repo id).
+                      Required for the first load; optional if re-activating a cached nickname.
+            target: Which transformer(s) to apply the LoRA to. One of:
+                - "all": Apply to all transformers (default)
+                - "transformer": Apply only to the primary transformer (high noise for Wan2.2)
+                - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
+                - "critic": Apply only to the critic model
+
+        Returns:
+            Dictionary containing the API response with status and message
+        """
+        if not lora_nickname:
+            raise ValueError("lora_nickname cannot be empty")
+
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "lora_nickname": lora_nickname,
+            "target": target,
+        }
+
+        # Add optional lora_path if provided
+        if lora_path:
+            payload["lora_path"] = lora_path
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/set_lora",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}")
+
+    def unset_lora(
+        self,
+        target: str = "all",
+    ) -> Dict[str, Any]:
+        """
+        Unset (unmerge) LoRA weights from the base model.
+
+        Args:
+            target: same as set_lora
+
+        Returns:
+            Dictionary containing the API response with status and message
+        """
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "target": target,
+        }
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/unmerge_lora_weights",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}")
+
+
+if __name__ == "__main__":
+    api = SGLDiffusionServerAPI(
+        base_url="http://localhost:30010/v1", api_key="sk-proj-1234567890"
+    )
+    model_info = api.get_model_info()
+    print(api.get_model_info())
+    if model_info.get("task_type") == "T2V" or model_info.get("task_type") == "I2V":
+        print(
+            api.generate_video(
+                prompt="A calico cat playing a piano on stage",
+                num_inference_steps=1,
+                size="480x480",
+            )
+        )
+    else:
+        print(
+            api.generate_image(
+                prompt="A calico cat playing a piano on stage", size="1024x1024"
+            )
+        )

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/__init__.py

@@ -0,0 +1,17 @@
+"""
+ComfyUI SGLang Diffusion executors package.
+Provides executor classes for different model types.
+"""
+
+from .base import SGLDiffusionExecutor
+from .flux import FluxExecutor
+from .qwen_image import QwenImageEditExecutor, QwenImageExecutor
+from .zimage import ZImageExecutor
+
+__all__ = [
+    "SGLDiffusionExecutor",
+    "FluxExecutor",
+    "ZImageExecutor",
+    "QwenImageExecutor",
+    "QwenImageEditExecutor",
+]

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py

@@ -0,0 +1,56 @@
+"""
+Base executor class for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+
+class SGLDiffusionExecutor(torch.nn.Module):
+    """Base executor class for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super(SGLDiffusionExecutor, self).__init__()
+        self.generator = generator
+        self.model_path = model_path
+        self.model = model
+        self.dtype = config.unet_config["dtype"]
+        self.config = config
+        self.loras = []
+
+    @staticmethod
+    def should_suppress_logs(timestep):
+        """Determine if logs should be suppressed based on timestep value."""
+        if torch.is_tensor(timestep):
+            return bool((timestep < 1.0).item())
+        return bool(timestep < 1.0)
+
+    def set_lora(self, lora_nickname=None, lora_path=None, strength=None, target=None):
+        """Set LoRA adapter using SGLang Diffusion API."""
+        if len(lora_nickname) > 0:
+            self.generator.set_lora(
+                lora_nickname=lora_nickname,
+                lora_path=lora_path,
+                strength=strength,
+                target=target,
+            )
+
+    def _unpack_latents(self, latents, height, width, channels):
+        """Unpack latents from packed format to standard format."""
+        batch_size = latents.shape[0]
+        latents = latents.view(batch_size, height // 2, width // 2, channels, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        latents = latents.reshape(batch_size, channels, height, width)
+
+        return latents
+
+    def _pack_latents(self, latents):
+        """Pack latents from standard format to packed format."""
+        batch_size, num_channels_latents, height, width = latents.shape
+        latents = latents.view(
+            batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+        )
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size, (height // 2) * (width // 2), num_channels_latents * 4
+        )
+        return latents

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py

@@ -0,0 +1,69 @@
+"""
+Flux executor for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+try:
+    from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+except ImportError:
+    print(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+from .base import SGLDiffusionExecutor
+
+
+class FluxExecutor(SGLDiffusionExecutor):
+    """Executor for Flux models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(self, x, timestep, context, y=None, guidance=None, **kwargs):
+        """Forward pass for Flux model."""
+        hidden_states = self._pack_latents(x)
+        timesteps = timestep * 1000.0
+        encoder_hidden_states = context
+        pooled_projections = y
+        guidance = guidance * 1000.0
+
+        B, C, H, W = x.shape
+        height = H * 8
+        width = W * 8
+        # Create SamplingParams
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=3.5,  # Flux typically uses embedded_cfg_scale=3.5
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        req.latents = hidden_states  # Set as [B, S, D] format directly
+        req.timesteps = timesteps  # ComfyUI's timesteps parameter
+        req.prompt_embeds = [pooled_projections, encoder_hidden_states]  # [CLIP, T5]
+        req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+        # Set pooled_projections (required by Flux)
+        req.pooled_embeds = [pooled_projections]  # List format as per Req definition
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        # Send request to scheduler
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+        return self._unpack_latents(noise_pred, H, W, C).to(x.device)

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/qwen_image.py

@@ -0,0 +1,172 @@
+"""
+QwenImage executor for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+try:
+    from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+except ImportError:
+    print(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+import comfy.ldm.common_dit
+
+from .base import SGLDiffusionExecutor
+
+
+class QwenImageExecutor(SGLDiffusionExecutor):
+    """Executor for QwenImage models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+        self.patch_size = 2
+
+    def _pack_latents(self, x):
+        """Process hidden states for QwenImage model."""
+        latents = comfy.ldm.common_dit.pad_to_patch_size(
+            x, (1, self.patch_size, self.patch_size)
+        )
+        orig_shape = latents.shape
+        latents = latents.view(
+            orig_shape[0],
+            orig_shape[1],
+            orig_shape[-3],
+            orig_shape[-2] // 2,
+            2,
+            orig_shape[-1] // 2,
+            2,
+        )
+        latents = latents.permute(0, 2, 3, 5, 1, 4, 6)
+        latents = latents.reshape(
+            orig_shape[0],
+            orig_shape[-3] * (orig_shape[-2] // 2) * (orig_shape[-1] // 2),
+            orig_shape[1] * 4,
+        )
+        return latents, orig_shape
+
+    def _unpack_latents(self, latents, num_embeds, orig_shape, x):
+        """Unpack hidden states from packed format to standard format."""
+        latents = latents[:, :num_embeds].view(
+            orig_shape[0],
+            orig_shape[-3],
+            orig_shape[-2] // 2,
+            orig_shape[-1] // 2,
+            orig_shape[1],
+            2,
+            2,
+        )
+        latents = latents.permute(0, 4, 1, 2, 5, 3, 6)
+        latents = latents.reshape(orig_shape)[:, :, :, : x.shape[-2], : x.shape[-1]]
+        return latents
+
+    def forward(self, x, timestep, context, **kwargs):
+        """Forward pass for QwenImage model."""
+        latents, orig_shape = self._pack_latents(x)
+        num_embeds = latents.shape[1]
+        height = orig_shape[-2] * 8
+        width = orig_shape[-1] * 8
+
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=1.0,
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents
+        req.timesteps = timestep * 1000.0
+        req.prompt_embeds = [context]
+        req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+
+        return self._unpack_latents(noise_pred, num_embeds, orig_shape, x)
+
+
+class QwenImageEditExecutor(QwenImageExecutor):
+    """Executor for QwenImageEdit models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(
+        self,
+        x,
+        timestep,
+        context,
+        attention_mask=None,
+        ref_latents=None,
+        additional_t_cond=None,
+        transformer_options={},
+        **kwargs,
+    ):
+        """Forward pass for QwenImageEdit model."""
+        latents, orig_shape = self._pack_latents(x)
+        num_embeds = latents.shape[1]
+        height = orig_shape[-2] * 8
+        width = orig_shape[-1] * 8
+
+        # Prepare vae_image_sizes for the condition image (ref_latents)
+        vae_image_sizes = []
+        pack_ref_latents = None
+
+        # TODO: sgld now don't support multiple condition images, so we only support one condition image for now.
+        if ref_latents is not None and len(ref_latents) > 0:
+            pack_ref_latents, orig_ref_shape = self._pack_latents(ref_latents[0])
+            vae_image_sizes = [(orig_ref_shape[-1], orig_ref_shape[-2])]
+
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=1.0,
+            image_path="",
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents
+        req.image_latent = pack_ref_latents
+        req.timesteps = timestep * 1000.0
+        req.vae_image_sizes = vae_image_sizes
+        req.prompt_embeds = [context]
+        req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+
+        return self._unpack_latents(noise_pred, num_embeds, orig_shape, x)

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/zimage.py

@@ -0,0 +1,64 @@
+"""
+ZImage executor for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+try:
+    from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+except ImportError:
+    print(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+from .base import SGLDiffusionExecutor
+
+
+class ZImageExecutor(SGLDiffusionExecutor):
+    """Executor for ZImage models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(self, x, timesteps, context, **kwargs):
+        """Forward pass for ZImage model."""
+        B, C, H, W = x.shape
+        height = H * 8
+        width = W * 8
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=1.0,
+            height=height,
+            width=width,
+            num_frames=1,  # For images
+            num_inference_steps=1,  # Single step for ComfyUI
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timesteps),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        latents = x.unsqueeze(2)
+        context = context.squeeze(0)
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents  # ComfyUI's x parameter
+        req.timesteps = timesteps * 1000.0  # ComfyUI's timesteps parameter
+        req.prompt_embeds = [
+            context
+        ]  # ComfyUI's context parameter (must be List[Tensor])
+        req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+
+        return noise_pred.permute(1, 0, 2, 3).to(x.device)

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/nodes.py

@@ -0,0 +1,715 @@
+"""
+ComfyUI nodes for SGLang Diffusion integration.
+Provides nodes for connecting to SGLang Diffusion server and generating images/videos.
+"""
+
+import os
+import uuid
+
+import folder_paths
+import torch
+
+from .core import SGLDiffusionGenerator, SGLDiffusionServerAPI
+from .utils import (
+    convert_b64_to_tensor_image,
+    convert_video_to_comfy_video,
+    get_image_path,
+    is_empty_image,
+)
+
+
+class SGLDOptions:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "optional": {
+                "model_type": (
+                    ["auto-detect", "qwen_image", "qwen_image_edit", "flux", "lumina2"],
+                    {"default": "auto-detect"},
+                ),
+                "enable_torch_compile": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "num_gpus": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "tp_size": ("INT", {"default": -1, "min": -1, "step": 1}),
+                "sp_degree": ("INT", {"default": -1, "min": -1, "step": 1}),
+                "ulysses_degree": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                    },
+                ),
+                "ring_degree": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                    },
+                ),
+                "dp_size": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "dp_degree": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "enable_cfg_parallel": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "attention_backend": (
+                    "STRING",
+                    {"default": ""},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_OPTIONS",)
+    RETURN_NAMES = ("sgld_options",)
+    FUNCTION = "create_options"
+    CATEGORY = "SGLDiffusion"
+
+    def create_options(
+        self,
+        model_type: str = "auto-detect",
+        enable_torch_compile: bool = False,
+        num_gpus: int = 1,
+        tp_size: int = -1,
+        sp_degree: int = -1,
+        ulysses_degree: int = -1,
+        ring_degree: int = -1,
+        dp_size: int = 1,
+        dp_degree: int = 1,
+        enable_cfg_parallel: bool = False,
+        attention_backend: str = "",
+    ):
+        """
+        Build a dictionary of SGLang Diffusion runtime options.
+        """
+        # Convert -1 to None for optional parameters (matching ServerArgs defaults)
+        ulysses_degree = None if ulysses_degree == -1 else ulysses_degree
+        ring_degree = None if ring_degree == -1 else ring_degree
+        attention_backend = None if attention_backend == "" else attention_backend
+
+        options = {
+            "model_type": model_type,
+            "enable_torch_compile": enable_torch_compile,
+            "num_gpus": num_gpus,
+            "tp_size": tp_size,
+            "sp_degree": sp_degree,
+            "ulysses_degree": ulysses_degree,
+            "ring_degree": ring_degree,
+            "dp_size": dp_size,
+            "dp_degree": dp_degree,
+            "enable_cfg_parallel": enable_cfg_parallel,
+            "attention_backend": attention_backend,
+        }
+
+        # Strip None to keep payload clean
+        options = {k: v for k, v in options.items() if v is not None}
+        return (options,)
+
+
+class SGLDLoraLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "lora_name": (folder_paths.get_filename_list("loras"),),
+                "strength_model": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0, "max": 10, "step": 0.01},
+                ),
+                "nickname": ("STRING", {"default": ""}),
+                "target": (
+                    ["all", "transformer", "transformer_2", "critic"],
+                    {"default": "all"},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_lora"
+
+    CATEGORY = "SGLDiffusion"
+
+    def load_lora(
+        self, model, lora_name, strength_model=1.0, nickname="", target="all"
+    ):
+        """Load LoRA adapter using SGLang Diffusion API."""
+        lora_path = folder_paths.get_full_path("loras", lora_name)
+        assert model is not None
+        bi = model.clone()
+        nickname = nickname if nickname != "" else str("lora" + str(uuid.uuid4()))
+        # set lora in the model
+        bi.patches[nickname] = (lora_path, strength_model, target)
+
+        # prepare input for the SGLang Diffusion API
+        lora_input = {
+            "lora_nickname": [],
+            "lora_path": [],
+            "strength": [],
+            "target": [],
+        }
+        for nickname, lora_info in bi.patches.items():
+            lora_input["lora_nickname"].append(nickname)
+            lora_input["lora_path"].append(lora_info[0])
+            lora_input["strength"].append(lora_info[1])
+            lora_input["target"].append(lora_info[2])
+
+        # call the SGLang Diffusion API
+        model.model.diffusion_model.set_lora(**lora_input)
+        return (model,)
+
+
+class SGLDUNETLoader:
+    def __init__(self):
+        self.generator = SGLDiffusionGenerator()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("diffusion_models"),),
+                "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
+            },
+            "optional": {
+                "sgld_options": ("SGLD_OPTIONS",),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_unet"
+
+    CATEGORY = "SGLDiffusion"
+
+    def load_unet(self, unet_name, weight_dtype, sgld_options: dict = None):
+        model_options = {}
+        if weight_dtype == "fp8_e4m3fn":
+            model_options["dtype"] = torch.float8_e4m3fn
+        elif weight_dtype == "fp8_e5m2":
+            model_options["dtype"] = torch.float8_e5m2
+
+        unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
+
+        model = self.generator.load_model(
+            unet_path, model_options=model_options, sgld_options=sgld_options
+        )
+        return (model,)
+
+
+class SGLDiffusionServerModel:
+    """Node to load and manage SGLang Diffusion server connection."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_url": (
+                    "STRING",
+                    {
+                        "default": "http://localhost:3000/v1",
+                        "multiline": False,
+                    },
+                ),
+                "api_key": (
+                    "STRING",
+                    {
+                        "default": "sk-proj-1234567890",
+                        "multiline": False,
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT", "STRING")
+    RETURN_NAMES = ("sgld_client", "model_info")
+    FUNCTION = "load_server"
+    CATEGORY = "SGLDiffusion"
+
+    def load_server(self, base_url: str, api_key: str):
+        """Initialize OpenAI client for SGLang Diffusion server."""
+        client = SGLDiffusionServerAPI(base_url=base_url, api_key=api_key)
+        try:
+            model_info = client.get_model_info()
+            # Format model_info as a readable string
+            info_lines = ["=== SGLDiffusion Model Info ==="]
+            for key, value in model_info.items():
+                info_lines.append(f"{key}: {value}")
+            model_info_str = "\n".join(info_lines)
+        except Exception as e:
+            model_info_str = f"Failed to get model info: {str(e)}"
+        return (client, model_info_str)
+
+
+class SGLDiffusionGenerateImage:
+    """Node to generate images using SGLang Diffusion."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "positive_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Text prompt for image generation",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Negative prompt to avoid certain elements",
+                    },
+                ),
+                "image": (
+                    "IMAGE",
+                    {
+                        "default": None,
+                        "tooltip": "input image to use for editing",
+                    },
+                ),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": -1,
+                        "max": 2**32 - 1,
+                    },
+                ),
+                "steps": (
+                    "INT",
+                    {
+                        "default": 6,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                    },
+                ),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 7.0,
+                        "min": 1.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                    },
+                ),
+                "width": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 64,
+                    },
+                ),
+                "height": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 64,
+                    },
+                ),
+                "enable_teacache": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate_image"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def generate_image(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        positive_prompt: str,
+        negative_prompt: str = "",
+        image: torch.Tensor = None,
+        seed: int = 1024,
+        steps: int = 6,
+        cfg: float = 7.0,
+        width: int = 1024,
+        height: int = 1024,
+        enable_teacache: bool = False,
+    ):
+        """Generate image using SGLang Diffusion API."""
+        if not positive_prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        size = f"{width}x{height}"
+
+        # Prepare request parameters
+        request_params = {
+            "prompt": positive_prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+
+        # Add optional parameters if provided
+        if negative_prompt:
+            request_params["negative_prompt"] = negative_prompt
+        if cfg is not None:
+            request_params["guidance_scale"] = cfg
+        if steps is not None:
+            request_params["num_inference_steps"] = steps
+        if seed is not None and seed >= 0:
+            request_params["seed"] = seed
+        if enable_teacache:
+            request_params["enable_teacache"] = True
+        if image is not None:
+            # If the image is empty, use the size of the image to generate the image
+            if is_empty_image(image):
+                width, height = image.shape[2], image.shape[1]
+                size = f"{width}x{height}"
+                request_params["size"] = size
+            else:
+                request_params["image_path"] = get_image_path(image)
+
+        # Call API
+        try:
+            response = sgld_client.generate_image(**request_params)
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate image: {str(e)}")
+
+        # Decode base64 image
+        if not response["data"] or not response["data"][0]["b64_json"]:
+            raise RuntimeError("No image data in response")
+        image_data = response["data"][0]["b64_json"]
+        image = convert_b64_to_tensor_image(image_data)
+
+        return (image,)
+
+
+class SGLDiffusionGenerateVideo:
+    """Node to generate videos using SGLang Diffusion."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "positive_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Text prompt for video generation",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Negative prompt to avoid certain elements",
+                    },
+                ),
+                "image": (
+                    "IMAGE",
+                    {
+                        "default": None,
+                        "tooltip": "input image to use for image-to-video",
+                    },
+                ),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": -1,
+                        "max": 2**32 - 1,
+                    },
+                ),
+                "steps": (
+                    "INT",
+                    {
+                        "default": 6,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                    },
+                ),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 7.0,
+                        "min": 1.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                    },
+                ),
+                "width": (
+                    "INT",
+                    {
+                        "default": 1280,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 1,
+                    },
+                ),
+                "height": (
+                    "INT",
+                    {
+                        "default": 720,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 1,
+                    },
+                ),
+                "num_frames": (
+                    "INT",
+                    {
+                        "default": 120,
+                        "min": 1,
+                        "max": 1000,
+                        "step": 1,
+                    },
+                ),
+                "fps": (
+                    "INT",
+                    {
+                        "default": 24,
+                        "min": 1,
+                        "max": 60,
+                        "step": 1,
+                    },
+                ),
+                "seconds": (
+                    "INT",
+                    {
+                        "default": 5,
+                        "min": 1,
+                        "max": 60,
+                        "step": 1,
+                    },
+                ),
+                "enable_teacache": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO", "STRING")
+    RETURN_NAMES = ("video", "video_path")
+    FUNCTION = "generate_video"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def generate_video(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        positive_prompt: str,
+        negative_prompt: str = "",
+        image: torch.Tensor = None,
+        seed: int = 1024,
+        steps: int = 6,
+        cfg: float = 7.0,
+        width: int = 1280,
+        height: int = 720,
+        num_frames: int = 120,
+        fps: int = 24,
+        seconds: int = 5,
+        enable_teacache: bool = False,
+    ):
+        """Generate video using SGLang Diffusion API."""
+        if not positive_prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        size = f"{width}x{height}"
+        output_dir = folder_paths.get_temp_directory()
+
+        # Prepare request parameters
+        request_params = {
+            "prompt": positive_prompt,
+            "size": size,
+            "seconds": seconds,
+            "fps": fps,
+            "output_path": output_dir,
+        }
+
+        # Add optional parameters if provided
+        if negative_prompt:
+            request_params["negative_prompt"] = negative_prompt
+        if cfg is not None:
+            request_params["guidance_scale"] = cfg
+        if steps is not None:
+            request_params["num_inference_steps"] = steps
+        if seed is not None and seed >= 0:
+            request_params["seed"] = seed
+        if enable_teacache:
+            request_params["enable_teacache"] = True
+        if num_frames is not None:
+            request_params["num_frames"] = num_frames
+        if image is not None:
+            # If the image is empty, use the size of the image to generate the video
+            if is_empty_image(image):
+                width, height = image.shape[2], image.shape[1]
+                size = f"{width}x{height}"
+                request_params["size"] = size
+            else:
+                request_params["input_reference"] = get_image_path(image)
+
+        # Call API
+        try:
+            response = sgld_client.generate_video(**request_params)
+            video_path = response.get("file_path", "")
+            video = convert_video_to_comfy_video(video_path, height, width)
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate video: {str(e)}")
+
+        return (video, video_path)
+
+
+class SGLDiffusionServerSetLora:
+    """Node to set LoRA adapter for SGLang Diffusion server."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "lora_name": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "The name of the LoRA adapter",
+                    },
+                ),
+            },
+            "optional": {
+                "lora_nickname": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "The nickname of the LoRA adapter",
+                    },
+                ),
+                "target": (
+                    [
+                        "all",
+                        "transformer",
+                        "transformer_2",
+                        "critic",
+                    ],
+                    {
+                        "default": "all",
+                        "tooltip": "Which transformer(s) to apply the LoRA to",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT",)
+    RETURN_NAMES = ("sgld_client",)
+    FUNCTION = "set_lora"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def set_lora(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        lora_name: str = "",
+        lora_nickname: str = "",
+        target: str = "all",
+    ):
+        """Set LoRA adapter using SGLang Diffusion API."""
+        if lora_nickname == "":
+            lora_nickname = os.path.splitext(lora_name)[0]
+
+        # Prepare request parameters
+        request_params = {
+            "lora_nickname": lora_nickname,
+            "lora_path": lora_name,
+            "target": target,
+        }
+
+        # Call API
+        try:
+            sgld_client.set_lora(**request_params)
+            return (sgld_client,)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}")
+
+
+class SGLDiffusionServerUnsetLora:
+    """Node to unset LoRA adapter for SGLang Diffusion server."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+            },
+            "optional": {
+                "target": (
+                    [
+                        "all",
+                        "transformer",
+                        "transformer_2",
+                        "critic",
+                    ],
+                    {
+                        "default": "all",
+                        "tooltip": "Which transformer(s) to unset the LoRA from",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT",)
+    RETURN_NAMES = ("sgld_client",)
+    FUNCTION = "unset_lora"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def unset_lora(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        target: str = "all",
+    ):
+        """Unset LoRA adapter using SGLang Diffusion API."""
+        try:
+            sgld_client.unset_lora(target=target)
+            return (sgld_client,)
+        except Exception as e:
+            raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}")
+
+
+# Register nodes
+NODE_CLASS_MAPPINGS = {
+    "SGLDiffusionServerModel": SGLDiffusionServerModel,
+    "SGLDiffusionGenerateImage": SGLDiffusionGenerateImage,
+    "SGLDiffusionGenerateVideo": SGLDiffusionGenerateVideo,
+    "SGLDiffusionServerSetLora": SGLDiffusionServerSetLora,
+    "SGLDiffusionServerUnsetLora": SGLDiffusionServerUnsetLora,
+    "SGLDUNETLoader": SGLDUNETLoader,
+    "SGLDOptions": SGLDOptions,
+    "SGLDLoraLoader": SGLDLoraLoader,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SGLDiffusionServerModel": "SGLDiffusion Server Model",
+    "SGLDiffusionGenerateImage": "SGLDiffusion Generate Image",
+    "SGLDiffusionGenerateVideo": "SGLDiffusion Generate Video",
+    "SGLDiffusionServerSetLora": "SGLDiffusion Server Set LoRA",
+    "SGLDiffusionServerUnsetLora": "SGLDiffusion Server Unset LoRA",
+    "SGLDUNETLoader": "SGLDiffusion UNET Loader",
+    "SGLDOptions": "SGLDiffusion Options",
+    "SGLDLoraLoader": "SGLDiffusion LoRA Loader",
+}

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/README.md

@@ -0,0 +1,66 @@
+# ComfyUI SGLDiffusion Pipeline Tests
+
+This directory contains tests for each ComfyUI pipeline integration.
+
+## Test Files
+
+- `test_zimage_pipeline.py` - Tests for ComfyUIZImagePipeline
+- `test_flux_pipeline.py` - Tests for ComfyUIFluxPipeline
+- `test_qwen_image_pipeline.py` - Tests for ComfyUIQwenImagePipeline
+- `test_qwen_image_edit_pipeline.py` - Tests for ComfyUIQwenImageEditPipeline (I2I/edit mode)
+
+## Running Tests
+
+### Run all tests
+
+```bash
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/ -v -s
+```
+
+### Run a specific test file
+
+```bash
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+```
+
+## Environment Variables
+
+You can configure model paths via environment variables. Model paths support two formats:
+- **Safetensors file**: Path to a single `.safetensors` file (e.g., `/path/to/model.safetensors`)
+- **Diffusers format**: HuggingFace model ID or local diffusers directory (e.g., `Tongyi-MAI/Z-Image-Turbo`)
+
+Environment variables:
+- `SGLANG_TEST_ZIMAGE_MODEL_PATH` - Path to ZImage model (default: `Tongyi-MAI/Z-Image-Turbo`)
+- `SGLANG_TEST_FLUX_MODEL_PATH` - Path to Flux model (default: `black-forest-labs/FLUX.1-dev`)
+- `SGLANG_TEST_QWEN_IMAGE_MODEL_PATH` - Path to QwenImage model (default: `Qwen/Qwen-Image`)
+- `SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH` - Path to QwenImageEdit model (default: `Qwen/Qwen-Image-Edit-2511`)
+
+Examples:
+
+```bash
+# Using HuggingFace model ID (diffusers format)
+export SGLANG_TEST_ZIMAGE_MODEL_PATH="Tongyi-MAI/Z-Image-Turbo"
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+
+# Using safetensors file
+export SGLANG_TEST_ZIMAGE_MODEL_PATH="/path/to/z_image_turbo_bf16.safetensors"
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+```
+
+## Test Structure
+
+Each test file follows a similar structure:
+
+1. **Setup**: Creates a `DiffGenerator` with the appropriate pipeline class
+2. **Input Preparation**: Creates dummy tensors for latents, timesteps, and embeddings
+3. **Request Preparation**: Uses `prepare_request` to convert `SamplingParams` to `Req`
+4. **ComfyUI Inputs**: Sets ComfyUI-specific inputs directly on the `Req` object
+5. **Execution**: Sends request to scheduler and waits for response
+6. **Validation**: Checks that `noise_pred` is retrieved from `OutputBatch`
+
+## Notes
+
+- These tests use `comfyui_mode=True` to enable ComfyUI-specific behavior
+- Tests use pre-processed inputs (latents, timesteps, embeddings) as ComfyUI would provide
+- The tests verify that `noise_pred` can be retrieved from the `OutputBatch` after processing
+- All tests use dummy/ones tensors for simplicity - in production, these would be actual model outputs

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/__init__.py

@@ -0,0 +1,9 @@
+"""
+Test suite for ComfyUI SGLDiffusion pipelines.
+
+This package contains tests for each ComfyUI pipeline integration:
+- ZImagePipeline
+- FluxPipeline
+- QwenImagePipeline
+- QwenImageEditPipeline
+"""

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py

@@ -0,0 +1,156 @@
+"""Test for ComfyUIFluxPipeline with pass-through scheduler."""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_flux_pipeline_direct() -> None:
+    """Test ComfyUIFluxPipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_FLUX_MODEL_PATH",
+        "black-forest-labs/FLUX.1-dev",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIFluxPipeline",
+        num_gpus=2,
+        comfyui_mode=True,
+    )
+
+    batch_size = 1
+    hidden_states_seq_len = 3600
+    hidden_states_dim = 64
+    height = 1280
+    width = 720
+
+    encoder_seq_len = 512
+    encoder_dim = 4096
+    pooled_dim = 768
+
+    hidden_states = torch.ones(
+        batch_size,
+        hidden_states_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    pooled_projections = torch.ones(
+        batch_size,
+        pooled_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt="a beautiful girl",
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        save_output=True,
+        return_trajectory_latents=True,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = hidden_states
+    req.timesteps = timesteps
+    req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+    clip_dim = 768
+    req.prompt_embeds = [pooled_projections, encoder_hidden_states]
+
+    if req.guidance_scale > 1.0:
+        dummy_neg_clip_embedding = torch.zeros(
+            batch_size,
+            77,
+            clip_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        negative_encoder_hidden_states = torch.ones(
+            batch_size,
+            encoder_seq_len,
+            encoder_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        req.negative_prompt_embeds = [
+            dummy_neg_clip_embedding,
+            negative_encoder_hidden_states,
+        ]
+    else:
+        req.negative_prompt_embeds = None
+
+    req.pooled_embeds = [pooled_projections]
+    req.neg_pooled_embeds = []
+
+    if (
+        req.guidance_scale > 1.0
+        and req.negative_prompt_embeds is not None
+        and len(req.negative_prompt_embeds) > 0
+    ):
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cuda" if generator_device == "cuda" else "cpu"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print("✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+    print(f"latents.shape: {latents.shape}")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py

@@ -0,0 +1,136 @@
+"""Test for ComfyUIQwenImageEditPipeline with pass-through scheduler (I2I/edit mode)."""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_qwen_image_edit_pipeline_direct() -> None:
+    """Test ComfyUIQwenImageEditPipeline with edit mode (I2I) and custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH",
+        "Qwen/Qwen-Image-Edit-2511",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIQwenImageEditPipeline",
+        num_gpus=1,
+        comfyui_mode=True,
+        dit_layerwise_offload=False,
+    )
+
+    batch_size = 1
+    noisy_image_seq_len = 3600
+    hidden_states_dim = 64
+    condition_image_seq_len = 6889
+    condition_image_dim = 64
+    encoder_seq_len = 45
+    encoder_dim = 3584
+    height = 720
+    width = 1280
+
+    vae_scale_factor = 8
+    condition_height_latent = 1328 // vae_scale_factor
+    condition_width_latent = 1328 // vae_scale_factor
+
+    noisy_image_latents = torch.ones(
+        batch_size,
+        noisy_image_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    condition_image_latents = torch.ones(
+        batch_size,
+        condition_image_seq_len,
+        condition_image_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt=" ",
+        guidance_scale=1.0,
+        height=height,
+        width=width,
+        image_path="",
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = noisy_image_latents
+    req.image_latent = condition_image_latents
+    req.timesteps = timesteps
+    req.prompt_embeds = [encoder_hidden_states]
+    req.negative_prompt_embeds = None
+    req.vae_image_sizes = [(condition_width_latent, condition_height_latent)]
+    req.raw_latent_shape = torch.tensor(noisy_image_latents.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print("✓ Successfully retrieved noise_pred from OutputBatch (Edit Mode)!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py

@@ -0,0 +1,120 @@
+"""Test for ComfyUIQwenImagePipeline with pass-through scheduler."""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_qwen_image_pipeline_direct() -> None:
+    """Test ComfyUIQwenImagePipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_QWEN_IMAGE_MODEL_PATH",
+        "Qwen/Qwen-Image",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIQwenImagePipeline",
+        num_gpus=2,
+        comfyui_mode=True,
+        dit_layerwise_offload=False,
+    )
+
+    batch_size = 1
+    hidden_states_seq_len = 6889
+    hidden_states_dim = 64
+    encoder_seq_len = 45
+    encoder_dim = 3584
+    height = 1328
+    width = 1328
+    dtype = torch.bfloat16
+
+    hidden_states = torch.ones(
+        batch_size,
+        hidden_states_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=dtype,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt=" ",
+        guidance_scale=3.0,
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = hidden_states
+    req.timesteps = timesteps
+    req.prompt_embeds = [encoder_hidden_states]
+    req.negative_prompt_embeds = [encoder_hidden_states]
+    req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print("✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py

@@ -0,0 +1,122 @@
+"""Test for ComfyUIZImagePipeline with pass-through scheduler."""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_zimage_pipeline_direct() -> None:
+    """Test ComfyUIZImagePipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_ZIMAGE_MODEL_PATH",
+        "Tongyi-MAI/Z-Image-Turbo",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIZImagePipeline",
+        num_gpus=1,
+        sp_degree=1,
+        comfyui_mode=True,
+    )
+
+    batch_size = 1
+    num_channels = 16
+    num_frames = 1
+    height = 720
+    width = 1280
+    latent_height = height // 8
+    latent_width = width // 8
+
+    latents = torch.ones(
+        batch_size,
+        num_channels,
+        num_frames,
+        latent_height,
+        latent_width,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    context_seq_len = 19
+    context_dim = 2560
+    context = torch.ones(
+        context_seq_len,
+        context_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt="a beautiful girl",
+        guidance_scale=1.0,
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = latents
+    req.timesteps = timesteps
+    req.prompt_embeds = [context]
+    req.negative_prompt_embeds = None
+    req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print("✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/utils.py

@@ -0,0 +1,176 @@
+import base64
+import io
+import os
+import shutil
+import time
+import uuid
+
+import folder_paths
+import numpy as np
+import torch
+from comfy_api.input import VideoInput
+from PIL import Image
+
+
+def _ensure_dir(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def _to_numpy_image(image: torch.Tensor) -> np.ndarray:
+    """Convert ComfyUI image tensor to uint8 numpy array (H, W, C)."""
+    if image.dim() == 4:
+        image = image[0]
+    if image.dim() == 3 and image.shape[0] in (1, 3, 4):
+        image = image.permute(1, 2, 0)
+    elif image.dim() == 2:
+        image = image.unsqueeze(-1)
+    np_img = image.detach().cpu().numpy()
+    np_img = np.clip(np_img, 0.0, 1.0)
+    np_img = (np_img * 255).astype(np.uint8)
+    if np_img.shape[-1] == 1:
+        np_img = np.repeat(np_img, 3, axis=-1)
+    return np_img
+
+
+def _to_hwc_tensor(image: torch.Tensor) -> torch.Tensor:
+    """Convert ComfyUI image tensor to HWC format (normalized [0, 1])."""
+    img = image.clone()
+    if img.dim() == 4:
+        img = img[0]
+    if img.dim() == 3 and img.shape[0] in (1, 3, 4):
+        img = img.permute(1, 2, 0)
+    elif img.dim() == 2:
+        img = img.unsqueeze(-1)
+
+    img = torch.clamp(img, 0.0, 1.0)
+    if img.shape[-1] == 1:
+        img = img.repeat(1, 1, 3)
+
+    return img
+
+
+def is_empty_image(image: torch.Tensor, tolerance: float = 1e-6) -> bool:
+    """
+    Check if the input image is an empty/solid color image (like ComfyUI's empty image).
+    Args:
+        image: Input tensor image in ComfyUI format (BCHW, CHW, HWC, etc.)
+        tolerance: Tolerance for floating point comparison (default: 1e-6)
+
+    Returns:
+        True if the image is empty (all pixels have same color), False otherwise
+    """
+    if image is None:
+        return True
+
+    # Convert to HWC format
+    img_hwc = _to_hwc_tensor(image)
+
+    # Get the first pixel's RGB values
+    first_pixel = img_hwc[0, 0, :]
+
+    h, w, c = img_hwc.shape
+    pixels = img_hwc.reshape(-1, c)
+
+    diff = torch.abs(pixels - first_pixel)
+    max_diff = torch.max(diff)
+
+    return max_diff.item() <= tolerance
+
+
+def get_image_path(image: torch.Tensor) -> str:
+    """
+    Save tensor image to ComfyUI temp directory as PNG and return the path.
+    """
+    temp_dir = folder_paths.get_temp_directory()
+
+    # Build file name
+    ts = time.strftime("%Y%m%d-%H%M%S")
+    unique = uuid.uuid4().hex[:8]
+    file_name = f"sgl_output_{ts}_{unique}.png"
+    file_path = os.path.join(temp_dir, file_name)
+
+    # Save image
+    np_img = _to_numpy_image(image)
+    img = Image.fromarray(np_img)
+    img.save(file_path, format="PNG")
+
+    return file_path
+
+
+def convert_b64_to_tensor_image(b64_image: str) -> torch.Tensor:
+    """
+    Convert base64 encoded image to ComfyUI IMAGE format (torch.Tensor).
+
+    Args:
+        b64_image: Base64 encoded image string
+
+    Returns:
+        torch.Tensor with shape [batch_size, height, width, channels] (BHWC format),
+        values normalized to [0, 1] range, RGB format (3 channels)
+    """
+    # Decode base64
+    image_bytes = base64.b64decode(b64_image)
+
+    # Open image and convert to RGB
+    pil_image = Image.open(io.BytesIO(image_bytes))
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+
+    # Convert to numpy array and normalize to [0, 1]
+    image_array = np.array(pil_image).astype(np.float32) / 255.0
+
+    # Add batch dimension: [height, width, channels] -> [1, height, width, channels]
+    image_array = image_array[np.newaxis, ...]
+
+    # Convert to torch.Tensor
+    tensor_image = torch.from_numpy(image_array)
+
+    return tensor_image
+
+
+class SGLDVideoInput(VideoInput):
+    def __init__(self, video_path: str, height: int, width: int):
+        super().__init__()
+
+        self.video_path = video_path
+        self.height = height
+        self.width = width
+
+    def get_dimensions(self) -> tuple[int, int]:
+        """
+        Returns the dimensions of the video input.
+
+        Returns:
+            Tuple of (width, height)
+        """
+        return self.width, self.height
+
+    def get_components(self):
+        """
+        Returns the components of the video input.
+        This is required by the VideoInput abstract base class.
+        """
+        return [self.video_path]
+
+    def save_to(self, path: str, format=None, codec=None, metadata=None):
+        """
+        Abstract method to save the video input to a file.
+        """
+        save_path = path
+        # Copy video file from video_path to save_path
+        if os.path.exists(self.video_path):
+            # Ensure destination directory exists
+            save_dir = os.path.dirname(save_path)
+            if save_dir:
+                os.makedirs(save_dir, exist_ok=True)
+            shutil.copy2(self.video_path, save_path)
+
+
+def convert_video_to_comfy_video(
+    video_path: str, height: int, width: int
+) -> VideoInput:
+    """
+    Convert video to ComfyUI VIDEO format (VideoInput).
+    """
+    video_input = SGLDVideoInput(video_path, height, width)
+    return video_input

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/flux_sgld_sp.json

@@ -0,0 +1,222 @@
+{
+    "8": {
+      "inputs": {
+        "samples": [
+          "40",
+          0
+        ],
+        "vae": [
+          "10",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "10": {
+      "inputs": {
+        "vae_name": "ae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "Load VAE"
+      }
+    },
+    "11": {
+      "inputs": {
+        "clip_name1": "t5xxl_fp16.safetensors",
+        "clip_name2": "clip_l.safetensors",
+        "type": "flux",
+        "device": "default"
+      },
+      "class_type": "DualCLIPLoader",
+      "_meta": {
+        "title": "DualCLIPLoader"
+      }
+    },
+    "17": {
+      "inputs": {
+        "scheduler": "normal",
+        "steps": 25,
+        "denoise": 1,
+        "model": [
+          "46",
+          0
+        ]
+      },
+      "class_type": "BasicScheduler",
+      "_meta": {
+        "title": "BasicScheduler"
+      }
+    },
+    "38": {
+      "inputs": {
+        "model": [
+          "46",
+          0
+        ],
+        "conditioning": [
+          "42",
+          0
+        ]
+      },
+      "class_type": "BasicGuider",
+      "_meta": {
+        "title": "BasicGuider"
+      }
+    },
+    "39": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "8",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "40": {
+      "inputs": {
+        "noise": [
+          "45",
+          0
+        ],
+        "guider": [
+          "38",
+          0
+        ],
+        "sampler": [
+          "47",
+          0
+        ],
+        "sigmas": [
+          "17",
+          0
+        ],
+        "latent_image": [
+          "44",
+          0
+        ]
+      },
+      "class_type": "SamplerCustomAdvanced",
+      "_meta": {
+        "title": "SamplerCustomAdvanced"
+      }
+    },
+    "42": {
+      "inputs": {
+        "guidance": 3.5,
+        "conditioning": [
+          "43",
+          0
+        ]
+      },
+      "class_type": "FluxGuidance",
+      "_meta": {
+        "title": "FluxGuidance"
+      }
+    },
+    "43": {
+      "inputs": {
+        "text": "beautiful photography of a gonger haired artist with Lots of Colorful coloursplashes in face and pn her hands, she is natural, having her hair in a casual bun, looking happily into camera, cinematic,",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "11",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Prompt)"
+      }
+    },
+    "44": {
+      "inputs": {
+        "width": 1024,
+        "height": 1024,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "45": {
+      "inputs": {
+        "noise_seed": 747172083610812
+      },
+      "class_type": "RandomNoise",
+      "_meta": {
+        "title": "RandomNoise"
+      }
+    },
+    "46": {
+      "inputs": {
+        "max_shift": 1.15,
+        "base_shift": 0.5,
+        "width": 1024,
+        "height": 1024,
+        "model": [
+          "51",
+          0
+        ]
+      },
+      "class_type": "ModelSamplingFlux",
+      "_meta": {
+        "title": "ModelSamplingFlux"
+      }
+    },
+    "47": {
+      "inputs": {
+        "sampler_name": "euler"
+      },
+      "class_type": "KSamplerSelect",
+      "_meta": {
+        "title": "KSamplerSelect"
+      }
+    },
+    "51": {
+      "inputs": {
+        "unet_name": "flux1-dev.safetensors",
+        "weight_dtype": "default",
+        "sgld_options": [
+          "52",
+          0
+        ]
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    },
+    "52": {
+      "inputs": {
+        "model_type": "auto-detect",
+        "enable_torch_compile": false,
+        "num_gpus": 2,
+        "tp_size": -1,
+        "sp_degree": -1,
+        "ulysses_degree": -1,
+        "ring_degree": -1,
+        "dp_size": 1,
+        "dp_degree": 1,
+        "enable_cfg_parallel": false,
+        "attention_backend": "",
+        "cache_strategy": "none"
+      },
+      "class_type": "SGLDOptions",
+      "_meta": {
+        "title": "SGLDiffusion Options"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/qwen_image_sgld.json

@@ -0,0 +1,165 @@
+{
+    "3": {
+      "inputs": {
+        "seed": 808633539418610,
+        "steps": 4,
+        "cfg": 1,
+        "sampler_name": "euler",
+        "scheduler": "simple",
+        "denoise": 1,
+        "model": [
+          "66",
+          0
+        ],
+        "positive": [
+          "6",
+          0
+        ],
+        "negative": [
+          "7",
+          0
+        ],
+        "latent_image": [
+          "58",
+          0
+        ]
+      },
+      "class_type": "KSampler",
+      "_meta": {
+        "title": "KSampler"
+      }
+    },
+    "6": {
+      "inputs": {
+        "text": "\"A vibrant, warm neon-lit street scene in Hong Kong at the afternoon, with a mix of colorful Chinese and English signs glowing brightly. The atmosphere is lively, cinematic, and rain-washed with reflections on the pavement. The colors are vivid, full of pink, blue, red, and green hues. Crowded buildings with overlapping neon signs. 1980s Hong Kong style. Signs include:\n\"龍鳳冰室\" \"金華燒臘\" \"HAPPY HAIR\" \"鴻運茶餐廳\" \"EASY BAR\" \"永發魚蛋粉\" \"添記粥麵\" \"SUNSHINE MOTEL\" \"美都餐室\" \"富記糖水\" \"太平館\" \"雅芳髮型屋\" \"STAR KTV\" \"銀河娛樂城\" \"百樂門舞廳\" \"BUBBLE CAFE\" \"萬豪麻雀館\" \"CITY LIGHTS BAR\" \"瑞祥香燭莊\" \"文記文具\" \"GOLDEN JADE HOTEL\" \"LOVELY BEAUTY\" \"合興百貨\" \"興旺電器\" And the background is warm yellow street and with all stores' lights on.",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "38",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Positive Prompt)"
+      }
+    },
+    "7": {
+      "inputs": {
+        "text": "",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "38",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Negative Prompt)"
+      }
+    },
+    "8": {
+      "inputs": {
+        "samples": [
+          "3",
+          0
+        ],
+        "vae": [
+          "39",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "38": {
+      "inputs": {
+        "clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors",
+        "type": "qwen_image",
+        "device": "default"
+      },
+      "class_type": "CLIPLoader",
+      "_meta": {
+        "title": "Load CLIP"
+      }
+    },
+    "39": {
+      "inputs": {
+        "vae_name": "qwen_image_vae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "Load VAE"
+      }
+    },
+    "58": {
+      "inputs": {
+        "width": 1328,
+        "height": 1328,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "60": {
+      "inputs": {
+        "filename_prefix": "ComfyUI"
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "66": {
+      "inputs": {
+        "shift": 3.1000000000000005,
+        "model": [
+          "78",
+          0
+        ]
+      },
+      "class_type": "ModelSamplingAuraFlow",
+      "_meta": {
+        "title": "ModelSamplingAuraFlow"
+      }
+    },
+    "77": {
+      "inputs": {
+        "unet_name": "qwen_image_2512_bf16.safetensors",
+        "weight_dtype": "default"
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    },
+    "78": {
+      "inputs": {
+        "lora_name": "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors",
+        "strength_model": 1,
+        "nickname": "",
+        "target": "all",
+        "model": [
+          "77",
+          0
+        ]
+      },
+      "class_type": "SGLDLoraLoader",
+      "_meta": {
+        "title": "SGLDiffusion LoRA Loader"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_image2video.json

@@ -0,0 +1,97 @@
+{
+    "1": {
+      "inputs": {
+        "base_url": "http://localhost:3000/v1",
+        "api_key": "sk-proj-1234567890"
+      },
+      "class_type": "SGLDiffusionServerModel",
+      "_meta": {
+        "title": "SGLDiffusion Server Model"
+      }
+    },
+    "3": {
+      "inputs": {
+        "prompt": "The girl turn the body and spin around in place.",
+        "main": "none",
+        "lighting": "none",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        }
+      },
+      "class_type": "easy prompt",
+      "_meta": {
+        "title": "Prompt"
+      }
+    },
+    "4": {
+      "inputs": {
+        "text": "",
+        "anything": [
+          "1",
+          1
+        ]
+      },
+      "class_type": "easy showAnything",
+      "_meta": {
+        "title": "Show Any"
+      }
+    },
+    "15": {
+      "inputs": {
+        "positive_prompt": [
+          "3",
+          0
+        ],
+        "negative_prompt": "",
+        "seed": 2435791308,
+        "steps": 50,
+        "cfg": 4,
+        "width": 704,
+        "height": 1280,
+        "num_frames": 16,
+        "fps": 16,
+        "seconds": 1,
+        "enable_teacache": false,
+        "sgld_client": [
+          "1",
+          0
+        ],
+        "image": [
+          "17",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionGenerateVideo",
+      "_meta": {
+        "title": "SGLDiffusion Generate Video"
+      }
+    },
+    "16": {
+      "inputs": {
+        "filename_prefix": "video/ComfyUI",
+        "format": "auto",
+        "codec": "auto",
+        "video-preview": "",
+        "video": [
+          "15",
+          0
+        ]
+      },
+      "class_type": "SaveVideo",
+      "_meta": {
+        "title": "save video"
+      }
+    },
+    "17": {
+      "inputs": {
+        "image": "tmpe_w0bd_0.jpg"
+      },
+      "class_type": "LoadImage",
+      "_meta": {
+        "title": "load image"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_text2img.json

@@ -0,0 +1,109 @@
+{
+    "1": {
+      "inputs": {
+        "base_url": "http://localhost:3000/v1",
+        "api_key": "sk-proj-1234567890"
+      },
+      "class_type": "SGLDiffusionServerModel",
+      "_meta": {
+        "title": "SGLDiffusion Server Model"
+      }
+    },
+    "3": {
+      "inputs": {
+        "prompt": "a bicycle, illustration in the style of SMPL, thick black lines on a white background",
+        "main": "none",
+        "lighting": "none",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        }
+      },
+      "class_type": "easy prompt",
+      "_meta": {
+        "title": "Prompt"
+      }
+    },
+    "4": {
+      "inputs": {
+        "text": "",
+        "anything": [
+          "1",
+          1
+        ]
+      },
+      "class_type": "easy showAnything",
+      "_meta": {
+        "title": "Show Any"
+      }
+    },
+    "5": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "6",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "save image"
+      }
+    },
+    "6": {
+      "inputs": {
+        "positive_prompt": [
+          "3",
+          0
+        ],
+        "negative_prompt": "",
+        "seed": 4215918563,
+        "steps": 50,
+        "cfg": 4,
+        "width": 512,
+        "height": 512,
+        "enable_teacache": false,
+        "sgld_client": [
+          "11",
+          0
+        ],
+        "image": [
+          "14",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionGenerateImage",
+      "_meta": {
+        "title": "SGLDiffusion Generate Image"
+      }
+    },
+    "11": {
+      "inputs": {
+        "lora_name": "dvyio/flux-lora-simple-illustration",
+        "lora_nickname": "",
+        "target": "all",
+        "sgld_client": [
+          "1",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionSetLora",
+      "_meta": {
+        "title": "SGLDiffusion Set LoRA"
+      }
+    },
+    "14": {
+      "inputs": {
+        "width": 512,
+        "height": 512,
+        "batch_size": 1,
+        "color": 0
+      },
+      "class_type": "EmptyImage",
+      "_meta": {
+        "title": "empty image"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/z-image_sgld.json

@@ -0,0 +1,140 @@
+{
+    "3": {
+      "inputs": {
+        "seed": 3338398,
+        "steps": 9,
+        "cfg": 1,
+        "sampler_name": "euler",
+        "scheduler": "simple",
+        "denoise": 1,
+        "model": [
+          "28",
+          0
+        ],
+        "positive": [
+          "6",
+          0
+        ],
+        "negative": [
+          "7",
+          0
+        ],
+        "latent_image": [
+          "13",
+          0
+        ]
+      },
+      "class_type": "KSampler",
+      "_meta": {
+        "title": "KSampler"
+      }
+    },
+    "6": {
+      "inputs": {
+        "text": "cute anime style girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black gold leaf pattern dress and a white apron, it is a postcard held by a hand in front of a beautiful realistic city at sunset and there is cursive writing that says \"ZImage, Now in ComfyUI\"",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "18",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Positive Prompt)"
+      }
+    },
+    "7": {
+      "inputs": {
+        "text": "blurry ugly bad",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "18",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Negative Prompt)"
+      }
+    },
+    "8": {
+      "inputs": {
+        "samples": [
+          "3",
+          0
+        ],
+        "vae": [
+          "17",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "9": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "8",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "13": {
+      "inputs": {
+        "width": 1024,
+        "height": 1024,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "17": {
+      "inputs": {
+        "vae_name": "ae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "VAE Loader"
+      }
+    },
+    "18": {
+      "inputs": {
+        "clip_name": "qwen_3_4b.safetensors",
+        "type": "lumina2",
+        "device": "default"
+      },
+      "class_type": "CLIPLoader",
+      "_meta": {
+        "title": "CLIP Loader"
+      }
+    },
+    "28": {
+      "inputs": {
+        "unet_name": "z_image_turbo_bf16.safetensors",
+        "weight_dtype": "default"
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    }
+}

python/sglang/multimodal_gen/apps/realtime_webui/README.md

@@ -0,0 +1,24 @@
+# SGLang Diffusion Realtime WebUI
+
+Standalone browser demo for `/v1/realtime_video/generate`.
+
+Open `index.html` directly in a browser, point it at an SGLang Diffusion server,
+and generate. The app sends msgpack init / event messages and renders lossless
+raw RGB frame batches on a canvas.
+
+The first version is intentionally static: no npm install, no build step, and no
+server-side dependencies. Presets are UI-side templates for prompt, LingBot
+example images, album artwork references, and session parameters. The default
+preset preloads a reference image so the demo can be tested without a file
+upload.
+
+By default, `Continuous session` is enabled for long-running camera control.
+Keyboard and pointer controls send state transitions instead of scripted preset
+actions. The telemetry `Chunk wait` measures request-to-chunk arrival time, not
+client-side RGB decode time. Continuous playback adapts to the measured chunk
+production rate so the canvas does not play a chunk at target FPS and then sit
+on the last frame while waiting for the next chunk.
+
+The interface shape follows camera-control-first video playgrounds such as
+Reactor LingBot: reference image, scene prompt, enhancement, clip controls,
+move/look camera controls, recordings history, and model telemetry.

python/sglang/multimodal_gen/apps/realtime_webui/decoder_worker.js

@@ -0,0 +1,172 @@
+const RAW_RGB_CONTENT_TYPE = "application/x-raw-rgb";
+const RAW_RGB_DELTA_GZIP_CONTENT_TYPE = "application/x-raw-rgb-delta-gzip";
+const RAW_RGBA_DELTA_GZIP_CONTENT_TYPE = "application/x-raw-rgba-delta-gzip";
+const WEBP_FRAME_CONTENT_TYPE = "image/webp";
+const JPEG_FRAME_CONTENT_TYPE = "image/jpeg";
+
+let lastFrame = null;
+
+function reset() {
+  lastFrame = null;
+}
+
+async function gunzipBytes(payload) {
+  if (typeof DecompressionStream === "undefined") {
+    throw new Error("This browser does not support gzip stream decoding");
+  }
+  const stream = new Blob([payload]).stream().pipeThrough(new DecompressionStream("gzip"));
+  return new Uint8Array(await new Response(stream).arrayBuffer());
+}
+
+async function restoreDeltaGzipFrames(header, payload) {
+  const frameBytes = Number(header.bytes_per_frame);
+  const count = Number(header.num_frames);
+  const expectedSize = frameBytes * count;
+  const restored = await gunzipBytes(payload);
+  if (restored.length !== expectedSize) {
+    throw new Error(`delta payload size mismatch: expected ${expectedSize}, got ${restored.length}`);
+  }
+
+  let previous = header.delta_reference === "previous-frame" ? lastFrame : null;
+  if (header.delta_reference === "previous-frame") {
+    if (!previous) throw new Error("Missing previous frame for delta payload");
+    if (previous.byteLength !== frameBytes) {
+      throw new Error("Previous frame size does not match current delta payload");
+    }
+  }
+
+  for (let f = 0; f < count; f++) {
+    const offset = f * frameBytes;
+    if (previous) {
+      for (let i = 0; i < frameBytes; i++) restored[offset + i] ^= previous[i];
+    }
+    previous = restored.slice(offset, offset + frameBytes);
+  }
+  lastFrame = previous;
+  return restored;
+}
+
+function rawFramesToRgbaBuffers(header, payload) {
+  const width = Number(header.width);
+  const height = Number(header.height);
+  const channels = Number(header.channels);
+  const count = Number(header.num_frames);
+  const frameBytes = Number(header.bytes_per_frame);
+  const pixels = width * height;
+  const buffers = [];
+
+  for (let f = 0; f < count; f++) {
+    const offset = f * frameBytes;
+    if (channels === 4) {
+      buffers.push(payload.buffer.slice(
+        payload.byteOffset + offset,
+        payload.byteOffset + offset + frameBytes,
+      ));
+      continue;
+    }
+
+    const rgba = new Uint8ClampedArray(pixels * 4);
+    let src = offset;
+    let dst = 0;
+    for (let p = 0; p < pixels; p++) {
+      rgba[dst++] = payload[src++];
+      rgba[dst++] = payload[src++];
+      rgba[dst++] = payload[src++];
+      src += channels - 3;
+      rgba[dst++] = 255;
+    }
+    buffers.push(rgba.buffer);
+  }
+  return buffers;
+}
+
+function splitEncodedPayload(header, payload) {
+  const bytes = payload instanceof Uint8Array ? payload : new Uint8Array(payload);
+  const lengths = Array.isArray(header.payload_lengths) && header.payload_lengths.length
+    ? header.payload_lengths.map(Number)
+    : [bytes.byteLength];
+  const payloads = [];
+  let offset = 0;
+  for (const length of lengths) {
+    payloads.push(bytes.buffer.slice(
+      bytes.byteOffset + offset,
+      bytes.byteOffset + offset + length,
+    ));
+    offset += length;
+  }
+  return payloads;
+}
+
+async function encodedFramesToImageBitmaps(header, payload) {
+  if (typeof createImageBitmap === "undefined") {
+    throw new Error("This browser does not support worker image decoding");
+  }
+
+  const frames = await Promise.all(splitEncodedPayload(header, payload).map((framePayload) => (
+    createImageBitmap(new Blob([framePayload], { type: header.content_type }))
+  )));
+  return {
+    width: frames[0]?.width || 0,
+    height: frames[0]?.height || 0,
+    frame_type: "bitmap",
+    frames,
+  };
+}
+
+async function decode(header, payload) {
+  let rawPayload;
+  if (
+    header.content_type === WEBP_FRAME_CONTENT_TYPE ||
+    header.content_type === JPEG_FRAME_CONTENT_TYPE
+  ) {
+    const decoded = await encodedFramesToImageBitmaps(header, payload);
+    return {
+      id: header.__decode_id,
+      width: decoded.width,
+      height: decoded.height,
+      chunk: Number(header.chunk_index),
+      frame_type: decoded.frame_type,
+      frames: decoded.frames,
+    };
+  } else if (header.content_type === RAW_RGB_CONTENT_TYPE) {
+    rawPayload = new Uint8Array(payload);
+    const frameBytes = Number(header.bytes_per_frame);
+    const count = Number(header.num_frames);
+    lastFrame = count > 0
+      ? rawPayload.slice((count - 1) * frameBytes, count * frameBytes)
+      : null;
+  } else if (
+    header.content_type === RAW_RGB_DELTA_GZIP_CONTENT_TYPE ||
+    header.content_type === RAW_RGBA_DELTA_GZIP_CONTENT_TYPE
+  ) {
+    rawPayload = await restoreDeltaGzipFrames(header, payload);
+  } else {
+    throw new Error(`Unsupported content type ${header.content_type}`);
+  }
+
+  return {
+    id: header.__decode_id,
+    width: Number(header.width),
+    height: Number(header.height),
+    chunk: Number(header.chunk_index),
+    frames: rawFramesToRgbaBuffers(header, rawPayload),
+  };
+}
+
+self.onmessage = async (event) => {
+  const message = event.data;
+  try {
+    if (message.type === "reset") {
+      reset();
+      return;
+    }
+    const result = await decode(message.header, message.payload);
+    self.postMessage({ type: "decoded", ...result }, result.frames);
+  } catch (error) {
+    self.postMessage({
+      type: "error",
+      id: message.header?.__decode_id,
+      message: error.message || "decode failed",
+    });
+  }
+};

python/sglang/multimodal_gen/apps/realtime_webui/index.html

@@ -0,0 +1,168 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <title>sglang-diffusion Realtime Studio</title>
+    <link rel="stylesheet" href="./styles.css?v=realtime-record-v49" />
+  </head>
+  <body>
+    <main class="shell">
+      <section class="panel controls" aria-label="Session controls">
+        <div class="brand">
+          <span>sglang-diffusion</span>
+          <strong>Realtime Studio</strong>
+        </div>
+        <label>Server<input id="serverUrl" value="ws://127.0.0.1:30000/v1/realtime_video/generate" /></label>
+        <label>Model<input id="model" value="" placeholder="auto from /v1/models" /></label>
+        <div class="section-title">Reference</div>
+        <label class="reference-upload">
+          <input id="firstFrame" type="file" accept="image/*" />
+          <canvas id="referencePreview" width="320" height="180"></canvas>
+          <span id="referenceName">Preset reference</span>
+        </label>
+        <div class="section-title">Generate the scene</div>
+        <label>Prompt<textarea id="prompt" rows="4">A cinematic handheld shot of a quiet city street at dusk, soft reflections, natural motion.</textarea></label>
+        <button id="enhanceBtn" class="wide">Enhance</button>
+        <div class="split">
+          <label>Size<input id="size" value="832x480" /></label>
+          <label>FPS<input id="fps" type="number" value="25" min="1" max="60" /></label>
+        </div>
+        <div class="split">
+          <label>Frames<input id="numFrames" type="number" value="9" min="5" step="4" /></label>
+          <label>Seed<input id="seed" type="number" value="42" /></label>
+        </div>
+        <div class="split">
+          <label>Steps<input id="steps" type="number" value="4" min="1" /></label>
+          <label>Guidance<input id="guidance" type="number" value="1" step="0.1" /></label>
+        </div>
+        <div class="split">
+          <label>Sink<input id="sinkSize" type="number" value="9" min="0" /></label>
+          <label>Window<input id="windowFrames" type="number" value="18" min="1" /></label>
+        </div>
+        <div class="split">
+          <label>Transport
+            <select id="transportFormat">
+              <option value="webp" selected>WebP preview</option>
+              <option value="jpeg">JPEG preview</option>
+              <option value="">Lossless delta</option>
+              <option value="raw">Raw RGB</option>
+            </select>
+          </label>
+          <label>Quality<input id="transportQuality" type="number" value="95" min="1" max="100" /></label>
+        </div>
+        <div class="split output-options">
+          <label class="toggle-row"><input id="superResolution" type="checkbox" />Super resolution</label>
+          <label>Scale
+            <select id="upscalingScale">
+              <option value="2" selected>2x</option>
+              <option value="4">4x</option>
+            </select>
+          </label>
+        </div>
+        <label>SR model
+          <select id="upscalingModel">
+            <option value="">Quality x2</option>
+            <option
+              value="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
+              selected
+            >
+              Fast general
+            </option>
+            <option value="/scratch/realesr-animevideov3.pth">Fast anime</option>
+          </select>
+        </label>
+        <label class="toggle-row"><input id="frameInterpolation" type="checkbox" />Smooth 2x frames</label>
+        <label class="toggle-row"><input id="continuous" type="checkbox" checked />Continuous session</label>
+        <div class="actions">
+          <button id="connectBtn" class="primary">Generate</button>
+          <button id="stopBtn">Close session</button>
+        </div>
+        <button id="sendPromptBtn" class="wide">Send prompt update</button>
+      </section>
+
+      <section class="workspace" aria-label="Realtime workspace">
+        <section class="stage" aria-label="Realtime preview">
+          <div class="topbar">
+            <span id="statusDot" class="dot"></span>
+            <span id="statusText">Idle</span>
+            <span id="chunkText">chunk -</span>
+            <button id="recordBtn" class="record-button" type="button" aria-pressed="false" title="Record preview">
+              <span class="record-button-icon" aria-hidden="true"></span>
+              <span id="recordLabel">Record</span>
+              <span id="recordDuration" class="record-button-duration">00:00</span>
+            </button>
+            <span class="topbar-spacer"></span>
+            <label class="preview-scale-control">Preview
+              <input id="previewScale" type="range" min="80" max="170" value="120" />
+              <b id="previewScaleText">120%</b>
+            </label>
+            <span class="stage-stat">output <b id="outputSizeText">832x480</b></span>
+            <span class="stage-stat">render <b id="renderFps">0</b> fps</span>
+            <span class="stage-stat">source <b id="theoreticalFpsText">-</b></span>
+            <span class="stage-stat">buffer <b id="stageLatencyText">-</b></span>
+          </div>
+          <div class="preview-frame">
+            <canvas id="viewport" width="1280" height="720"></canvas>
+            <div id="previewOverlay" class="preview-overlay" aria-hidden="true">
+              <span class="preview-loader"></span>
+            </div>
+          </div>
+          <div class="stage-controls" aria-label="Camera controls">
+            <div class="control-cluster" aria-label="Move camera">
+              <span class="control-title">Move</span>
+              <div class="camera-pad move-pad">
+                <span></span>
+                <button data-action="w" data-key="W">Forward</button>
+                <span></span>
+                <button data-action="a" data-key="A">Left</button>
+                <button data-action="s" data-key="S">Back</button>
+                <button data-action="d" data-key="D">Right</button>
+              </div>
+            </div>
+            <div class="control-cluster" aria-label="Look around">
+              <span class="control-title">Look</span>
+              <div class="camera-pad look-pad">
+                <span></span>
+                <button data-action="i" data-key="↑">Pitch +</button>
+                <span></span>
+                <button data-action="j" data-key="←">Yaw -</button>
+                <button data-action="k" data-key="↓">Pitch -</button>
+                <button data-action="l" data-key="→">Yaw +</button>
+              </div>
+            </div>
+          </div>
+          <div class="timeline">
+            <span id="queueText">queue 0</span>
+            <span id="frameText">frames 0</span>
+            <span id="byteText">0 MB</span>
+          </div>
+          <div class="telemetry stage-telemetry">
+            <span>Payload<b id="payloadMode">webp</b></span>
+            <span>Server send<b id="serverSendText">-</b></span>
+            <span>Chunk bytes<b id="chunkPayloadText">-</b></span>
+            <span>Chunk wait<b id="latencyText">-</b></span>
+            <span>Decode<b id="decodeText">-</b></span>
+            <span>Display lag<b id="displayLagText">-</b></span>
+          </div>
+        </section>
+
+        <section class="panel presets" aria-label="Presets and camera">
+          <div class="section-title">LingBot</div>
+          <div class="spec-grid">
+            <span><b>25 fps</b> target</span>
+            <span><b>chunked</b> stream</span>
+            <span><b>480p/720p</b></span>
+            <span><b>Cam + Act</b></span>
+          </div>
+          <div class="section-title">Presets</div>
+          <div id="presetList" class="preset-list"></div>
+          <div class="section-title">History</div>
+          <div id="historyList" class="history-list"></div>
+        </section>
+      </section>
+    </main>
+  <script src="./playback_controller.js?v=realtime-playback-v13"></script>
+  <script src="./app.js?v=realtime-record-v75"></script>
+</body>
+</html>

python/sglang/multimodal_gen/apps/realtime_webui/playback_controller.js

@@ -0,0 +1,517 @@
+(function attachRealtimePlaybackController(global) {
+  const DEFAULT_CONFIG = {
+    targetFps: 25,
+    minSourceFps: 1,
+    serverFpsAlphaUp: 0.28,
+    serverFpsAlphaDown: 0.2,
+    deliveryFpsAlphaUp: 0.08,
+    deliveryFpsAlphaDown: 0.55,
+    targetLeadChunkRatio: 1.5,
+    minTargetLeadMs: 1500,
+    maxTargetLeadMs: 2600,
+    maxLeadExtraChunkRatio: 8.0,
+    startLeadChunkRatio: 1.85,
+    minStartLeadMs: 1700,
+    resumeLeadChunkRatio: 2.5,
+    minResumeLeadMs: 1000,
+    maxResumeLeadMs: 1800,
+    rebufferLeadBoostMs: 250,
+    rebufferLeadBoostDecayMsPerSecond: 120,
+    deliveryLeadBoostDecayMsPerSecond: 80,
+    maxDeliveryLeadBoostMs: 2000,
+    deliveryStallExpectedMultiplier: 1.25,
+    receiveStallPlaybackRateMin: 0.65,
+    receiveStallPlaybackRateSlewPerSecond: 0.5,
+    lowWaterRatio: 0.4,
+    playbackRateGain: 0.14,
+    playbackRateMin: 0.92,
+    playbackRateMax: 1.08,
+    emergencyPlaybackRateMin: 0.9,
+    emergencyPlaybackRateMax: 1.12,
+    playbackRateSlewPerSecond: 0.08,
+    eventCutoverMaxMs: 420,
+    eventCutoverMaxFrames: 10,
+    settleEventCutoverMaxMs: 720,
+    settleEventCutoverMaxFrames: 18,
+    startupWarmupMinMs: 1500,
+    startupWarmupExpectedMultiplier: 3,
+  };
+
+  function clamp(value, min, max) {
+    return Math.min(max, Math.max(min, value));
+  }
+
+  function finitePositive(value) {
+    return Number.isFinite(value) && value > 0;
+  }
+
+  class RealtimePlaybackController {
+    constructor(config = {}) {
+      this.config = { ...DEFAULT_CONFIG, ...config };
+      this.reset({ targetFps: this.config.targetFps });
+    }
+
+    reset({ targetFps } = {}) {
+      this.targetFps = Math.max(1, Number(targetFps || this.config.targetFps));
+      this.sourceFps = this.targetFps;
+      this.serverFps = this.targetFps;
+      this.deliveryFps = this.targetFps;
+      this.hasServerSample = false;
+      this.hasDeliverySample = false;
+      this.latestChunkDurationMs = 1000 / this.targetFps;
+      this.latestChunkFrames = 1;
+      this.playbackRate = 1;
+      this.renderFps = this.targetFps;
+      this.queue = [];
+      this.lastDrawAt = 0;
+      this.lastRateUpdateAt = 0;
+      this.renderedFrames = 0;
+      this.droppedFrames = 0;
+      this.buffering = true;
+      this.pendingEventId = 0;
+      this.pendingEventSentAt = 0;
+      this.pendingEventCutoverMode = "motion";
+      this.lastDropReason = "";
+      this.lastDropAt = 0;
+      this.lastDropCount = 0;
+      this.rebufferLeadBoostMs = 0;
+      this.deliveryLeadBoostMs = 0;
+      this.chunkReceives = new Map();
+      this.serverStatChunks = new Set();
+      this.lastFinalReceiveAt = 0;
+      this.receiveStalled = false;
+    }
+
+    setTargetFps(targetFps) {
+      const nextTargetFps = Math.max(1, Number(targetFps || this.config.targetFps));
+      this.targetFps = nextTargetFps;
+      if (!this.hasServerSample && !this.hasDeliverySample) {
+        this.serverFps = nextTargetFps;
+        this.deliveryFps = nextTargetFps;
+        this.sourceFps = nextTargetFps;
+        this.renderFps = nextTargetFps;
+      } else {
+        this.serverFps = clamp(this.serverFps, this.config.minSourceFps, nextTargetFps);
+        this.deliveryFps = clamp(this.deliveryFps, this.config.minSourceFps, nextTargetFps);
+        this.sourceFps = clamp(this.sourceFps, this.config.minSourceFps, nextTargetFps);
+        this.renderFps = this.sourceFps * this.playbackRate;
+      }
+      this.latestChunkDurationMs = Math.max(this.latestChunkDurationMs, 1000 / this.targetFps);
+    }
+
+    clear() {
+      const frames = this.queue.splice(0);
+      this.lastDrawAt = 0;
+      this.buffering = true;
+      return frames;
+    }
+
+    noteInputEvent(eventId, now, { cutoverMode = "motion" } = {}) {
+      this.pendingEventId = Number(eventId || 0);
+      this.pendingEventSentAt = Number(now || 0);
+      this.pendingEventCutoverMode = cutoverMode;
+    }
+
+    observeServerStats(stats, now) {
+      const chunkIndex = Number(stats.chunk_index || 0);
+      const numFrames = Number(stats.num_frames || 0);
+      const chunkTotalMs = Number(stats.chunk_total_ms || 0);
+      if (numFrames > 0 && chunkTotalMs > 0) {
+        this.serverStatChunks.add(chunkIndex);
+        if (this.serverStatChunks.size > 128) {
+          this.serverStatChunks.delete(this.serverStatChunks.values().next().value);
+        }
+        const expectedMs = numFrames / Math.max(1, this.targetFps) * 1000;
+        const isStartupWarmup =
+          chunkIndex === 0 &&
+          chunkTotalMs > Math.max(
+            this.config.startupWarmupMinMs,
+            expectedMs * this.config.startupWarmupExpectedMultiplier,
+          );
+        if (isStartupWarmup) return this.snapshot();
+        this.#observeFpsSample("server", {
+          fps: numFrames / (chunkTotalMs / 1000),
+          frameCount: numFrames,
+          durationMs: chunkTotalMs,
+          now,
+        });
+      }
+      return this.snapshot();
+    }
+
+    enqueueDecodedFrames(header, frames, now) {
+      const chunkIndex = Number(header.chunk_index || 0);
+      const eventId = Number(header.event_id || 0);
+      const receivedAt = Number(header.__received_at || now);
+      const preparedFrames = frames.map((frame) => ({
+        ...frame,
+        chunk: Number(frame.chunk ?? chunkIndex),
+        chunkIndex,
+        eventId,
+      }));
+      const droppedFrames = [];
+      let cutover = null;
+
+      if (this.pendingEventId && eventId >= this.pendingEventId) {
+        const oldEventFrameCount = this.#oldEventFrameCount(eventId);
+        const graceFrames = this.#eventGraceFrames();
+        const dropCount = Math.max(0, oldEventFrameCount - graceFrames);
+        if (dropCount > 0) {
+          droppedFrames.push(...this.queue.splice(graceFrames, dropCount));
+          this.#recordDrop(dropCount, "event cutover", now);
+        }
+        cutover = {
+          eventId,
+          latencyMs: this.pendingEventSentAt ? now - this.pendingEventSentAt : 0,
+        };
+        this.pendingEventId = 0;
+        this.pendingEventSentAt = 0;
+        this.pendingEventCutoverMode = "motion";
+      }
+
+      this.queue.push(...preparedFrames);
+      this.#observeChunkArrival(header, preparedFrames.length, receivedAt, now);
+      droppedFrames.push(...this.#trimBacklog(now));
+      return { droppedFrames, cutover, snapshot: this.snapshot() };
+    }
+
+    render(now, { hasPendingInput = true } = {}) {
+      this.#decayRebufferBoost(now);
+      this.#updateReceiveStallGuard(now);
+      const droppedFrames = this.#trimBacklog(now);
+      if (!this.queue.length) {
+        if (this.renderedFrames && hasPendingInput && !this.buffering) {
+          this.buffering = true;
+          this.rebufferLeadBoostMs = Math.max(
+            this.rebufferLeadBoostMs,
+            this.config.rebufferLeadBoostMs,
+          );
+        }
+        return { action: "hold", droppedFrames, snapshot: this.snapshot() };
+      }
+
+      const bufferMs = this.bufferDurationMs;
+      if (
+        hasPendingInput &&
+        this.receiveStalled &&
+        this.renderedFrames &&
+        bufferMs < this.targetLeadMs
+      ) {
+        this.buffering = true;
+        this.lastDrawAt = 0;
+        return { action: "hold", droppedFrames, snapshot: this.snapshot() };
+      }
+
+      if (
+        hasPendingInput &&
+        this.buffering &&
+        bufferMs < (this.renderedFrames ? this.#resumeLeadMs() : this.#startLeadMs())
+      ) {
+        this.buffering = true;
+        this.lastDrawAt = 0;
+        return { action: "hold", droppedFrames, snapshot: this.snapshot() };
+      }
+
+      if (this.buffering) {
+        this.buffering = false;
+        this.lastDrawAt = 0;
+      }
+
+      this.#updatePlaybackRate(now);
+      const targetMs = 1000 / Math.max(1, this.renderFps);
+      const elapsedMs = this.lastDrawAt ? now - this.lastDrawAt : targetMs;
+      if (elapsedMs < targetMs) {
+        return { action: "wait", droppedFrames, snapshot: this.snapshot() };
+      }
+
+      const frame = this.queue.shift();
+      this.renderedFrames += 1;
+      this.lastDrawAt = !this.lastDrawAt || elapsedMs > targetMs * 4
+        ? now
+        : now - (elapsedMs % targetMs);
+      return { action: "draw", frame, droppedFrames, snapshot: this.snapshot() };
+    }
+
+    get queuedFrames() {
+      return this.queue.length;
+    }
+
+    get bufferDurationMs() {
+      return this.queue.length / Math.max(1, this.sourceFps) * 1000;
+    }
+
+    get targetLeadMs() {
+      const base = clamp(
+        this.latestChunkDurationMs * this.config.targetLeadChunkRatio,
+        this.config.minTargetLeadMs,
+        this.config.maxTargetLeadMs,
+      );
+      return clamp(
+        base + this.rebufferLeadBoostMs + this.deliveryLeadBoostMs,
+        this.config.minTargetLeadMs,
+        this.config.maxTargetLeadMs +
+          this.config.rebufferLeadBoostMs +
+          this.config.maxDeliveryLeadBoostMs,
+      );
+    }
+
+    get maxLeadMs() {
+      return this.targetLeadMs + this.latestChunkDurationMs * this.config.maxLeadExtraChunkRatio;
+    }
+
+    snapshot() {
+      return {
+        queueFrames: this.queue.length,
+        bufferMs: this.bufferDurationMs,
+        targetLeadMs: this.targetLeadMs,
+        maxLeadMs: this.maxLeadMs,
+        sourceFps: this.sourceFps,
+        serverFps: this.serverFps,
+        deliveryFps: this.deliveryFps,
+        targetFps: this.targetFps,
+        renderFps: this.renderFps,
+        playbackRate: this.playbackRate,
+        droppedFrames: this.droppedFrames,
+        lastDropAt: this.lastDropAt,
+        lastDropCount: this.lastDropCount,
+        buffering: this.buffering,
+        lastDropReason: this.lastDropReason,
+      };
+    }
+
+    #observeFpsSample(kind, { fps, frameCount, durationMs, now }) {
+      if (!finitePositive(fps)) return;
+      const cappedFps = clamp(fps, this.config.minSourceFps, this.targetFps);
+      const isDelivery = kind === "delivery";
+      const currentFps = isDelivery ? this.deliveryFps : this.serverFps;
+      const hasSample = isDelivery ? this.hasDeliverySample : this.hasServerSample;
+      let nextFps;
+      if (!hasSample) {
+        nextFps = cappedFps;
+      } else {
+        const alpha = cappedFps > currentFps
+          ? (isDelivery ? this.config.deliveryFpsAlphaUp : this.config.serverFpsAlphaUp)
+          : (isDelivery ? this.config.deliveryFpsAlphaDown : this.config.serverFpsAlphaDown);
+        nextFps = currentFps * (1 - alpha) + cappedFps * alpha;
+      }
+      if (isDelivery) {
+        this.deliveryFps = nextFps;
+        this.hasDeliverySample = true;
+        this.#observeDeliveryJitter(frameCount, durationMs);
+      } else {
+        this.serverFps = nextFps;
+        this.hasServerSample = true;
+      }
+      const effectiveFps = this.hasServerSample
+        ? this.serverFps
+        : (this.hasDeliverySample ? this.deliveryFps : this.targetFps);
+      this.sourceFps = clamp(effectiveFps, this.config.minSourceFps, this.targetFps);
+      if (!isDelivery || !this.hasServerSample) {
+        this.latestChunkFrames = Math.max(1, Number(frameCount || this.latestChunkFrames));
+        this.latestChunkDurationMs = clamp(
+          Number(durationMs || (this.latestChunkFrames / Math.max(1, this.sourceFps) * 1000)),
+          1000 / Math.max(1, this.targetFps),
+          2500,
+        );
+      }
+      this.#updatePlaybackRate(now);
+    }
+
+    #observeDeliveryJitter(frameCount, durationMs) {
+      if (!this.hasServerSample || !finitePositive(durationMs)) return;
+      const expectedMs = Number(frameCount || 0) / Math.max(1, this.serverFps) * 1000;
+      if (expectedMs <= 0) return;
+      if (durationMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
+      const boostMs = clamp(
+        durationMs - expectedMs,
+        0,
+        this.config.maxDeliveryLeadBoostMs,
+      );
+      this.deliveryLeadBoostMs = Math.max(this.deliveryLeadBoostMs, boostMs);
+    }
+
+    #updateReceiveStallGuard(now) {
+      this.receiveStalled = false;
+      if (!this.lastFinalReceiveAt || !this.hasServerSample) return;
+      const elapsedMs = now - this.lastFinalReceiveAt;
+      const expectedMs = Math.max(
+        this.latestChunkDurationMs,
+        this.latestChunkFrames / Math.max(1, this.serverFps) * 1000,
+      );
+      if (elapsedMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
+      this.receiveStalled = true;
+      this.deliveryLeadBoostMs = Math.max(
+        this.deliveryLeadBoostMs,
+        clamp(elapsedMs - expectedMs, 0, this.config.maxDeliveryLeadBoostMs),
+      );
+    }
+
+    #observeChunkArrival(header, frameCount, receivedAt, now) {
+      const chunkIndex = Number(header.chunk_index || 0);
+      const state = this.chunkReceives.get(chunkIndex) || {
+        firstAt: receivedAt,
+        frames: 0,
+      };
+      state.frames += Number(frameCount || 0);
+      state.lastAt = receivedAt;
+      this.chunkReceives.set(chunkIndex, state);
+
+      const frameBatchIndex = Number(header.frame_batch_index || 0);
+      const numFrameBatches = Number(header.num_frame_batches || 1);
+      const isFinalFrameBatch =
+        Boolean(header.is_final_frame_batch) ||
+        frameBatchIndex + 1 >= numFrameBatches;
+      if (!isFinalFrameBatch) return;
+      const durationMs = this.lastFinalReceiveAt
+        ? receivedAt - this.lastFinalReceiveAt
+        : 0;
+      this.lastFinalReceiveAt = receivedAt;
+      if (state.frames > 0 && durationMs > 0) {
+        this.#observeFpsSample("delivery", {
+          fps: state.frames / (durationMs / 1000),
+          frameCount: state.frames,
+          durationMs,
+          now,
+        });
+      }
+      this.chunkReceives.delete(chunkIndex);
+    }
+
+    #updatePlaybackRate(now) {
+      const bufferMs = this.bufferDurationMs;
+      const targetLeadMs = Math.max(1, this.targetLeadMs);
+      const error = (bufferMs - targetLeadMs) / targetLeadMs;
+      const emergency =
+        bufferMs > this.maxLeadMs ||
+        bufferMs < targetLeadMs * this.config.lowWaterRatio ||
+        (this.receiveStalled && bufferMs < targetLeadMs);
+      const minRate = emergency
+        ? (
+            this.receiveStalled
+              ? this.config.receiveStallPlaybackRateMin
+              : this.config.emergencyPlaybackRateMin
+          )
+        : this.config.playbackRateMin;
+      const maxRate = this.receiveStalled && bufferMs < targetLeadMs
+        ? 1
+        : emergency
+        ? this.config.emergencyPlaybackRateMax
+        : this.config.playbackRateMax;
+      const desiredRate = clamp(
+        1 + error * this.config.playbackRateGain,
+        minRate,
+        maxRate,
+      );
+
+      if (!this.lastRateUpdateAt) {
+        this.playbackRate = desiredRate;
+      } else {
+        const dtSeconds = Math.max(0.001, (now - this.lastRateUpdateAt) / 1000);
+        const slewPerSecond = this.receiveStalled
+          ? this.config.receiveStallPlaybackRateSlewPerSecond
+          : this.config.playbackRateSlewPerSecond;
+        const maxDelta = slewPerSecond * dtSeconds;
+        this.playbackRate = clamp(
+          desiredRate,
+          this.playbackRate - maxDelta,
+          this.playbackRate + maxDelta,
+        );
+      }
+      this.lastRateUpdateAt = now;
+      this.renderFps = clamp(
+        this.sourceFps * this.playbackRate,
+        this.config.minSourceFps,
+        this.targetFps * this.config.emergencyPlaybackRateMax,
+      );
+    }
+
+    #trimBacklog(now) {
+      const droppedFrames = [];
+      while (this.queue.length && this.bufferDurationMs > this.maxLeadMs) {
+        const firstChunk = this.queue[0].chunkIndex;
+        let dropCount = 0;
+        while (
+          dropCount < this.queue.length &&
+          this.queue[dropCount].chunkIndex === firstChunk
+        ) {
+          dropCount += 1;
+        }
+        if (!dropCount || dropCount >= this.queue.length) break;
+        droppedFrames.push(...this.queue.splice(0, dropCount));
+        this.#recordDrop(dropCount, "backlog", now);
+      }
+      return droppedFrames;
+    }
+
+    #oldEventFrameCount(nextEventId) {
+      let count = 0;
+      while (count < this.queue.length && this.queue[count].eventId < nextEventId) {
+        count += 1;
+      }
+      return count;
+    }
+
+    #eventGraceFrames() {
+      const byTime = Math.max(
+        1,
+        Math.round(this.sourceFps * this.#eventCutoverMaxMs() / 1000),
+      );
+      const byChunkRatio = this.pendingEventCutoverMode === "settle" ? 1.5 : 0.85;
+      const byChunk = Math.max(1, Math.round(this.latestChunkFrames * byChunkRatio));
+      return Math.min(this.#eventCutoverMaxFrames(), byTime, byChunk);
+    }
+
+    #eventCutoverMaxMs() {
+      return this.pendingEventCutoverMode === "settle"
+        ? this.config.settleEventCutoverMaxMs
+        : this.config.eventCutoverMaxMs;
+    }
+
+    #eventCutoverMaxFrames() {
+      return this.pendingEventCutoverMode === "settle"
+        ? this.config.settleEventCutoverMaxFrames
+        : this.config.eventCutoverMaxFrames;
+    }
+
+    #startLeadMs() {
+      return Math.max(
+        this.config.minStartLeadMs,
+        this.latestChunkDurationMs * this.config.startLeadChunkRatio,
+        this.targetLeadMs,
+      );
+    }
+
+    #resumeLeadMs() {
+      return clamp(
+        this.latestChunkDurationMs * this.config.resumeLeadChunkRatio,
+        this.config.minResumeLeadMs,
+        this.config.maxResumeLeadMs,
+      );
+    }
+
+    #decayRebufferBoost(now) {
+      if ((!this.rebufferLeadBoostMs && !this.deliveryLeadBoostMs) || !this.lastRateUpdateAt) return;
+      const dtSeconds = Math.max(0, (now - this.lastRateUpdateAt) / 1000);
+      this.rebufferLeadBoostMs = Math.max(
+        0,
+        this.rebufferLeadBoostMs - dtSeconds * this.config.rebufferLeadBoostDecayMsPerSecond,
+      );
+      this.deliveryLeadBoostMs = Math.max(
+        0,
+        this.deliveryLeadBoostMs - dtSeconds * this.config.deliveryLeadBoostDecayMsPerSecond,
+      );
+    }
+
+    #recordDrop(count, reason, now) {
+      this.droppedFrames += count;
+      this.lastDropAt = Number(now || 0);
+      this.lastDropCount = count;
+      this.lastDropReason = reason;
+    }
+  }
+
+  global.RealtimePlaybackController = RealtimePlaybackController;
+  if (typeof module !== "undefined" && module.exports) {
+    module.exports = { RealtimePlaybackController };
+  }
+})(typeof globalThis !== "undefined" ? globalThis : window);

python/sglang/multimodal_gen/apps/realtime_webui/playback_controller_test.js

@@ -0,0 +1,116 @@
+const assert = require("node:assert/strict");
+const { RealtimePlaybackController } = require("./playback_controller.js");
+
+function frames(count, chunk) {
+  return Array.from({ length: count }, (_, index) => ({
+    image: { close() {} },
+    chunk,
+    index,
+  }));
+}
+
+function enqueueChunk(controller, {
+  chunk,
+  eventId = 0,
+  frameCount = 12,
+  durationMs = 480,
+  now,
+}) {
+  controller.observeServerStats({
+    chunk_index: chunk,
+    num_frames: frameCount,
+    chunk_total_ms: durationMs,
+  }, now);
+  return controller.enqueueDecodedFrames({
+    chunk_index: chunk,
+    event_id: eventId,
+    num_frames: frameCount,
+    __received_at: now,
+    is_final_frame_batch: true,
+  }, frames(frameCount, chunk), now);
+}
+
+function renderFor(controller, startMs, durationMs) {
+  let rendered = 0;
+  for (let now = startMs; now <= startMs + durationMs; now += 16.67) {
+    const decision = controller.render(now, { hasPendingInput: true });
+    if (decision.action === "draw") rendered += 1;
+  }
+  return rendered;
+}
+
+function stableSourceDoesNotDrop() {
+  const controller = new RealtimePlaybackController({ targetFps: 25 });
+  let now = 0;
+  for (let chunk = 0; chunk < 8; chunk += 1) {
+    now += 480;
+    enqueueChunk(controller, { chunk, now });
+    renderFor(controller, now, 480);
+  }
+  const snapshot = controller.snapshot();
+  assert.equal(snapshot.droppedFrames, 0);
+  assert.ok(snapshot.sourceFps > 24 && snapshot.sourceFps <= 25);
+}
+
+function slowServerCapsRenderFps() {
+  const controller = new RealtimePlaybackController({ targetFps: 25 });
+  let now = 0;
+  for (let chunk = 0; chunk < 8; chunk += 1) {
+    now += 1360;
+    enqueueChunk(controller, { chunk, durationMs: 1360, now });
+    renderFor(controller, now, 1360);
+  }
+  const snapshot = controller.snapshot();
+  assert.ok(snapshot.sourceFps > 8 && snapshot.sourceFps < 10);
+  assert.ok(snapshot.renderFps <= 10);
+}
+
+function backlogDropsContiguousOldFrames() {
+  const controller = new RealtimePlaybackController({ targetFps: 25 });
+  let now = 100;
+  for (let chunk = 0; chunk < 13; chunk += 1) {
+    enqueueChunk(controller, { chunk, now, durationMs: 480 });
+    now += 20;
+  }
+  const snapshot = controller.snapshot();
+  assert.ok(snapshot.droppedFrames > 0);
+  assert.equal(snapshot.lastDropReason, "backlog");
+}
+
+function eventCutoverKeepsShortGrace() {
+  const controller = new RealtimePlaybackController({ targetFps: 25 });
+  enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
+  controller.noteInputEvent(5, 1050);
+  const result = enqueueChunk(controller, {
+    chunk: 2,
+    eventId: 5,
+    frameCount: 12,
+    durationMs: 480,
+    now: 1150,
+  });
+  assert.ok(result.cutover);
+  assert.ok(result.droppedFrames.length >= 14);
+  assert.equal(controller.queue[0].chunk, 1);
+  assert.equal(controller.queue[0].index, 0);
+}
+
+function settleEventCutoverKeepsWiderGrace() {
+  const controller = new RealtimePlaybackController({ targetFps: 25 });
+  enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
+  controller.noteInputEvent(5, 1050, { cutoverMode: "settle" });
+  const result = enqueueChunk(controller, {
+    chunk: 2,
+    eventId: 5,
+    frameCount: 12,
+    durationMs: 480,
+    now: 1150,
+  });
+  assert.ok(result.cutover);
+  assert.ok(result.droppedFrames.length <= 12);
+}
+
+stableSourceDoesNotDrop();
+slowServerCapsRenderFps();
+backlogDropsContiguousOldFrames();
+eventCutoverKeepsShortGrace();
+settleEventCutoverKeepsWiderGrace();

python/sglang/multimodal_gen/apps/realtime_webui/styles.css

@@ -0,0 +1,714 @@
+:root {
+  --paper: #eef1ec;
+  --panel: #fbfaf5;
+  --ink: #171a16;
+  --muted: #687164;
+  --line: #cbd2c4;
+  --accent: #b9543c;
+  --green: #4d765f;
+  --blue: #3f607c;
+  --pressed: #8c9288;
+  --pressed-border: #aeb4aa;
+  --pressed-ring: rgba(238, 241, 236, 0.2);
+  --shadow: 0 18px 60px rgba(23, 26, 22, 0.12);
+}
+
+* { box-sizing: border-box; }
+
+body {
+  margin: 0;
+  overflow-x: hidden;
+  min-height: 100vh;
+  background:
+    linear-gradient(90deg, rgba(23, 26, 22, 0.035) 1px, transparent 1px),
+    linear-gradient(180deg, rgba(23, 26, 22, 0.035) 1px, transparent 1px),
+    var(--paper);
+  background-size: 28px 28px;
+  color: var(--ink);
+  font-family: ui-sans-serif, "Avenir Next", "Helvetica Neue", sans-serif;
+}
+
+button, input, textarea, select { font: inherit; }
+button:disabled { cursor: wait; opacity: 0.64; transform: none; }
+
+.shell {
+  display: grid;
+  grid-template-columns: minmax(260px, 320px) minmax(0, 1fr);
+  gap: 18px;
+  width: 100%;
+  max-width: 100vw;
+  min-height: 100vh;
+  padding: 18px;
+}
+
+.panel {
+  background: color-mix(in oklch, var(--panel), white 20%);
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  box-shadow: var(--shadow);
+  padding: 18px;
+}
+
+.brand {
+  display: flex;
+  align-items: baseline;
+  gap: 10px;
+  margin-bottom: 22px;
+}
+
+.brand span {
+  color: var(--panel);
+  background: var(--ink);
+  border-radius: 4px;
+  padding: 3px 7px;
+  font-size: 12px;
+  letter-spacing: 0;
+}
+
+.brand strong { font-size: 18px; font-weight: 650; }
+
+label {
+  display: grid;
+  gap: 7px;
+  margin: 12px 0;
+  color: var(--muted);
+  font-size: 12px;
+}
+
+.label-row {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 8px;
+}
+
+.help-tooltip {
+  position: relative;
+  display: inline-grid;
+  place-items: center;
+  width: 18px;
+  height: 18px;
+  border: 1px solid var(--line);
+  border-radius: 50%;
+  color: var(--muted);
+  background: #fffdf7;
+  cursor: help;
+  font-size: 11px;
+  line-height: 1;
+}
+
+.help-tooltip::after {
+  content: attr(aria-label);
+  position: absolute;
+  right: 0;
+  bottom: calc(100% + 8px);
+  z-index: 20;
+  width: 280px;
+  max-width: min(280px, calc(100vw - 48px));
+  padding: 9px 10px;
+  border-radius: 6px;
+  background: var(--ink);
+  box-shadow: 0 12px 36px rgba(23, 26, 22, 0.24);
+  color: var(--panel);
+  font-size: 11px;
+  font-weight: 400;
+  line-height: 1.4;
+  opacity: 0;
+  pointer-events: none;
+  transform: translateY(4px);
+  transition: opacity 120ms ease, transform 120ms ease;
+}
+
+.help-tooltip:hover::after,
+.help-tooltip:focus-visible::after {
+  opacity: 1;
+  transform: translateY(0);
+}
+
+input, textarea, select {
+  width: 100%;
+  border: 1px solid var(--line);
+  border-radius: 6px;
+  background: #fffdf7;
+  color: var(--ink);
+  padding: 10px 11px;
+  outline: none;
+}
+
+textarea { resize: vertical; line-height: 1.45; }
+input:focus, textarea:focus, select:focus { border-color: var(--accent); box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.12); }
+
+.split { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; }
+.output-options {
+  align-items: end;
+}
+.output-options .toggle-row {
+  min-height: 40px;
+  margin: 12px 0;
+}
+.actions { display: grid; grid-template-columns: 1fr 0.7fr; gap: 10px; margin-top: 16px; }
+.toggle-row {
+  display: flex;
+  align-items: center;
+  gap: 9px;
+  margin-top: 14px;
+}
+.toggle-row input {
+  width: 16px;
+  height: 16px;
+  accent-color: var(--ink);
+}
+
+button {
+  border: 1px solid var(--line);
+  border-radius: 6px;
+  color: var(--ink);
+  background: #fffdf7;
+  min-height: 38px;
+  padding: 0 12px;
+  cursor: pointer;
+  transition:
+    background-color 120ms ease,
+    border-color 120ms ease,
+    box-shadow 120ms ease,
+    color 120ms ease,
+    transform 120ms ease;
+}
+
+button:hover:not(:disabled) {
+  border-color: var(--ink);
+  background: color-mix(in oklch, #fffdf7, var(--green) 10%);
+  box-shadow: 0 8px 18px rgba(23, 26, 22, 0.08);
+  transform: translateY(-1px);
+}
+button:active:not(:disabled),
+button.is-pressed:not(:disabled) {
+  border-color: var(--pressed-border);
+  background: var(--pressed);
+  color: #fffdf7;
+  box-shadow:
+    inset 0 0 0 1px rgba(255, 253, 247, 0.18),
+    inset 0 2px 7px rgba(23, 26, 22, 0.16);
+  transform: translateY(0);
+}
+button.is-key-active:not(:disabled) {
+  border-color: var(--pressed-border);
+  background: var(--pressed);
+  color: #fffdf7;
+  box-shadow:
+    inset 0 0 0 1px rgba(255, 253, 247, 0.22),
+    0 0 0 3px var(--pressed-ring),
+    0 10px 22px rgba(23, 26, 22, 0.18);
+  transform: none;
+}
+button:focus-visible {
+  outline: none;
+  box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.18);
+}
+.primary { background: var(--ink); color: var(--panel); border-color: var(--ink); }
+.primary:hover:not(:disabled) {
+  background: color-mix(in oklch, var(--ink), var(--green) 18%);
+  color: var(--panel);
+}
+.primary:active:not(:disabled),
+.primary.is-pressed:not(:disabled) {
+  background: var(--pressed);
+  border-color: var(--pressed-border);
+  color: var(--panel);
+}
+.wide { width: 100%; margin-top: 10px; }
+
+.workspace {
+  display: grid;
+  gap: 18px;
+  min-width: 0;
+  max-width: 100%;
+}
+
+.stage {
+  position: relative;
+  display: grid;
+  grid-template-rows: auto auto auto auto auto;
+  align-self: start;
+  justify-self: center;
+  min-width: 0;
+  width: 100%;
+  max-width: min(1500px, 100%);
+  overflow: hidden;
+  border: 1px solid #11140f;
+  border-radius: 8px;
+  background: #11140f;
+  box-shadow: var(--shadow);
+}
+
+.preview-frame {
+  position: relative;
+  display: grid;
+  place-items: center;
+  justify-self: center;
+  width: min(calc(1040px * var(--preview-scale, 1.2)), 100%);
+  overflow: hidden;
+  background: #11140f;
+  contain: paint;
+  isolation: isolate;
+}
+
+.preview-frame::before {
+  content: "";
+  position: absolute;
+  inset: 0;
+  z-index: 0;
+  pointer-events: none;
+  background: linear-gradient(
+    180deg,
+    rgba(238, 241, 236, 0.045),
+    transparent 34%,
+    rgba(0, 0, 0, 0.18)
+  );
+}
+
+.preview-frame::after {
+  content: none;
+}
+
+.stage[data-preview-state="waiting"] .preview-frame::after {
+  animation: none;
+}
+
+.topbar, .timeline {
+  display: flex;
+  align-items: center;
+  gap: 10px;
+  min-width: 0;
+  height: 44px;
+  padding: 0 14px;
+  color: #e8eadf;
+  background: rgba(17, 20, 15, 0.88);
+  font-size: 12px;
+  font-variant-numeric: tabular-nums;
+  white-space: nowrap;
+}
+
+.topbar > * {
+  flex: 0 0 auto;
+  align-self: center;
+}
+
+.topbar-spacer { flex: 1; }
+.record-button {
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  gap: 6px;
+  flex: 0 0 118px;
+  width: 118px;
+  min-height: 28px;
+  height: 28px;
+  padding: 0 9px;
+  border-color: rgba(232, 234, 223, 0.22);
+  background: rgba(238, 241, 236, 0.08);
+  color: #e8eadf;
+  font-size: 11px;
+  font-variant-numeric: tabular-nums;
+}
+.record-button:hover:not(:disabled) {
+  border-color: rgba(232, 234, 223, 0.44);
+  background: rgba(238, 241, 236, 0.14);
+  box-shadow: none;
+  transform: none;
+}
+.record-button:active:not(:disabled),
+.record-button.is-pressed:not(:disabled),
+.record-button:focus-visible {
+  transform: none;
+}
+.record-button.is-recording {
+  border-color: color-mix(in oklch, var(--accent), white 18%);
+  background: var(--accent);
+  color: #fffdf7;
+}
+.record-button.is-saving {
+  cursor: wait;
+  opacity: 0.76;
+}
+#recordLabel {
+  flex: 0 0 36px;
+  text-align: left;
+}
+.record-button-icon {
+  flex: 0 0 9px;
+  width: 9px;
+  height: 9px;
+  min-width: 9px;
+  border-radius: 50%;
+  background: var(--accent);
+  box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.16);
+}
+.record-button.is-recording .record-button-icon {
+  border-radius: 2px;
+  background: #fffdf7;
+  box-shadow: none;
+}
+.record-button-duration {
+  display: inline-block;
+  flex: 0 0 34px;
+  min-width: 34px;
+  text-align: right;
+  color: rgba(232, 234, 223, 0.7);
+}
+.record-button.is-recording .record-button-duration {
+  color: rgba(255, 253, 247, 0.86);
+}
+.preview-scale-control {
+  display: inline-flex;
+  align-items: center;
+  gap: 8px;
+  flex: 0 0 170px;
+  min-width: 170px;
+  margin: 0;
+  color: rgba(232, 234, 223, 0.72);
+  font-size: 11px;
+  line-height: 1;
+}
+.preview-scale-control input {
+  width: 92px;
+  min-width: 72px;
+  padding: 0;
+  border: 0;
+  background: transparent;
+  accent-color: #eef1ec;
+}
+.preview-scale-control b {
+  min-width: 36px;
+  color: #fffdf7;
+  font-weight: 650;
+}
+#statusText {
+  display: inline-block;
+  min-width: 92px;
+  line-height: 1;
+}
+#chunkText {
+  display: inline-block;
+  min-width: 70px;
+  line-height: 1;
+}
+.stage-stat {
+  display: inline-flex;
+  align-items: center;
+  gap: 5px;
+  flex: 0 1 auto;
+  min-width: 0;
+  color: rgba(232, 234, 223, 0.72);
+  line-height: 1;
+}
+.stage-stat b {
+  display: inline-block;
+  color: #fffdf7;
+  font-weight: 650;
+  font-variant-numeric: tabular-nums;
+}
+
+#outputSizeText { min-width: 206px; }
+#renderFps { min-width: 2ch; text-align: right; }
+#theoreticalFpsText { min-width: 116px; }
+#stageLatencyText { min-width: 120px; }
+
+@media (max-width: 1180px) {
+  .topbar {
+    flex-wrap: wrap;
+    height: auto;
+    min-height: 44px;
+    padding: 8px 14px;
+    row-gap: 7px;
+  }
+  .topbar-spacer { display: none; }
+  .preview-scale-control { flex-basis: 170px; min-width: 170px; }
+  #outputSizeText { min-width: 156px; }
+  #theoreticalFpsText { min-width: 100px; }
+  #stageLatencyText { min-width: 108px; }
+}
+
+.timeline { justify-content: flex-end; border-top: 1px solid rgba(232, 234, 223, 0.12); }
+.dot { width: 8px; height: 8px; border-radius: 50%; background: var(--muted); }
+.dot.live { background: #8ecf9d; box-shadow: 0 0 0 4px rgba(142, 207, 157, 0.14); }
+.dot.error { background: var(--accent); }
+
+#viewport {
+  position: relative;
+  z-index: 1;
+  display: block;
+  width: 100%;
+  height: auto;
+  max-height: min(calc(56vh * var(--preview-scale, 1.2)), 82vh);
+  min-height: 0;
+  object-fit: contain;
+  image-rendering: auto;
+  transform: translateZ(0);
+}
+
+.preview-overlay {
+  position: absolute;
+  inset: 0;
+  z-index: 3;
+  display: none;
+  place-items: center;
+  pointer-events: none;
+  background: transparent;
+}
+
+.stage[data-preview-state="waiting"] .preview-overlay {
+  display: grid;
+}
+
+.preview-loader {
+  width: 18px;
+  height: 18px;
+  border-radius: 50%;
+  border: 2px solid rgba(238, 241, 236, 0.22);
+  border-top-color: rgba(238, 241, 236, 0.82);
+  animation: previewProgressSpin 0.8s linear infinite;
+}
+
+.stage-controls {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(180px, 1fr));
+  gap: 12px;
+  padding: 12px 14px 13px;
+  border-top: 1px solid rgba(232, 234, 223, 0.12);
+  background: #151912;
+}
+
+.control-cluster {
+  display: grid;
+  grid-template-columns: 46px 1fr;
+  gap: 10px;
+  align-items: center;
+}
+
+.control-title {
+  color: rgba(232, 234, 223, 0.62);
+  font-size: 11px;
+  text-transform: uppercase;
+  letter-spacing: 0.08em;
+}
+
+.stage-controls .camera-pad {
+  margin: 0;
+}
+
+.stage-controls .camera-pad button {
+  position: relative;
+  border-color: rgba(232, 234, 223, 0.18);
+  background: #eef1ec;
+  color: #11140f;
+}
+
+.stage-controls .camera-pad button:active:not(:disabled),
+.stage-controls .camera-pad button.is-pressed:not(:disabled),
+.stage-controls .camera-pad button.is-key-active:not(:disabled) {
+  border-color: var(--pressed-border);
+  background: var(--pressed);
+  color: #fffdf7;
+  box-shadow:
+    inset 0 0 0 1px rgba(255, 253, 247, 0.22),
+    0 0 0 3px var(--pressed-ring),
+    0 10px 22px rgba(23, 26, 22, 0.18);
+}
+
+.stage-controls .camera-pad button::after {
+  content: attr(data-key);
+  position: absolute;
+  right: 7px;
+  top: 5px;
+  color: color-mix(in oklch, var(--muted), var(--ink) 18%);
+  font-size: 10px;
+  font-weight: 650;
+}
+
+.stage-controls .camera-pad button:active::after,
+.stage-controls .camera-pad button.is-pressed::after,
+.stage-controls .camera-pad button.is-key-active::after {
+  color: rgba(255, 253, 247, 0.78);
+}
+
+.section-title {
+  margin: 16px 0 10px;
+  color: var(--muted);
+  font-size: 12px;
+  text-transform: uppercase;
+  letter-spacing: 0.08em;
+}
+
+.reference-upload {
+  margin-top: 0;
+}
+
+.reference-upload input {
+  border: 1px dashed var(--line);
+}
+
+#referencePreview {
+  display: block;
+  width: 100%;
+  aspect-ratio: 16 / 9;
+  min-height: 0;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: #e5e7df;
+}
+
+#referenceName {
+  min-height: 18px;
+  color: var(--muted);
+  font-size: 11px;
+}
+
+.spec-grid {
+  display: grid;
+  grid-template-columns: repeat(4, minmax(120px, 1fr));
+  gap: 8px;
+  margin-bottom: 12px;
+}
+
+.spec-grid span {
+  display: grid;
+  gap: 2px;
+  min-height: 46px;
+  align-content: center;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: #fffdf7;
+  padding: 9px;
+  color: var(--muted);
+  font-size: 11px;
+}
+
+.spec-grid b {
+  color: var(--ink);
+  font-size: 14px;
+}
+
+.presets {
+  position: static;
+  max-height: none;
+  overflow: visible;
+  scrollbar-gutter: stable;
+}
+
+.preset-list {
+  display: grid;
+  grid-template-columns: repeat(auto-fill, minmax(220px, 1fr));
+  gap: 7px;
+  min-height: 0;
+  max-height: 230px;
+  margin-bottom: 12px;
+  overflow: auto;
+  padding-right: 3px;
+}
+.preset {
+  display: grid;
+  grid-template-columns: 72px minmax(0, 1fr);
+  gap: 4px 9px;
+  align-items: center;
+  padding: 8px;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: #fffdf7;
+  text-align: left;
+}
+.preset-thumb {
+  display: block;
+  grid-row: span 2;
+  width: 72px;
+  height: 46px;
+  min-height: 0;
+  object-fit: cover;
+  border-radius: 5px;
+  border: 1px solid color-mix(in oklch, var(--line), var(--ink) 8%);
+}
+.preset b { min-width: 0; font-size: 13px; }
+.preset span { min-width: 0; color: var(--muted); font-size: 11px; line-height: 1.25; }
+.preset[data-tone="green"] { border-left: 4px solid var(--green); }
+.preset[data-tone="blue"] { border-left: 4px solid var(--blue); }
+.preset[data-tone="accent"] { border-left: 4px solid var(--accent); }
+.preset:hover:not(:disabled) {
+  background: color-mix(in oklch, #fffdf7, var(--blue) 9%);
+}
+
+.camera-pad {
+  display: grid;
+  grid-template-columns: repeat(3, 1fr);
+  gap: 6px;
+  margin-bottom: 8px;
+}
+
+.camera-pad span { min-height: 36px; }
+.camera-pad button { min-height: 36px; font-size: 12px; }
+.telemetry { display: grid; gap: 7px; margin-top: 10px; }
+.telemetry span {
+  display: flex;
+  justify-content: space-between;
+  border-bottom: 1px solid var(--line);
+  padding-bottom: 8px;
+  color: var(--muted);
+  font-size: 12px;
+}
+.telemetry b { color: var(--ink); font-weight: 650; }
+
+.stage-telemetry {
+  grid-template-columns: repeat(3, minmax(0, 1fr));
+  gap: 0;
+  margin-top: 0;
+  border-top: 1px solid rgba(232, 234, 223, 0.12);
+  background: #11140f;
+}
+
+.stage-telemetry span {
+  min-height: 36px;
+  align-items: center;
+  gap: 8px;
+  border-right: 1px solid rgba(232, 234, 223, 0.1);
+  border-bottom: 1px solid rgba(232, 234, 223, 0.1);
+  padding: 0 14px;
+  color: rgba(232, 234, 223, 0.62);
+  font-size: 11px;
+}
+
+.stage-telemetry b {
+  color: #fffdf7;
+  font-size: 12px;
+}
+
+.history-list {
+  display: grid;
+  gap: 7px;
+  max-height: 92px;
+  overflow: auto;
+}
+
+.history-list span {
+  display: block;
+  border-left: 3px solid var(--blue);
+  background: #fffdf7;
+  padding: 8px 9px;
+  color: var(--muted);
+  font-size: 12px;
+}
+
+@media (max-width: 980px) {
+  .shell { grid-template-columns: 1fr; }
+  .presets { position: static; max-height: none; overflow: visible; }
+  .spec-grid { grid-template-columns: repeat(2, 1fr); }
+  .preset-list { min-height: 260px; max-height: 420px; }
+  #viewport { max-height: 420px; }
+  .stage-controls { grid-template-columns: 1fr; }
+  .stage-telemetry { grid-template-columns: repeat(2, minmax(0, 1fr)); }
+  .topbar { flex-wrap: wrap; height: auto; min-height: 44px; padding: 10px 14px; }
+  .topbar-spacer { display: none; }
+  .preview-scale-control { min-width: 160px; }
+}
+
+@keyframes previewProgressSpin {
+  to { transform: rotate(360deg); }
+}

python/sglang/multimodal_gen/apps/webui/README.md

@@ -0,0 +1,58 @@
+# SGLang Diffusion WebUI User Guide
+
+SGLang Diffusion WebUI provides an intuitive Gradio-based interface for image and video generation, supporting parameter
+tuning and real-time previews.
+
+## Prerequisites
+
+The WebUI runs on Gradio. To get started, install Gradio first:
+
+```bash
+pip install gradio==6.1.0
+```
+
+## Launch WebUI Service
+
+SGLang Diffusion now includes an integrated WebUI. Simply add the `--webui` parameter when starting the service.
+
+### Launch Text-to-Image Service
+
+```bash
+sglang serve --model-path black-forest-labs/FLUX.1-dev --num-gpus 1 --webui --webui-port 2333
+```
+
+### Launch Text-to-Video Service
+
+```bash
+sglang serve --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers --num-gpus 1 --webui --webui-port 2333
+```
+
+### Launch Image-to-Image Service
+```bash
+sglang serve --model-path Qwen/Qwen-Image-Edit-2511 --num-gpus 1 --webui --webui-port 2333
+```
+
+### Launch Image-to-Video Service
+```bash
+sglang serve --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers --num-gpus 1 --webui --webui-port 2333
+```
+
+## Port Forwarding
+
+Once the WebUI service is running, you need to use **SSH port forwarding** to securely access the remote service from
+your local machine.
+
+In most cases: Your IDE (like VS Code, Cursor, etc.) can handle this automatically. Check your IDE's remote development
+or port forwarding features. Otherwise, execute this command manually.
+
+```bash
+ssh -L ${WEBUI_PORT}:localhost:${WEBUI_PORT} user_name@machine_name
+```
+
+Learn more about port forwarding: [Port Forwarding](https://en.wikipedia.org/wiki/Port_forwarding).
+
+## Interface Instructions
+
+You can view your model path and task name directly in the UI. We'd appreciate any feedback you'd like to share.
+
+Once launched, access the interface at `http://localhost:${WEBUI_PORT}` in your browser.

python/sglang/multimodal_gen/apps/webui/__init__.py

@@ -0,0 +1,3 @@
+from .main import run_sgl_diffusion_webui
+
+__all__ = ["run_sgl_diffusion_webui"]

python/sglang/multimodal_gen/apps/webui/main.py

@@ -0,0 +1,265 @@
+import argparse
+import os
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    post_process_sample,
+    prepare_request,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.environ import envs
+
+logger = init_logger(__name__)
+
+
+def add_webui_args(parser: argparse.ArgumentParser):
+    """Add the arguments for the generate command."""
+    parser = ServerArgs.add_cli_args(parser)
+    parser = SamplingParams.add_cli_args(parser)
+    return parser
+
+
+def run_sgl_diffusion_webui(server_args: ServerArgs):
+    # import gradio in function to avoid CI crash
+
+    import gradio as gr
+
+    def resolve_model_repo_id(model_path: str) -> str:
+        from pathlib import Path
+
+        from huggingface_hub.utils import HFValidationError, validate_repo_id
+
+        try:
+            validate_repo_id(model_path)
+            return model_path
+        except HFValidationError:
+            pass
+
+        p = Path(model_path).expanduser()
+        parts = p.parts
+
+        if len(parts) < 2:
+            raise ValueError(f"Invalid model_path: {model_path}")
+
+        candidate = f"{parts[-2]}/{parts[-1]}"
+        validate_repo_id(candidate)  # let it raise if invalid
+        return candidate
+
+    # Prefer the hub pipeline tag for Hub models; fall back to the loaded pipeline's
+    # own task_type for local checkpoints (e.g. a single .safetensors path), which
+    # have no hub repo to query.
+    task_name = None
+    try:
+        repo_id = resolve_model_repo_id(server_args.model_path)
+        if envs.SGLANG_USE_MODELSCOPE.get():
+            from modelscope.hub.api import HubApi
+
+            api = HubApi()
+            model_info_obj = api.model_info(repo_id)
+            task_name = model_info_obj.tasks[0]["Name"].replace("-synthesis", "")
+        else:
+            from huggingface_hub import model_info
+
+            task_name = model_info(repo_id).pipeline_tag
+    except Exception as e:
+        logger.info(
+            "Could not resolve task from the model hub (%s); using the loaded "
+            "pipeline's task_type.",
+            e,
+        )
+
+    # init client
+    sync_scheduler_client.initialize(server_args)
+
+    if task_name in ("text-to-video", "image-to-video", "video-to-video"):
+        task_type = "video"
+    elif task_name in ("text-to-image", "image-to-image"):
+        task_type = "image"
+    else:
+        task_type = (
+            "image" if server_args.pipeline_config.task_type.is_image_gen() else "video"
+        )
+        task_name = task_name or server_args.pipeline_config.task_type.name
+    video_visible_only = task_type == "video"
+    image_visible_only = task_type == "image"
+
+    # server_args will be reused in gradio_generate function
+    def gradio_generate(
+        prompt,
+        negative_prompt,
+        reference_image_paths_str,
+        seed,
+        num_frames,
+        frames_per_second,
+        width,
+        height,
+        num_inference_steps,
+        guidance_scale,
+        enable_teacache,
+    ):
+        """
+        NOTE: The input and output of function which is called by gradio button must be gradio components
+        So we use global variable sampling_params_kwargs to avoid pass this param, because gradio does not support this.
+        return [ np.ndarray, None ] | [None, np.ndarray]
+        """
+        if reference_image_paths_str:
+            if "，" in reference_image_paths_str:
+                logger.warning(
+                    f"Warning: please use English comma to separate the reference image paths, and the reference image paths is: {reference_image_paths_str}"
+                )
+                reference_image_paths_str = reference_image_paths_str.replace("，", ",")
+            image_path = [path.strip() for path in reference_image_paths_str.split(",")]
+        else:
+            image_path = None
+
+        sampling_params_kwargs = dict(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image_path=image_path,
+            seed=seed,
+            num_frames=num_frames,
+            fps=frames_per_second,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            enable_teacache=enable_teacache,
+            return_file_paths_only=False,
+        )
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            server_args.model_path,
+            server_args=server_args,
+            **sampling_params_kwargs,
+        )
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling_params,
+        )
+        result = sync_scheduler_client.forward([batch])
+        save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
+        if result.output is None:
+            sampling_params_str = "\n".join(
+                [f"{key}: {value}" for key, value in sampling_params_kwargs.items()]
+            )
+            no_output_msg = f"No output is generated by client, and their sampling params is: {sampling_params_str}"
+
+            if batch.data_type == DataType.VIDEO:
+                if os.path.exists(save_file_path):
+                    logger.warning(no_output_msg)
+                    return None, save_file_path
+                else:
+                    no_output_msg += f"\nAnd the expected output file was not found at: {save_file_path}"
+                    raise ValueError(no_output_msg)
+            else:
+                raise ValueError(no_output_msg)
+
+        frames = post_process_sample(
+            result.output[0],
+            batch.data_type,
+            batch.fps,
+            batch.save_output,
+            save_file_path,
+        )
+        if batch.data_type == DataType.VIDEO:
+            # gradio video need video path to show video
+            return None, save_file_path
+        else:
+            return frames[0], None
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# 🚀 SGLang Diffusion Application")
+        with gr.Row():
+            gr.Textbox(label="Model", value=server_args.model_path)
+            gr.Textbox(label="Task name", value=task_name)
+
+        with gr.Row():
+            with gr.Column(scale=4):
+                prompt = gr.Textbox(label="Prompt", value="A curious raccoon")
+                negative_prompt = gr.Textbox(
+                    label="Negative_prompt",
+                    value="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
+                )
+            with gr.Column(scale=1):
+                seed = gr.Number(label="seed", precision=0, value=1234)
+                run_btn = gr.Button("Generate", variant="primary", size="lg")
+
+        with gr.Row():
+            with gr.Column():
+                width = gr.Number(label="width", precision=0, value=720)
+                height = gr.Number(label="height", precision=0, value=480)
+                num_inference_steps = gr.Slider(
+                    minimum=0, maximum=50, value=20, step=1, label="num_inference_steps"
+                )
+                guidance_scale = gr.Slider(
+                    minimum=0.0, maximum=10, value=5, step=0.01, label="guidance_scale"
+                )
+                num_frames = gr.Slider(
+                    minimum=1,
+                    maximum=181,
+                    value=81,
+                    step=1,
+                    label="num_frames",
+                    visible=video_visible_only,
+                )
+                frames_per_second = gr.Slider(
+                    minimum=4,
+                    maximum=60,
+                    value=16,
+                    step=1,
+                    label="frames_per_second",
+                    visible=video_visible_only,
+                )
+                reference_image_paths_str = gr.Textbox(
+                    label="reference images",
+                    placeholder="Examples: 'image1.png, image2.png' or 'https://example.com/image1.png, https://example.com/image2.png'",
+                )
+                enable_teacache = gr.Checkbox(label="enable_teacache", value=False)
+
+            with gr.Column():
+                image_out = gr.Image(
+                    label="Generated Image", visible=image_visible_only, format="png"
+                )
+                video_out = gr.Video(
+                    label="Generated Video", visible=video_visible_only
+                )
+
+        run_btn.click(
+            fn=gradio_generate,
+            inputs=[
+                prompt,
+                negative_prompt,
+                reference_image_paths_str,
+                seed,
+                num_frames,
+                frames_per_second,
+                width,
+                height,
+                num_inference_steps,
+                guidance_scale,
+                enable_teacache,
+            ],
+            outputs=[image_out, video_out],
+        )
+
+        _, local_url, _ = demo.launch(
+            server_port=server_args.webui_port,
+            quiet=True,
+            prevent_thread_lock=True,
+            show_error=True,
+        )
+
+        # print banner
+        delimiter = "=" * 80
+        url = local_url or f"http://localhost:{server_args.webui_port}"
+        print(f"""
+{delimiter}
+\033[1mSGLang Diffusion WebUI available at:\033[0m \033[1;4;92m{url}\033[0m
+{delimiter}
+""")
+
+        demo.block_thread()

python/sglang/multimodal_gen/benchmarks/bench_offline_throughput.py

@@ -0,0 +1,540 @@
+"""
+Benchmark offline throughput for multimodal generation models (Image/Video Generation).
+
+This script benchmarks generation throughput without running a server, using low-level APIs.
+It provides detailed metrics on throughput, latency, and resource utilization.
+
+# Usage Examples
+
+## Text-to-Video with VBench dataset
+python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\
+    --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\
+    --dataset vbench \\
+    --num-prompts 20 \\
+    --batch-size 1 \\
+    --width 512 --height 512 --num-frames 16
+
+## Random dataset for stress testing
+python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\
+    --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\
+    --dataset random \\
+    --num-prompts 100 \\
+    --batch-size 1 \\
+    --num-inference-steps 20 \\
+    --output-file results.json
+"""
+
+import argparse
+import dataclasses
+import json
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from tqdm import tqdm
+
+from sglang.multimodal_gen.benchmarks.datasets import RandomDataset, VBenchDataset
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    configure_logger,
+    init_logger,
+)
+from sglang.multimodal_gen.test.test_utils import print_divider, print_value_formatted
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class BatchOutput:
+    """Container for batch generation results."""
+
+    latency: float = 0.0
+    latency_per_sample: float = 0.0
+    num_samples: int = 0
+    total_frames: int = 0
+    peak_memory_mb: float = 0.0
+    success: bool = False
+    error: str = ""
+
+
+@dataclass
+class BenchArgs:
+    """Benchmark configuration for multimodal generation."""
+
+    # Diffusion Model Configuration
+    num_inference_steps: int = 20
+    guidance_scale: float = 7.5
+    seed: int = 42
+    disable_safety_checker: bool = False
+
+    # Output Configuration
+    width: int = 32
+    height: int = 32
+    num_frames: int = 1
+    fps: int = 24
+
+    # Dataset & Benchmark
+    dataset: str = "random"
+    dataset_path: str = ""
+    task_name: str = "unknown"
+    num_prompts: int = 10
+    num_outputs_per_prompt: int = 1
+    batch_size: int = 1
+    random_request_config: str = None
+    random_request_seed: int = 42
+
+    # Benchmark Execution
+    skip_warmup: bool = False
+    output_file: str = ""
+    disable_tqdm: bool = False
+
+    # Profiling
+    profile: bool = False
+    num_profiled_timesteps: int = 5
+    profile_all_stages: bool = False
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser):
+        """Add benchmark-specific CLI arguments."""
+        # Diffusion Model Configuration
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=20,
+            help="Number of denoising steps",
+        )
+        parser.add_argument(
+            "--guidance-scale",
+            type=float,
+            default=7.5,
+            help="Classifier-free guidance scale",
+        )
+        parser.add_argument("--seed", type=int, default=42, help="Random seed")
+        parser.add_argument(
+            "--disable-safety-checker",
+            action="store_true",
+            help="Disable NSFW detection",
+        )
+
+        # Output Configuration
+        parser.add_argument("--width", type=int, default=32, help="Image/video width")
+        parser.add_argument("--height", type=int, default=32, help="Image/video height")
+        parser.add_argument(
+            "--num-frames", type=int, default=1, help="Number of frames for video"
+        )
+        parser.add_argument("--fps", type=int, default=24, help="FPS for video")
+
+        # Dataset & Benchmark
+        parser.add_argument(
+            "--dataset",
+            type=str,
+            default="random",
+            choices=["vbench", "random"],
+            help="Dataset to use",
+        )
+        parser.add_argument(
+            "--dataset-path",
+            type=str,
+            default="",
+            help="Path to dataset (prompts file or image directory)",
+        )
+        parser.add_argument(
+            "--task-name",
+            type=str,
+            default="unknown",
+            help="Task name for benchmark identification",
+        )
+        parser.add_argument(
+            "--num-prompts",
+            type=int,
+            default=10,
+            help="Total number of prompts to benchmark",
+        )
+        parser.add_argument(
+            "--num-outputs-per-prompt",
+            type=int,
+            default=1,
+            help="Number of generated outputs requested per prompt",
+        )
+        parser.add_argument(
+            "--batch-size",
+            type=int,
+            default=1,
+            help="Batch size per generation call (currently only bs=1 is supported)",
+        )
+
+        parser.add_argument(
+            "--random-request-config",
+            type=str,
+            default=None,
+            help=(
+                "JSON string defining random request profiles. "
+                "Each profile may contain: width, height, num_inference_steps, etc. "
+                "The 'weight' field controls sampling probability (relative weight)."
+            ),
+        )
+        parser.add_argument(
+            "--random-request-seed",
+            type=int,
+            default=42,
+            help="Random seed for sampling request profiles (default: 42).",
+        )
+
+        # Benchmark Execution
+        parser.add_argument(
+            "--skip-warmup", action="store_true", help="Skip warmup batch"
+        )
+        parser.add_argument(
+            "--output-file",
+            type=str,
+            default="",
+            help="Output JSON file for results (append mode)",
+        )
+        parser.add_argument(
+            "--disable-tqdm",
+            action="store_true",
+            help="Disable progress bar",
+        )
+        parser.add_argument(
+            "--profile",
+            action="store_true",
+            help=(
+                "Enable PyTorch profiler for diffusion generation. "
+                "Set SGLANG_DIFFUSION_TORCH_PROFILER_DIR to control trace output directory."
+            ),
+        )
+        parser.add_argument(
+            "--num-profiled-timesteps",
+            type=int,
+            default=5,
+            help=(
+                "Number of denoising timesteps to profile after warmup. "
+                "Use -1 to profile all denoising timesteps."
+            ),
+        )
+        parser.add_argument(
+            "--profile-all-stages",
+            action="store_true",
+            help="Profile all diffusion pipeline stages instead of only denoising steps.",
+        )
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace):
+        """Create BenchArgs from parsed CLI arguments."""
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        return cls(**{attr: getattr(args, attr) for attr in attrs})
+
+
+def initialize_engine(server_args: ServerArgs) -> DiffGenerator:
+    """Initialize diffusion pipeline engine."""
+    logger.info("Initializing engine...")
+    engine = DiffGenerator.from_server_args(server_args, local_mode=True)
+    logger.info("Engine initialized successfully")
+    return engine
+
+
+def generate_batch(
+    engine: DiffGenerator,
+    bench_args: BenchArgs,
+    prompts: List[str],
+    user_sampling_params: List[Dict[str, Any]],
+) -> BatchOutput:
+    """Generate batch of images/videos synchronously."""
+    assert len(user_sampling_params) == len(prompts), (
+        f"user_sampling_params length ({len(user_sampling_params)}) must match "
+        f"prompts length ({len(prompts)})"
+    )
+
+    output = BatchOutput()
+    start_time = time.perf_counter()
+
+    torch.get_device_module().reset_peak_memory_stats()
+
+    for prompt, params in zip(prompts, user_sampling_params):
+        try:
+            sampling_params_kwargs = dict(params)
+            sampling_params_kwargs["prompt"] = prompt
+            result = engine.generate(sampling_params_kwargs=sampling_params_kwargs)
+
+            if result is not None:
+                if isinstance(result, list):
+                    output.total_frames += len(result)
+                else:
+                    output.total_frames += 1
+            output.num_samples += 1
+        except Exception as e:
+            logger.error(f"Generation failed for prompt '{prompt[:50]}...': {e}")
+            output.error = str(e)
+
+    output.latency = time.perf_counter() - start_time
+    output.latency_per_sample = output.latency / len(prompts) if prompts else 0.0
+    output.success = output.num_samples > 0
+    output.peak_memory_mb = torch.get_device_module().max_memory_allocated() / (
+        1024 * 1024
+    )
+
+    logger.debug(
+        f"Batch generated: {output.num_samples}/{len(prompts)} samples in {output.latency:.2f}s"
+    )
+
+    return output
+
+
+def calculate_metrics(
+    outputs: List[BatchOutput],
+    total_duration: float,
+    resolution: Tuple[int, int, int],
+    num_requests: int,
+    all_sampling_params: Optional[List[Dict[str, Any]]] = None,
+) -> Dict[str, Any]:
+    """Calculate generation-specific throughput metrics."""
+    successful = [o for o in outputs if o.success]
+    num_success = sum(o.num_samples for o in successful)
+    total_frames = sum(o.total_frames for o in successful)
+    peak_memory = max((o.peak_memory_mb for o in outputs), default=0)
+
+    width, height, frames = resolution
+    if all_sampling_params:
+        total_pixels = sum(
+            p.get("width", width)
+            * p.get("height", height)
+            * p.get("num_frames", frames)
+            for p in all_sampling_params[:num_success]
+        )
+    else:
+        total_pixels = num_success * width * height * frames
+
+    metrics = {
+        "num_requests": num_requests,
+        "successful_requests": num_success,
+        "failed_requests": num_requests - num_success,
+        "total_duration_seconds": total_duration,
+        "total_frames_generated": total_frames,
+        "total_pixels_generated": total_pixels,
+        "images_per_second": num_success / total_duration if total_duration > 0 else 0,
+        "frames_per_second": total_frames / total_duration if total_duration > 0 else 0,
+        "megapixels_per_second": (
+            total_pixels / (total_duration * 1e6) if total_duration > 0 else 0
+        ),
+        "requests_per_second": (
+            num_success / total_duration if total_duration > 0 else 0
+        ),
+        "latency_per_request_seconds": (
+            total_duration / num_success if num_success > 0 else 0
+        ),
+        "peak_memory_mb": peak_memory,
+    }
+
+    return metrics
+
+
+def throughput_test(
+    server_args: ServerArgs,
+    bench_args: BenchArgs,
+) -> Dict[str, Any]:
+    """Main throughput benchmark function."""
+    configure_logger(server_args=server_args)
+    logger.info("Starting offline throughput benchmark...")
+
+    engine = initialize_engine(server_args)
+
+    if bench_args.random_request_config and bench_args.dataset != "random":
+        raise ValueError(
+            "--random-request-config can only be used with --dataset random"
+        )
+
+    if bench_args.num_outputs_per_prompt != 1:
+        raise ValueError(
+            "bench_offline_throughput currently supports only --num-outputs-per-prompt 1"
+        )
+
+    logger.info(f"Loading {bench_args.dataset} dataset...")
+    if bench_args.dataset == "vbench":
+        bench_args.task_name = str(engine.server_args.pipeline_config.task_type)
+        dataset = VBenchDataset(bench_args)
+    elif bench_args.dataset == "random":
+        dataset = RandomDataset(bench_args)
+    else:
+        raise ValueError(f"Unknown dataset: {bench_args.dataset}")
+
+    _sampling_params = {
+        "guidance_scale": bench_args.guidance_scale,
+        "num_inference_steps": bench_args.num_inference_steps,
+        "height": bench_args.height,
+        "width": bench_args.width,
+        "num_frames": bench_args.num_frames,
+        "num_outputs_per_prompt": bench_args.num_outputs_per_prompt,
+        "seed": bench_args.seed,
+        "profile": bench_args.profile,
+        "num_profiled_timesteps": bench_args.num_profiled_timesteps,
+        "profile_all_stages": bench_args.profile_all_stages,
+    }
+    if bench_args.disable_safety_checker:
+        _sampling_params["safety_checker"] = None
+
+    total_count = min(bench_args.num_prompts, len(dataset))
+    all_prompts = [dataset[i].prompt for i in range(total_count)]
+
+    if bench_args.random_request_config:
+        all_sampling_params = []
+        for i in range(total_count):
+            params = dict(_sampling_params)
+            params.update(dataset.get_sampling_params(i))
+            all_sampling_params.append(params)
+    else:
+        all_sampling_params = [_sampling_params] * total_count
+
+    if not bench_args.skip_warmup:
+        logger.info("Running warmup batch...")
+        warmup_count = min(bench_args.batch_size, total_count)
+        warmup_prompts = all_prompts[:warmup_count]
+        warmup_sampling_params = [
+            {**p, "profile": False} for p in all_sampling_params[:warmup_count]
+        ]
+        generate_batch(engine, bench_args, warmup_prompts, warmup_sampling_params)
+
+    logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...")
+    outputs: List[BatchOutput] = []
+
+    start_time = time.perf_counter()
+
+    num_batches = (total_count + bench_args.batch_size - 1) // bench_args.batch_size
+    pbar = tqdm(
+        total=num_batches,
+        disable=bench_args.disable_tqdm,
+        desc="Benchmark",
+    )
+
+    for batch_start in range(0, total_count, bench_args.batch_size):
+        batch_end = min(batch_start + bench_args.batch_size, total_count)
+        batch_prompts = all_prompts[batch_start:batch_end]
+        batch_sampling_params = all_sampling_params[batch_start:batch_end]
+
+        batch_output = generate_batch(
+            engine, bench_args, batch_prompts, batch_sampling_params
+        )
+        outputs.append(batch_output)
+
+        pbar.update(1)
+
+    pbar.close()
+    total_duration = time.perf_counter() - start_time
+
+    resolution = (bench_args.width, bench_args.height, bench_args.num_frames)
+    metrics = calculate_metrics(
+        outputs,
+        total_duration,
+        resolution=resolution,
+        num_requests=total_count,
+        all_sampling_params=all_sampling_params,
+    )
+
+    display_results(
+        metrics,
+        bench_args,
+        model_path=server_args.model_path,
+    )
+
+    if bench_args.output_file:
+        save_results(metrics, bench_args, server_args)
+
+    return metrics
+
+
+def display_results(
+    metrics: Dict[str, Any],
+    bench_args: BenchArgs,
+    model_path: str,
+):
+    """Display benchmark results in console."""
+    print(
+        "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=110, c="=")
+    )
+    print_value_formatted("Model:", model_path)
+    print_value_formatted("Dataset:", bench_args.dataset)
+    print_value_formatted(
+        "Resolution:",
+        f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}",
+    )
+    print_value_formatted("Num Inference Steps:", bench_args.num_inference_steps)
+    print_divider(75)
+    print_value_formatted("Total Requests:", metrics["num_requests"])
+    print_value_formatted("Successful Requests:", metrics["successful_requests"])
+    print_value_formatted("Failed Requests:", metrics["failed_requests"])
+    print_value_formatted(
+        "Total Duration (seconds):", metrics["total_duration_seconds"]
+    )
+    print_divider(75)
+    print_value_formatted("Frames Generated:", metrics["total_frames_generated"])
+    print_value_formatted(
+        "Megapixels Generated:", metrics["total_pixels_generated"] / 1e6
+    )
+    print_divider(75)
+    print_value_formatted(
+        "Frame Throughput (frames/sec):", metrics["frames_per_second"]
+    )
+    print_value_formatted("MP Throughput (MP/sec):", metrics["megapixels_per_second"])
+    print_value_formatted("Requests Per Second:", metrics["requests_per_second"])
+    print_value_formatted(
+        "Latency Per Request (sec):", metrics["latency_per_request_seconds"]
+    )
+    print_value_formatted("Peak Memory (MB):", metrics["peak_memory_mb"])
+    print_divider(110, "=")
+
+
+def save_results(
+    metrics: Dict[str, Any],
+    bench_args: BenchArgs,
+    server_args: ServerArgs,
+):
+    """Save benchmark results to JSON file."""
+    result = {
+        "metadata": {
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            "model_path": server_args.model_path,
+            "task_type": bench_args.task_name,
+            "backend": "engine",
+        },
+        "configuration": {
+            "num_inference_steps": bench_args.num_inference_steps,
+            "guidance_scale": bench_args.guidance_scale,
+            "seed": bench_args.seed,
+            "batch_size": bench_args.batch_size,
+            "num_prompts": bench_args.num_prompts,
+            "resolution": f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}",
+            "dataset": bench_args.dataset,
+        },
+        "results": metrics,
+    }
+
+    with open(bench_args.output_file, "a") as f:
+        f.write(json.dumps(result) + "\n")
+
+    logger.info(f"Results saved to {bench_args.output_file}")
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(
+        description="Offline throughput benchmark for multimodal generation models"
+    )
+
+    ServerArgs.add_cli_args(parser)
+    BenchArgs.add_cli_args(parser)
+
+    args, unknown_args = parser.parse_known_args()
+
+    server_args = ServerArgs.from_cli_args(args, unknown_args)
+    bench_args = BenchArgs.from_cli_args(args)
+
+    set_global_server_args(server_args)
+
+    result = throughput_test(server_args, bench_args)
+
+    return result
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/benchmarks/bench_serving.py

@@ -0,0 +1,842 @@
+"""
+Benchmark online serving for diffusion models (Image/Video Generation).
+
+
+Usage:
+    # launch a server and benchmark on it
+
+    # T2V or T2I or any other multimodal generation model
+    sglang serve --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers --num-gpus 1 --port 1231
+
+    # benchmark it and make sure the port is the same as the server's port
+    python3 -m sglang.multimodal_gen.benchmarks.bench_serving --dataset vbench --num-prompts 20 --port 1231
+
+    # benchmark with SLO metrics enabled
+    python3 -m sglang.multimodal_gen.benchmarks.bench_serving --dataset vbench --num-prompts 20 --port 1231 --slo --slo-scale 3.0 --warmup-requests 2
+"""
+
+import argparse
+import asyncio
+import json
+import os
+import time
+from dataclasses import replace
+from typing import Any, Dict, List, Optional
+
+import aiohttp
+import numpy as np
+import requests
+from tqdm.asyncio import tqdm
+
+from sglang.multimodal_gen.benchmarks.datasets import (
+    RandomDataset,
+    RequestFuncInput,
+    RequestFuncOutput,
+    VBenchDataset,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    configure_logger,
+    init_logger,
+)
+from sglang.multimodal_gen.test.test_utils import print_divider, print_value_formatted
+from sglang.srt.utils.network import NetworkAddress
+
+logger = init_logger(__name__)
+
+# Patch size used for computing area units (e.g. in latent diffusion models).
+PATCH_SIZE = 16
+PATCH_AREA = PATCH_SIZE * PATCH_SIZE
+
+
+def _get_response_output_count(resp_json: Dict[str, Any]) -> int:
+    if isinstance(resp_json.get("num_outputs"), int):
+        return resp_json["num_outputs"]
+    if isinstance(resp_json.get("data"), list):
+        return len(resp_json["data"])
+    if isinstance(resp_json.get("file_paths"), list):
+        return len(resp_json["file_paths"])
+    if isinstance(resp_json.get("urls"), list):
+        return len(resp_json["urls"])
+    if resp_json.get("file_path") or resp_json.get("url"):
+        return 1
+    return 0
+
+
+def _compute_scale_factor(req: RequestFuncInput, args) -> Optional[float]:
+    """Computes the composite scale factor (area × frames × steps) for a request."""
+    width = req.width or args.width
+    height = req.height or args.height
+    if None in (width, height):
+        return None
+    frames = req.num_frames or args.num_frames
+    steps = req.num_inference_steps or args.num_inference_steps
+
+    frame_scale = frames if isinstance(frames, int) and frames > 0 else 1
+    step_scale = steps if isinstance(steps, int) and steps > 0 else 1
+
+    area_units = max((float(width) * float(height)) / float(PATCH_AREA), 1.0)
+    return area_units * float(frame_scale) * float(step_scale)
+
+
+def _compute_expected_latency_ms_from_base(
+    req: RequestFuncInput, args, base_time_ms: Optional[float]
+) -> Optional[float]:
+    """Scales latency linearly by pixel area, frame count, and inference steps."""
+    if base_time_ms is None:
+        return None
+    scale = _compute_scale_factor(req, args)
+    if scale is None:
+        return None
+    return float(base_time_ms) * scale
+
+
+def _infer_slo_base_time_ms_from_warmups(
+    warmup_pairs: List[tuple], args
+) -> Optional[float]:
+    """Derives median base latency from successful warmup runs."""
+    candidates_ms: List[float] = []
+    for req, out in warmup_pairs:
+        if not out.success or out.latency <= 0:
+            logger.warning(
+                f"Skipping warmup result: success={out.success}, latency={out.latency:.3f}"
+            )
+            continue
+
+        scale = _compute_scale_factor(req, args)
+        if scale is None or scale <= 0:
+            continue
+
+        candidates_ms.append((out.latency * 1000.0) / scale)
+
+    return float(np.median(candidates_ms)) if candidates_ms else None
+
+
+def _populate_slo_ms_from_warmups(
+    requests_list: List[RequestFuncInput], warmup_pairs: List[tuple], args
+) -> List[RequestFuncInput]:
+    """Assigns estimated SLO targets to requests lacking them."""
+    if not any(req.slo_ms is None for req in requests_list):
+        return requests_list
+
+    base_time_ms = _infer_slo_base_time_ms_from_warmups(warmup_pairs, args)
+    if base_time_ms is None:
+        return requests_list
+
+    slo_scale = float(getattr(args, "slo_scale", 3.0))
+    if slo_scale <= 0:
+        raise ValueError(f"slo_scale must be positive, got {slo_scale}.")
+
+    updated: List[RequestFuncInput] = []
+    for req in requests_list:
+        if req.slo_ms is not None:
+            updated.append(req)
+            continue
+        expected_ms = _compute_expected_latency_ms_from_base(req, args, base_time_ms)
+        if expected_ms is not None:
+            # Create a new RequestFuncInput with updated slo_ms
+            updated.append(replace(req, slo_ms=expected_ms * slo_scale))
+        else:
+            updated.append(req)
+
+    return updated
+
+
+async def async_request_image_sglang(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: Optional[tqdm] = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Check if we need to use multipart (for image edits with input images)
+    if input.image_paths and len(input.image_paths) > 0:
+        # Use multipart/form-data for image edits
+        data = aiohttp.FormData()
+        data.add_field("model", input.model)
+        data.add_field("prompt", input.prompt)
+        data.add_field("response_format", "b64_json")
+        data.add_field("n", str(input.num_outputs_per_prompt))
+
+        if input.width and input.height:
+            data.add_field("size", f"{input.width}x{input.height}")
+
+        # Merge extra parameters
+        for key, value in input.extra_body.items():
+            data.add_field(key, str(value))
+
+        # Add image file(s)
+        for idx, img_path in enumerate(input.image_paths):
+            if os.path.exists(img_path):
+                data.add_field(
+                    "image",
+                    open(img_path, "rb"),
+                    filename=os.path.basename(img_path),
+                    content_type="application/octet-stream",
+                )
+            else:
+                output.error = f"Image file not found: {img_path}"
+                output.success = False
+                if pbar:
+                    pbar.update(1)
+                return output
+
+        try:
+            async with session.post(input.api_url, data=data) as response:
+                if response.status == 200:
+                    resp_json = await response.json()
+                    output.response_body = resp_json
+                    output.success = True
+                    output.output_count = _get_response_output_count(resp_json)
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
+                else:
+                    output.error = f"HTTP {response.status}: {await response.text()}"
+                    output.success = False
+        except Exception as e:
+            output.error = str(e)
+            output.success = False
+    else:
+        # Use JSON for text-to-image generation
+        payload = {
+            "model": input.model,
+            "prompt": input.prompt,
+            "n": input.num_outputs_per_prompt,
+            "response_format": "b64_json",
+        }
+
+        if input.width and input.height:
+            payload["size"] = f"{input.width}x{input.height}"
+        if input.num_inference_steps:
+            payload["num_inference_steps"] = input.num_inference_steps
+
+        # Merge extra parameters
+        payload.update(input.extra_body)
+
+        try:
+            async with session.post(input.api_url, json=payload) as response:
+                if response.status == 200:
+                    resp_json = await response.json()
+                    output.response_body = resp_json
+                    output.success = True
+                    output.output_count = _get_response_output_count(resp_json)
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
+                else:
+                    output.error = f"HTTP {response.status}: {await response.text()}"
+                    output.success = False
+        except Exception as e:
+            output.error = str(e)
+            output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    # Check SLO if defined
+    if input.slo_ms is not None and output.success:
+        output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_video_sglang(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: Optional[tqdm] = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # 1. Submit Job
+    job_id = None
+    # Check if we need to upload images (Multipart) or just send JSON
+    if input.image_paths and len(input.image_paths) > 0:
+        # Use multipart/form-data
+        data = aiohttp.FormData()
+        data.add_field("model", input.model)
+        data.add_field("prompt", input.prompt)
+        data.add_field("num_outputs_per_prompt", str(input.num_outputs_per_prompt))
+
+        if input.width and input.height:
+            data.add_field("size", f"{input.width}x{input.height}")
+
+        # Add extra body fields to form data if possible, or assume simple key-values
+        # Note: Nested dicts in extra_body might need JSON serialization if API expects it stringified
+        if input.extra_body:
+            data.add_field("extra_body", json.dumps(input.extra_body))
+
+        # Explicitly add fps/num_frames if they are not in extra_body (bench_serving logic overrides)
+        if input.num_frames:
+            data.add_field("num_frames", str(input.num_frames))
+        if input.fps:
+            data.add_field("fps", str(input.fps))
+
+        # Add image file
+        # Currently only support single image upload as 'input_reference' per API spec
+        img_path = input.image_paths[0]
+        if os.path.exists(img_path):
+            data.add_field(
+                "input_reference",
+                open(img_path, "rb"),
+                filename=os.path.basename(img_path),
+                content_type="application/octet-stream",
+            )
+        else:
+            output.error = f"Image file not found: {img_path}"
+            output.success = False
+            if pbar:
+                pbar.update(1)
+            return output
+
+        try:
+            async with session.post(input.api_url, data=data) as response:
+                if response.status == 200:
+                    resp_json = await response.json()
+                    job_id = resp_json.get("id")
+                else:
+                    output.error = (
+                        f"Submit failed HTTP {response.status}: {await response.text()}"
+                    )
+                    output.success = False
+                    if pbar:
+                        pbar.update(1)
+                    return output
+        except Exception as e:
+            output.error = f"Submit exception: {str(e)}"
+            output.success = False
+            if pbar:
+                pbar.update(1)
+            return output
+
+    else:
+        # Use JSON
+        payload: Dict[str, Any] = {
+            "model": input.model,
+            "prompt": input.prompt,
+            "num_outputs_per_prompt": input.num_outputs_per_prompt,
+        }
+        if input.width and input.height:
+            payload["size"] = f"{input.width}x{input.height}"
+        if input.num_frames:
+            payload["num_frames"] = input.num_frames
+        if input.num_inference_steps:
+            payload["num_inference_steps"] = input.num_inference_steps
+        if input.fps:
+            payload["fps"] = input.fps
+
+        payload.update(input.extra_body)
+
+        try:
+            async with session.post(input.api_url, json=payload) as response:
+                if response.status == 200:
+                    resp_json = await response.json()
+                    job_id = resp_json.get("id")
+                else:
+                    output.error = (
+                        f"Submit failed HTTP {response.status}: {await response.text()}"
+                    )
+                    output.success = False
+                    if pbar:
+                        pbar.update(1)
+                    return output
+        except Exception as e:
+            output.error = f"Submit exception: {str(e)}"
+            output.success = False
+            if pbar:
+                pbar.update(1)
+            return output
+
+    if not job_id:
+        output.error = "No job_id returned"
+        output.success = False
+        if pbar:
+            pbar.update(1)
+        return output
+
+    # 2. Poll for completion
+    # Assuming the API returns a 'status' field.
+    # We construct the check URL. Assuming api_url is like .../v1/videos
+    # The check url should be .../v1/videos/{id}
+    check_url = f"{input.api_url}/{job_id}"
+
+    while True:
+        try:
+            async with session.get(check_url) as response:
+                if response.status == 200:
+                    status_data = await response.json()
+                    status = status_data.get("status")
+                    if status == "completed":
+                        output.success = True
+                        output.response_body = status_data
+                        output.output_count = _get_response_output_count(status_data)
+                        if "peak_memory_mb" in status_data:
+                            output.peak_memory_mb = status_data["peak_memory_mb"]
+                        break
+                    elif status == "failed":
+                        output.success = False
+                        output.error = f"Job failed: {status_data.get('error')}"
+                        break
+                    else:
+                        # queued or processing
+                        await asyncio.sleep(1.0)
+                else:
+                    output.success = False
+                    output.error = (
+                        f"Poll failed HTTP {response.status}: {await response.text()}"
+                    )
+                    break
+        except Exception as e:
+            output.success = False
+            output.error = f"Poll exception: {str(e)}"
+            break
+
+    output.latency = time.perf_counter() - output.start_time
+
+    # Check SLO if defined
+    if input.slo_ms is not None and output.success:
+        output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+def calculate_metrics(
+    outputs: List[RequestFuncOutput],
+    total_duration: float,
+    requests_list: List[RequestFuncInput],
+    args,
+    slo_enabled: bool,
+):
+    success_outputs = [o for o in outputs if o.success]
+    error_outputs = [o for o in outputs if not o.success]
+
+    num_success = len(success_outputs)
+    latencies = [o.latency for o in success_outputs]
+    completed_outputs = sum(o.output_count for o in success_outputs)
+    peak_memories = [
+        o.peak_memory_mb
+        for o in success_outputs
+        if o.peak_memory_mb is not None and o.peak_memory_mb > 0
+    ]
+
+    metrics = {
+        "duration": total_duration,
+        "completed_requests": num_success,
+        "completed_outputs": completed_outputs,
+        "failed_requests": len(error_outputs),
+        "throughput_qps": num_success / total_duration if total_duration > 0 else 0,
+        "output_throughput_ops": (
+            completed_outputs / total_duration if total_duration > 0 else 0
+        ),
+        "latency_mean": np.mean(latencies) if latencies else 0,
+        "latency_median": np.median(latencies) if latencies else 0,
+        "latency_p50": np.percentile(latencies, 50) if latencies else 0,
+        "latency_p90": np.percentile(latencies, 90) if latencies else 0,
+        "latency_p95": np.percentile(latencies, 95) if latencies else 0,
+        "latency_p99": np.percentile(latencies, 99) if latencies else 0,
+        "num_outputs_per_prompt": args.num_outputs_per_prompt,
+        "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
+    }
+
+    if slo_enabled:
+        slo_defined_total = 0
+        slo_met_success = 0
+
+        for req, out in zip(requests_list, outputs):
+            if req.slo_ms is None:
+                continue
+            slo_defined_total += 1
+            if out.slo_achieved:
+                slo_met_success += 1
+
+        slo_attain_all = (
+            (slo_met_success / slo_defined_total) if slo_defined_total > 0 else 0.0
+        )
+
+        metrics.update(
+            {
+                "slo_attainment_rate": slo_attain_all,
+                "slo_met_success": slo_met_success,
+                "slo_scale": getattr(args, "slo_scale", 3.0),
+            }
+        )
+
+    return metrics
+
+
+def wait_for_service(base_url: str, timeout: int = 1200) -> None:
+    logger.info(f"Waiting for service at {base_url}...")
+    start_time = time.time()
+    while True:
+        try:
+            # Try /health endpoint first
+            resp = requests.get(f"{base_url}/health", timeout=1)
+            if resp.status_code == 200:
+                logger.info("Service is ready.")
+                break
+        except requests.exceptions.RequestException:
+            pass
+
+        if time.time() - start_time > timeout:
+            raise TimeoutError(
+                f"Service at {base_url} did not start within {timeout} seconds."
+            )
+
+        time.sleep(1)
+
+
+async def benchmark(args):
+    from huggingface_hub import model_info
+
+    # Construct base_url if not provided
+    if args.base_url is None:
+        args.base_url = NetworkAddress(args.host, args.port).to_url()
+
+    # Wait for service
+    wait_for_service(args.base_url)
+
+    # Fetch model info
+    try:
+        resp = requests.get(f"{args.base_url}/v1/model_info", timeout=5)
+        if resp.status_code == 200:
+            info = resp.json()
+            if "model_path" in info and info["model_path"]:
+                args.model = info["model_path"]
+                logger.info(f"Updated model name from server: {args.model}")
+    except Exception as e:
+        logger.info(f"Failed to fetch model info: {e}. Using default: {args.model}")
+
+    valid_tasks = (
+        "text-to-video",
+        "image-to-video",
+        "video-to-video",
+        "text-to-image",
+        "image-to-image",
+    )
+
+    # Resolve task_name with priority: args.task > local config > HF pipeline_tag
+    if args.task:
+        task_name = args.task
+        logger.info(f"Using task from --task: {task_name}")
+    elif os.path.exists(args.model):
+        config_path = os.path.join(args.model, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                config = json.load(f)
+            task_name = config.get("pipeline_tag", "text-to-image")
+            logger.info(f"Inferred task from local config.json: {task_name}")
+        else:
+            task_name = "text-to-image"
+            logger.info(f"No config.json found, defaulting task to: {task_name}")
+    else:
+        task_name = model_info(args.model).pipeline_tag
+        logger.info(f"Inferred task from HuggingFace pipeline_tag: {task_name}")
+
+    if task_name not in valid_tasks:
+        raise ValueError(
+            f"Task '{task_name}' is not a valid multimodal generation task. "
+            f"Use --task to specify one of: {', '.join(valid_tasks)}"
+        )
+
+    if task_name in ("text-to-video", "image-to-video", "video-to-video"):
+        api_url = f"{args.base_url}/v1/videos"
+        request_func = async_request_video_sglang
+    else:  # text-to-image or image-to-image
+        api_url = (
+            f"{args.base_url}/v1/images/edits"
+            if task_name == "image-to-image"
+            else f"{args.base_url}/v1/images/generations"
+        )
+        request_func = async_request_image_sglang
+
+    setattr(args, "task_name", task_name)
+
+    if args.random_request_config and args.dataset != "random":
+        raise ValueError(
+            "--random-request-config can only be used with --dataset random"
+        )
+
+    if args.dataset == "vbench":
+        dataset = VBenchDataset(args, api_url, args.model)
+    elif args.dataset == "random":
+        dataset = RandomDataset(args, api_url, args.model)
+    else:
+        raise ValueError(f"Unknown dataset: {args.dataset}")
+
+    logger.info("Loading requests...")
+    requests_list = dataset.get_requests()
+    logger.info(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
+
+    # Limit concurrency
+    if args.max_concurrency is not None:
+        semaphore = asyncio.Semaphore(args.max_concurrency)
+    else:
+        semaphore = None
+
+    async def limited_request_func(req, session, pbar):
+        if semaphore:
+            async with semaphore:
+                return await request_func(req, session, pbar)
+        else:
+            return await request_func(req, session, pbar)
+
+    async with aiohttp.ClientSession() as session:
+        # Run warmup requests
+        warmup_pairs: List[tuple] = []
+        if args.warmup_requests and requests_list:
+            # The server always overrides warmup requests to use
+            # num_inference_steps=1 (see Req.set_as_warmup), so we match
+            # that here to keep the benchmark's SLO estimation consistent.
+            warmup_steps = 1
+            logger.info(
+                f"Running {args.warmup_requests} warmup request(s) with "
+                f"num_inference_steps={warmup_steps}..."
+            )
+            for i in range(args.warmup_requests):
+                warm_req = requests_list[i % len(requests_list)]
+                warm_req = replace(
+                    warm_req,
+                    num_inference_steps=warmup_steps,
+                )
+                warm_out = await limited_request_func(warm_req, session, None)
+                warmup_pairs.append((warm_req, warm_out))
+                logger.info(
+                    f"Warmup {i+1}/{args.warmup_requests}: "
+                    f"latency={warm_out.latency:.2f}s, success={warm_out.success}"
+                )
+
+        # Populate SLO values from warmups if enabled
+        if args.slo:
+            requests_list = _populate_slo_ms_from_warmups(
+                requests_list=requests_list, warmup_pairs=warmup_pairs, args=args
+            )
+
+        # Run benchmark
+        pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
+        start_time = time.perf_counter()
+        tasks = []
+        for req in requests_list:
+            if args.request_rate != float("inf"):
+                # Poisson process: inter-arrival times follow exponential distribution
+                interval = np.random.exponential(1.0 / args.request_rate)
+                await asyncio.sleep(interval)
+
+            task = asyncio.create_task(limited_request_func(req, session, pbar))
+            tasks.append(task)
+
+        outputs = await asyncio.gather(*tasks)
+        total_duration = time.perf_counter() - start_time
+
+        pbar.close()
+
+    # Calculate metrics
+    metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
+
+    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
+
+    # Section 1: Configuration
+    print_value_formatted("Task:", task_name)
+    print_value_formatted("Model:", args.model)
+    print_value_formatted("Dataset:", args.dataset)
+
+    # Section 2: Execution & Traffic
+    print_divider(50)
+    print_value_formatted("Benchmark duration (s):", metrics["duration"])
+    print_value_formatted("Request rate:", str(args.request_rate))
+    print_value_formatted(
+        "Max request concurrency:",
+        str(args.max_concurrency) if args.max_concurrency else "not set",
+    )
+    print_value_formatted(
+        "Successful requests:",
+        f"{metrics['completed_requests']}/{len(requests_list)}",
+    )
+    print_value_formatted("Completed outputs:", metrics["completed_outputs"])
+    print_value_formatted("Outputs per prompt:", metrics["num_outputs_per_prompt"])
+
+    # Section 3: Performance Metrics
+    print_divider(50)
+
+    print_value_formatted("Request throughput (req/s):", metrics["throughput_qps"])
+    print_value_formatted(
+        "Output throughput (outputs/s):", metrics["output_throughput_ops"]
+    )
+
+    print_value_formatted("Latency Mean (s):", metrics["latency_mean"])
+    print_value_formatted("Latency Median (s):", metrics["latency_median"])
+    print_value_formatted("Latency P90 (s):", metrics["latency_p90"])
+    print_value_formatted("Latency P95 (s):", metrics["latency_p95"])
+    print_value_formatted("Latency P99 (s):", metrics["latency_p99"])
+
+    if metrics["peak_memory_mb_max"] > 0:
+        print_divider(50)
+        print_value_formatted("Peak Memory Max (MB):", metrics["peak_memory_mb_max"])
+        print_value_formatted("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"])
+        print_value_formatted(
+            "Peak Memory Median (MB):", metrics["peak_memory_mb_median"]
+        )
+
+    if args.slo and "slo_attainment_rate" in metrics:
+        print_divider(50)
+        print(
+            "{:<40} {:<15.2%}".format(
+                "SLO Attainment Rate:", metrics["slo_attainment_rate"]
+            )
+        )
+        print("{:<40} {:<15}".format("SLO Met (Success):", metrics["slo_met_success"]))
+        print("{:<40} {:<15.2f}".format("SLO Scale:", metrics["slo_scale"]))
+
+    print_divider(60)
+
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Benchmark serving for diffusion models."
+    )
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default=None,
+        help="DEPRECATED: --task is deprecated and will be ignored. The task will be inferred from --model.",
+    )
+    parser.add_argument(
+        "--base-url",
+        type=str,
+        default=None,
+        help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.",
+    )
+    parser.add_argument("--host", type=str, default="localhost", help="Server host.")
+    parser.add_argument("--port", type=int, default=30000, help="Server port.")
+    parser.add_argument("--model", type=str, default="default", help="Model name.")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="vbench",
+        choices=["vbench", "random"],
+        help="Dataset to use.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        choices=[
+            "text-to-video",
+            "image-to-video",
+            "text-to-image",
+            "image-to-image",
+            "video-to-video",
+        ],
+        default=None,
+        help="The task will be inferred from huggingface pipeline_tag. When huggingface pipeline_tag is not provided, --task will be used.",
+    )
+    parser.add_argument(
+        "--dataset-path",
+        type=str,
+        default=None,
+        help="Path to local dataset file (optional).",
+    )
+    parser.add_argument(
+        "--num-prompts", type=int, default=10, help="Number of prompts to benchmark."
+    )
+    parser.add_argument(
+        "--num-outputs-per-prompt",
+        type=int,
+        default=1,
+        help="Number of generated outputs requested per prompt.",
+    )
+    parser.add_argument(
+        "--max-concurrency",
+        type=int,
+        default=1,
+        help="Maximum number of concurrent requests, default to `1`. This can be used "
+        "to help simulate an environment where a higher level component "
+        "is enforcing a maximum number of concurrent requests. While the "
+        "--request-rate argument controls the rate at which requests are "
+        "initiated, this argument will control how many are actually allowed "
+        "to execute at a time. This means that when used in combination, the "
+        "actual request rate may be lower than specified with --request-rate, "
+        "if the server is not processing requests fast enough to keep up.",
+    )
+    parser.add_argument(
+        "--request-rate",
+        type=float,
+        default=float("inf"),
+        help="Number of requests per second. If this is inf, then all the requests are sent at time 0. "
+        "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.",
+    )
+    parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
+    parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
+    parser.add_argument(
+        "--random-request-config",
+        type=str,
+        default=None,
+        help=(
+            "JSON string defining random request profiles. "
+            "Each profile may contain: width, height, num_inference_steps, "
+            "num_outputs_per_prompt, etc. "
+            "The 'weight' field controls sampling probability (relative weight). "
+            "Example: "
+            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
+            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
+        ),
+    )
+    parser.add_argument(
+        "--random-request-seed",
+        type=int,
+        default=42,
+        help="Random seed for sampling request profiles (default: 42).",
+    )
+    parser.add_argument(
+        "--num-frames", type=int, default=None, help="Number of frames (for video)."
+    )
+    parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
+    parser.add_argument(
+        "--output-file", type=str, default=None, help="Output JSON file for metrics."
+    )
+    parser.add_argument(
+        "--disable-tqdm", action="store_true", help="Disable progress bar."
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Log level.",
+    )
+    parser.add_argument(
+        "--slo",
+        action="store_true",
+        help="Enable SLO calculation. Uses trace-provided slo_ms or infers from warmups.",
+    )
+    parser.add_argument(
+        "--slo-scale",
+        type=float,
+        default=3.0,
+        help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
+    )
+    parser.add_argument(
+        "--warmup-requests",
+        type=int,
+        default=1,
+        help="Number of warmup requests to run before measurement.",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=None,
+        help="Number of inference steps for diffusion models.",
+    )
+
+    args = parser.parse_args()
+
+    configure_logger(args)
+
+    asyncio.run(benchmark(args))

python/sglang/multimodal_gen/benchmarks/compare_perf.py

@@ -0,0 +1,301 @@
+import argparse
+import json
+import os
+import re
+from datetime import datetime
+from typing import Any, Dict, List, Tuple
+
+
+def calculate_diff(base: float, new: float) -> Tuple[float, float]:
+    """Returns (diff, diff_percent)."""
+    diff = new - base
+    if base == 0:
+        percent = 0.0
+    else:
+        percent = (diff / base) * 100
+    return diff, percent
+
+
+def calculate_upper_bound(baseline: float, rel_tol: float, min_abs_tol: float) -> float:
+    """Calculates the upper bound for performance regression check."""
+    rel_limit = baseline * (1 + rel_tol)
+    abs_limit = baseline + min_abs_tol
+    return max(rel_limit, abs_limit)
+
+
+def calculate_lower_bound(baseline: float, rel_tol: float, min_abs_tol: float) -> float:
+    """Calculates the lower bound for performance improvement check."""
+    rel_lower = baseline * (1 - rel_tol)
+    abs_lower = baseline - min_abs_tol
+    return min(rel_lower, abs_lower)
+
+
+def get_perf_status_emoji(
+    baseline: float,
+    new: float,
+    rel_tol: float = 0.1,
+    min_abs_tol: float = 120.0,
+) -> str:
+    """
+    Determines the status emoji based on performance difference.
+
+    Logic:
+      Upper bound (Slower): max(baseline * (1 + rel_tol), baseline + min_abs_tol)
+      Lower bound (Faster): min(baseline * (1 - rel_tol), baseline - min_abs_tol)
+    """
+    upper_bound = calculate_upper_bound(baseline, rel_tol, min_abs_tol)
+    lower_bound = calculate_lower_bound(baseline, rel_tol, min_abs_tol)
+
+    if new > upper_bound:
+        return "🔴"
+    elif new < lower_bound:
+        return "🟢"
+    else:
+        return "⚪️"
+
+
+def consolidate_steps(
+    steps_list: List[Dict[str, Any]],
+) -> Tuple[Dict[str, float], List[str], Dict[str, int]]:
+    """
+    Aggregates specific repeating steps (like denoising_step_*) into groups.
+    Returns:
+        - aggregated_durations: {name: duration_ms}
+        - ordered_names: list of names in execution order
+        - counts: {name: count_of_steps_aggregated}
+    """
+    durations = {}
+    counts = {}
+    ordered_names = []
+    seen_names = set()
+
+    # Regex for steps to group
+    # Group "denoising_step_0", "denoising_step_1" -> "Denoising Loop"
+    denoise_pattern = re.compile(r"^denoising_step_(\d+)$")
+    denoising_group_name = "Denoising Loop"
+
+    for step in steps_list:
+        name = step.get("name", "unknown")
+        dur = step.get("duration_ms", 0.0)
+
+        match = denoise_pattern.match(name)
+        if match:
+            key = denoising_group_name
+            if key not in durations:
+                durations[key] = 0.0
+                counts[key] = 0
+                if key not in seen_names:
+                    ordered_names.append(key)
+                    seen_names.add(key)
+            durations[key] += dur
+            counts[key] += 1
+        else:
+            # Standard stage (preserve order)
+            if name not in durations:
+                durations[name] = 0.0
+                counts[name] = 0
+                if name not in seen_names:
+                    ordered_names.append(name)
+                    seen_names.add(name)
+            durations[name] += dur
+            counts[name] += 1
+
+    return durations, ordered_names, counts
+
+
+def _load_benchmark_file(file_path: str) -> Dict[str, Any]:
+    """Loads a benchmark JSON file."""
+    with open(file_path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _get_status_emoji_from_diff_percent(diff_pct):
+    if diff_pct < -2.0:
+        return "✅"
+    elif diff_pct > 2.0:
+        return "❌"
+    else:
+        return "⚪️"
+
+
+def _print_single_comparison_report(
+    others_data, base_e2e, combined_order, base_durations, others_processed, base_counts
+):
+    new_data = others_data[0]
+    new_e2e = new_data.get("total_duration_ms", 0)
+    diff_ms, diff_pct = calculate_diff(base_e2e, new_e2e)
+    status = _get_status_emoji_from_diff_percent(diff_pct)
+
+    print("#### 1. High-level Summary")
+    print("| Metric | Baseline | New | Diff | Status |")
+    print("| :--- | :--- | :--- | :--- | :--- |")
+    print(
+        f"| **E2E Latency** | {base_e2e:.2f} ms | {new_e2e:.2f} ms | **{diff_ms:+.2f} ms ({diff_pct:+.1f}%)** | {status} |"
+    )
+    print(
+        f"| **Throughput** | {1000 / base_e2e if base_e2e else 0:.2f} req/s | {1000 / new_e2e if new_e2e else 0:.2f} req/s | - | - |"
+    )
+    print("\n")
+
+    print("#### 2. Stage Breakdown")
+    print("| Stage Name | Baseline (ms) | New (ms) | Diff (ms) | Diff (%) | Status |")
+    print("| :--- | :--- | :--- | :--- | :--- | :--- |")
+
+    new_durations, _, new_counts = others_processed[0]
+
+    for stage in combined_order:
+        b_val = base_durations.get(stage, 0.0)
+        n_val = new_durations.get(stage, 0.0)
+        b_count = base_counts.get(stage, 1)
+        n_count = new_counts.get(stage, 1)
+
+        s_diff, s_pct = calculate_diff(b_val, n_val)
+
+        count_str = ""
+        if stage == "Denoising Loop":
+            count_str = (
+                f" ({n_count} steps)"
+                if n_count == b_count
+                else f" ({b_count}->{n_count} steps)"
+            )
+
+        status_emoji = get_perf_status_emoji(b_val, n_val)
+        print(
+            f"| {stage}{count_str} | {b_val:.2f} | {n_val:.2f} | {s_diff:+.2f} | {s_pct:+.1f}% | {status_emoji} |"
+        )
+
+
+def _print_multi_comparison_report(
+    base_e2e,
+    others_data,
+    other_labels,
+    combined_order,
+    base_durations,
+    others_processed,
+):
+    print("#### 1. High-level Summary")
+    header = "| Metric | Baseline | " + " | ".join(other_labels) + " |"
+    sep = "| :--- | :--- | " + " | ".join([":---"] * len(other_labels)) + " |"
+    print(header)
+    print(sep)
+
+    # E2E Row
+    row_e2e = f"| **E2E Latency** | {base_e2e:.2f} ms |"
+    for i, d in enumerate(others_data):
+        val = d.get("total_duration_ms", 0)
+        diff_ms, diff_pct = calculate_diff(base_e2e, val)
+
+        status = _get_status_emoji_from_diff_percent(diff_pct)
+
+        row_e2e += f" {val:.2f} ms ({diff_pct:+.1f}%) {status} |"
+    print(row_e2e)
+    print("\n")
+
+    print("#### 2. Stage Breakdown")
+    # Header: Stage | Baseline | Label1 | Label2 ...
+    header = "| Stage Name | Baseline | " + " | ".join(other_labels) + " |"
+    sep = "| :--- | :--- | " + " | ".join([":---"] * len(other_labels)) + " |"
+    print(header)
+    print(sep)
+
+    for stage in combined_order:
+        b_val = base_durations.get(stage, 0.0)
+        row_str = f"| {stage} | {b_val:.2f} |"
+
+        for i, (n_durations, _, n_counts) in enumerate(others_processed):
+            n_val = n_durations.get(stage, 0.0)
+            _, s_pct = calculate_diff(b_val, n_val)
+            status_emoji = get_perf_status_emoji(b_val, n_val)
+
+            row_str += f" {n_val:.2f} ({s_pct:+.1f}%) {status_emoji} |"
+        print(row_str)
+
+
+def compare_benchmarks(file_paths: List[str], output_format: str = "markdown"):
+    """
+    Compares benchmark JSON files and prints a report.
+    First file is baseline, others will be compared against it.
+    """
+    if len(file_paths) < 2:
+        print("Error: Need at least 2 files to compare.")
+        return
+
+    try:
+        data_list = [_load_benchmark_file(f) for f in file_paths]
+    except Exception as e:
+        print(f"Error loading benchmark files: {e}")
+        return
+
+    base_data = data_list[0]
+    others_data = data_list[1:]
+
+    # Use filenames as labels if multiple comparisons, else just "New"
+    other_labels = [os.path.basename(p) for p in file_paths[1:]]
+
+    base_e2e = base_data.get("total_duration_ms", 0)
+
+    base_durations, base_order, base_counts = consolidate_steps(
+        base_data.get("steps", [])
+    )
+
+    others_processed = []
+    for d in others_data:
+        dur, order, counts = consolidate_steps(d.get("steps", []))
+        others_processed.append((dur, order, counts))
+
+    combined_order = []
+    # Collect all unique stages maintaining order from newest to baseline
+    for _, order, _ in reversed(others_processed):
+        for name in order:
+            if name not in combined_order:
+                combined_order.append(name)
+    for name in base_order:
+        if name not in combined_order:
+            combined_order.append(name)
+
+    if output_format == "markdown":
+        print("### Performance Comparison Report\n")
+
+        if len(others_data) == 1:
+            _print_single_comparison_report(
+                others_data,
+                base_e2e,
+                combined_order,
+                base_durations,
+                others_processed,
+                base_counts,
+            )
+        else:
+            _print_multi_comparison_report(
+                base_e2e,
+                others_data,
+                other_labels,
+                combined_order,
+                base_durations,
+                others_processed,
+            )
+
+        print("\n")
+        # Metadata
+        print("<details>")
+        print("<summary>Metadata</summary>\n")
+        print(f"- Baseline Commit: `{base_data.get('commit_hash', 'N/A')}`")
+        for i, d in enumerate(others_data):
+            label = "New" if len(others_data) == 1 else other_labels[i]
+            print(f"- {label} Commit: `{d.get('commit_hash', 'N/A')}`")
+        print(f"- Timestamp: {datetime.now().isoformat()}")
+        print("</details>")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Compare sglang-diffusion performance JSON files."
+    )
+    parser.add_argument(
+        "files",
+        nargs="+",
+        help="List of JSON files. First is baseline, others are compared against it.",
+    )
+    args = parser.parse_args()
+
+    compare_benchmarks(args.files)

python/sglang/multimodal_gen/benchmarks/datasets.py

@@ -0,0 +1,361 @@
+import glob
+import json
+import os
+import random
+import re
+import subprocess
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import requests
+from PIL import Image
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class RequestFuncInput:
+    prompt: str
+    api_url: str = ""
+    model: str = ""
+    num_outputs_per_prompt: int = 1
+    width: Optional[int] = None
+    height: Optional[int] = None
+    num_frames: Optional[int] = None
+    fps: Optional[int] = None
+    extra_body: Dict[str, Any] = field(default_factory=dict)
+    image_paths: Optional[List[str]] = None
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    slo_ms: Optional[float] = None
+    num_inference_steps: Optional[int] = None
+
+
+@dataclass
+class RequestFuncOutput:
+    success: bool = False
+    latency: float = 0.0
+    error: str = ""
+    start_time: float = 0.0
+    response_body: Dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
+    slo_achieved: Optional[bool] = None
+    output_count: int = 0
+
+
+def is_dir_not_empty(path: str) -> bool:
+    return os.path.isdir(path) and bool(os.listdir(path))
+
+
+class BaseDataset(ABC):
+    def __init__(self, args, api_url: str = "", model: str = ""):
+        self.args = args
+        self.api_url = api_url
+        self.model = model
+        self.items: List[Dict[str, Any]] = []
+
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @abstractmethod
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        pass
+
+    def get_requests(self) -> List[RequestFuncInput]:
+        return [self[i] for i in range(len(self))]
+
+
+class VBenchDataset(BaseDataset):
+    """
+    Dataset loader for VBench prompts.
+    Supports t2v, i2v.
+    """
+
+    T2V_PROMPT_URL = "https://raw.githubusercontent.com/Vchitect/VBench/master/prompts/prompts_per_dimension/subject_consistency.txt"
+    I2V_DOWNLOAD_SCRIPT_URL = "https://raw.githubusercontent.com/Vchitect/VBench/master/vbench2_beta_i2v/download_data.sh"
+
+    def __init__(self, args, api_url: str = "", model: str = ""):
+        super().__init__(args, api_url, model)
+        self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "sglang")
+        self.items = self._load_data()
+
+    @staticmethod
+    def _normalize_task_name(task_name: Any) -> Any:
+        """Normalize enum-style task values to legacy benchmark task-name strings."""
+        enum_to_task_name = {
+            "T2V": "text-to-video",
+            "I2V": "image-to-video",
+            "TI2V": "image-to-video",
+            "T2I": "text-to-image",
+            "I2I": "image-to-image",
+            "TI2I": "image-to-image",
+        }
+
+        # Handle Enum-like objects, e.g., ModelTaskType.T2I
+        enum_name = getattr(task_name, "name", None)
+        if isinstance(enum_name, str):
+            return enum_to_task_name.get(enum_name, task_name)
+
+        # Handle direct string inputs or enum string repr
+        if isinstance(task_name, str):
+            if task_name in enum_to_task_name:
+                return enum_to_task_name[task_name]
+            if "." in task_name:
+                suffix = task_name.split(".")[-1]
+                return enum_to_task_name.get(suffix, task_name)
+
+        return task_name
+
+    def _load_data(self) -> List[Dict[str, Any]]:
+        task_name = self._normalize_task_name(self.args.task_name)
+
+        if task_name in ("text-to-video", "text-to-image", "video-to-video"):
+            return self._load_t2v_prompts()
+        elif task_name in ("image-to-video", "image-to-image"):
+            return self._load_i2v_data()
+        else:
+            raise ValueError(
+                f"Illegal task name is found in VBenchDataset {self.args.task_name}"
+            )
+
+    def _download_file(self, url: str, dest_path: str) -> None:
+        """Download a file from URL to destination path."""
+        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+        resp = requests.get(url)
+        resp.raise_for_status()
+        with open(dest_path, "w") as f:
+            f.write(resp.text)
+
+    def _load_t2v_prompts(self) -> List[Dict[str, Any]]:
+        path = self.args.dataset_path
+
+        if not path:
+            path = os.path.join(self.cache_dir, "vbench_subject_consistency.txt")
+            if not os.path.exists(path):
+                logger.info(f"Downloading VBench T2V prompts to {path}...")
+                try:
+                    self._download_file(self.T2V_PROMPT_URL, path)
+                except Exception as e:
+                    logger.info(f"Failed to download VBench prompts: {e}")
+                    return [{"prompt": "A cat sitting on a bench"}] * 50
+
+        prompts = []
+        with open(path, "r") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    prompts.append({"prompt": line})
+
+        return self._resize_data(prompts)
+
+    def _auto_download_i2v_dataset(self) -> Optional[str]:
+        """Auto-download VBench I2V dataset and return the dataset directory."""
+        vbench_i2v_dir = os.path.join(self.cache_dir, "vbench_i2v", "vbench2_beta_i2v")
+        info_json_path = os.path.join(vbench_i2v_dir, "data", "i2v-bench-info.json")
+        crop_dir = os.path.join(vbench_i2v_dir, "data", "crop")
+        origin_dir = os.path.join(vbench_i2v_dir, "data", "origin")
+
+        if (
+            os.path.exists(info_json_path)
+            and is_dir_not_empty(crop_dir)
+            and is_dir_not_empty(origin_dir)
+        ):
+            return vbench_i2v_dir
+
+        logger.info(f"Downloading VBench I2V dataset to {vbench_i2v_dir}...")
+        try:
+            cache_root = os.path.join(self.cache_dir, "vbench_i2v")
+            script_path = os.path.join(cache_root, "download_data.sh")
+
+            self._download_file(self.I2V_DOWNLOAD_SCRIPT_URL, script_path)
+            os.chmod(script_path, 0o755)
+
+            logger.info("Executing download_data.sh (this may take a while)...")
+
+            result = subprocess.run(
+                ["bash", script_path],
+                cwd=cache_root,
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode != 0:
+                raise RuntimeError(f"Download script failed: {result.stderr}")
+            missing_packages = re.findall(r"(\S+): command not found", result.stderr)
+            if missing_packages:
+                missing_packages = list(set(missing_packages))
+                package_list = ", ".join(f"'{cmd}'" for cmd in missing_packages)
+                raise RuntimeError(
+                    f"Download script failed because the following commands are not installed: {package_list}.\n"
+                    "Please install them (e.g., on Ubuntu: `sudo apt install ...`) and try again."
+                )
+            logger.info(
+                f"Successfully downloaded VBench I2V dataset to {vbench_i2v_dir}"
+            )
+        except Exception as e:
+            logger.info(f"Failed to download VBench I2V dataset: {e}")
+            logger.info("Please manually download following instructions at:")
+            logger.info(
+                "https://github.com/Vchitect/VBench/tree/master/vbench2_beta_i2v#22-download"
+            )
+            return None
+
+        return vbench_i2v_dir if os.path.exists(info_json_path) else None
+
+    def _load_from_i2v_json(self, json_path: str) -> List[Dict[str, Any]]:
+        """Load I2V data from i2v-bench-info.json format."""
+        with open(json_path, "r") as f:
+            items = json.load(f)
+
+        base_dir = os.path.dirname(
+            os.path.dirname(json_path)
+        )  # Go up to vbench2_beta_i2v
+        origin_dir = os.path.join(base_dir, "data", "origin")
+
+        data = []
+        for item in items:
+            img_path = os.path.join(origin_dir, item.get("file_name", ""))
+            if os.path.exists(img_path):
+                data.append({"prompt": item.get("caption", ""), "image_path": img_path})
+            else:
+                logger.warning(f"Image not found: {img_path}")
+
+        logger.info(f"Loaded {len(data)} I2V samples from VBench I2V dataset")
+        return data
+
+    def _scan_directory_for_images(self, path: str) -> List[Dict[str, Any]]:
+        """Scan directory for image files."""
+        exts = ["*.jpg", "*.jpeg", "*.png", "*.webp"]
+        files = []
+
+        for ext in exts:
+            files.extend(glob.glob(os.path.join(path, ext)))
+            files.extend(glob.glob(os.path.join(path, ext.upper())))
+
+            origin_dir = os.path.join(path, "data", "origin")
+            if os.path.exists(origin_dir):
+                files.extend(glob.glob(os.path.join(origin_dir, ext)))
+                files.extend(glob.glob(os.path.join(origin_dir, ext.upper())))
+
+        return [
+            {"prompt": os.path.splitext(os.path.basename(f))[0], "image_path": f}
+            for f in files
+        ]
+
+    def _create_dummy_data(self) -> List[Dict[str, Any]]:
+        """Create dummy data with a placeholder image in cache directory."""
+        logger.info("No I2V data found. Using dummy placeholders.")
+
+        dummy_image = os.path.join(self.cache_dir, "dummy_image.jpg")
+        if not os.path.exists(dummy_image):
+            os.makedirs(self.cache_dir, exist_ok=True)
+            img = Image.new("RGB", (100, 100), color="red")
+            img.save(dummy_image)
+            logger.info(f"Created dummy image at {dummy_image}")
+
+        return [{"prompt": "A moving cat", "image_path": dummy_image}] * 10
+
+    def _load_i2v_data(self) -> List[Dict[str, Any]]:
+        """Load I2V data from VBench I2V dataset or user-provided path."""
+        path = self.args.dataset_path
+        if not path:
+            path = self._auto_download_i2v_dataset()
+            if not path:
+                return self._resize_data(self._create_dummy_data())
+
+        info_json_candidates = [
+            os.path.join(path, "data", "i2v-bench-info.json"),
+            path if path.endswith(".json") else None,
+        ]
+
+        for json_path in info_json_candidates:
+            if json_path and os.path.exists(json_path):
+                try:
+                    return self._resize_data(self._load_from_i2v_json(json_path))
+                except Exception as e:
+                    logger.info(f"Failed to load {json_path}: {e}")
+
+        if os.path.isdir(path):
+            data = self._scan_directory_for_images(path)
+            if data:
+                return self._resize_data(data)
+
+        return self._resize_data(self._create_dummy_data())
+
+    def _resize_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Resize data to match num_prompts."""
+        if not self.args.num_prompts:
+            return data
+
+        if len(data) < self.args.num_prompts:
+            factor = (self.args.num_prompts // len(data)) + 1
+            data = data * factor
+
+        return data[: self.args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        item = self.items[idx]
+        return RequestFuncInput(
+            prompt=item.get("prompt", ""),
+            api_url=self.api_url,
+            model=self.model,
+            num_outputs_per_prompt=self.args.num_outputs_per_prompt,
+            width=self.args.width,
+            height=self.args.height,
+            num_frames=self.args.num_frames,
+            fps=self.args.fps,
+            num_inference_steps=self.args.num_inference_steps,
+            image_paths=[item["image_path"]] if "image_path" in item else None,
+        )
+
+
+class RandomDataset(BaseDataset):
+    def __init__(self, args, api_url: str = "", model: str = ""):
+        super().__init__(args, api_url, model)
+        self.num_prompts = args.num_prompts or 100
+
+        self.random_request_config = args.random_request_config
+        if self.random_request_config:
+            self.random_request_config = json.loads(self.random_request_config)
+            weights = [p.pop("weight") for p in self.random_request_config]
+            seed = args.random_request_seed
+            rng = random.Random(seed)
+            self._sampled_requests = rng.choices(
+                self.random_request_config, weights=weights, k=self.num_prompts
+            )
+        else:
+            self._sampled_requests = None
+
+    def get_sampling_params(self, idx: int) -> dict:
+        """Return the per-request sampling profile dict, or empty dict if not mix-diffusion."""
+        if self._sampled_requests:
+            return self._sampled_requests[idx]
+        return {}
+
+    def __len__(self) -> int:
+        return self.num_prompts
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        profile = self._sampled_requests[idx] if self._sampled_requests else {}
+
+        return RequestFuncInput(
+            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            api_url=self.api_url,
+            model=self.model,
+            num_outputs_per_prompt=profile.get(
+                "num_outputs_per_prompt", self.args.num_outputs_per_prompt
+            ),
+            width=profile.get("width", self.args.width),
+            height=profile.get("height", self.args.height),
+            num_frames=profile.get("num_frames", self.args.num_frames),
+            num_inference_steps=profile.get(
+                "num_inference_steps", self.args.num_inference_steps
+            ),
+            fps=profile.get("fps", self.args.fps),
+        )

python/sglang/multimodal_gen/configs/__init__.py

@@ -0,0 +1,3 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# Configs for pipelines, and pipeline modules (in models folder)

python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_448_832.json

@@ -0,0 +1,16 @@
+{
+    "temporal_chunk_size": 2,
+    "temporal_topk": 2,
+    "spatial_chunk_size": [4, 13],
+    "spatial_topk": 6,
+    "st_chunk_size": [4, 4, 13],
+    "st_topk": 18,
+    "moba_select_mode": "topk",
+    "moba_threshold": 0.25,
+    "moba_threshold_type": "query_head",
+    "first_full_layer": 0,
+    "first_full_step": 12,
+    "temporal_layer": 1,
+    "spatial_layer": 1,
+    "st_layer": 1
+}

python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_480_832.json

@@ -0,0 +1,16 @@
+{
+    "temporal_chunk_size": 2,
+    "temporal_topk": 3,
+    "spatial_chunk_size": [3, 4],
+    "spatial_topk": 20,
+    "st_chunk_size": [4, 6, 4],
+    "st_topk": 15,
+    "moba_select_mode": "threshold",
+    "moba_threshold": 0.25,
+    "moba_threshold_type": "query_head",
+    "first_full_layer": 0,
+    "first_full_step": 12,
+    "temporal_layer": 1,
+    "spatial_layer": 1,
+    "st_layer": 1
+}

python/sglang/multimodal_gen/configs/models/__init__.py

@@ -0,0 +1,8 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.configs.models.base import ModelConfig
+from sglang.multimodal_gen.configs.models.dits.base import DiTConfig
+from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig
+from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
+
+__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]

python/sglang/multimodal_gen/configs/models/adapter/base.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Any
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class AdapterArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=list)
+    _compile_conditions: list = field(default_factory=list)
+
+    # convert weights name from HF-format to SGLang-dit-format
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN,
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.AITER,
+            AttentionBackendEnum.AITER_SAGE,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.SAGE_ATTN_3,
+            AttentionBackendEnum.LASER_ATTN,
+            AttentionBackendEnum.BLOCK_SPARSE_ATTN,
+            AttentionBackendEnum.RAIN_FUSION_ATTN,
+        }
+    )
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+    exclude_lora_layers: list[str] = field(default_factory=list)
+    boundary_ratio: float | None = None
+
+    def __post_init__(self) -> None:
+        if not self._compile_conditions:
+            self._compile_conditions = self._fsdp_shard_conditions.copy()
+
+
+@dataclass
+class AdapterConfig(ModelConfig):
+    arch_config: AdapterArchConfig = field(default_factory=AdapterArchConfig)
+
+    # sglang-diffusion Adapter-specific parameters
+    prefix: str = ""
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for AdapterConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=AdapterConfig.prefix,
+            help="Prefix for the Adapter",
+        )
+
+        return parser

python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.adapter.base import (
+    AdapterArchConfig,
+    AdapterConfig,
+)
+
+
+@dataclass
+class LTX2ConnectorArchConfig(AdapterArchConfig):
+    audio_connector_attention_head_dim: int = 128
+    audio_connector_num_attention_heads: int = 30
+    audio_connector_num_layers: int = 2
+    audio_connector_num_learnable_registers: int = 128
+    audio_feature_extractor_out_features: int = 0
+    caption_channels: int = 3840
+    causal_temporal_positioning: bool = False
+    connector_rope_base_seq_len: int = 4096
+    connector_apply_gated_attention: bool = False
+    feature_extractor_in_features: int = 0
+    rope_double_precision: bool = True
+    rope_theta: float = 10000.0
+    rope_type: str = "split"
+    text_proj_in_factor: int = 49
+    video_feature_extractor_out_features: int = 0
+    video_connector_attention_head_dim: int = 128
+    video_connector_num_attention_heads: int = 30
+    video_connector_num_layers: int = 2
+    video_connector_num_learnable_registers: int = 128
+
+
+@dataclass
+class LTX2ConnectorConfig(AdapterConfig):
+
+    arch_config: AdapterArchConfig = field(default_factory=LTX2ConnectorArchConfig)
+
+    prefix: str = "LTX2"

python/sglang/multimodal_gen/configs/models/base.py

@@ -0,0 +1,100 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field, fields
+from typing import Any, Dict
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
+# 2. ArchConfig should be inherited & overridden by each model arch_config
+# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
+@dataclass
+class ArchConfig:
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names
+    extra_attrs: Dict[str, Any] = field(default_factory=dict)
+
+    def __getattr__(self, name: str):
+        d = object.__getattribute__(self, "__dict__")
+        extras = d.get("extra_attrs")
+        if extras is not None and name in extras:
+            return extras[name]
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{name}'"
+        )
+
+    def __setattr__(self, key, value):
+        if key in type(self).__dataclass_fields__:
+            object.__setattr__(self, key, value)
+        else:
+            d = object.__getattribute__(self, "__dict__")
+            extras = d.get("extra_attrs")
+            if extras is None:
+                extras = {}
+                d["extra_attrs"] = extras
+            extras[key] = value
+
+
+@dataclass
+class ModelConfig:
+    # Every model config parameter can be categorized into either ArchConfig or everything else
+    # Diffuser/Transformer parameters
+    arch_config: ArchConfig = field(default_factory=ArchConfig)
+
+    # sglang-diffusion-specific parameters here
+    # i.e. STA, quantization, teacache
+
+    def __getattr__(self, name):
+        # Only called if 'name' is not found in ModelConfig directly
+        if hasattr(self.arch_config, name):
+            return getattr(self.arch_config, name)
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )
+
+    def __getstate__(self):
+        # Return a dictionary of attributes to pickle
+        # Convert to dict and exclude any problematic attributes
+        state = self.__dict__.copy()
+        return state
+
+    def __setstate__(self, state):
+        # Restore instance attributes from the unpickled state
+        self.__dict__.update(state)
+
+    # This should be used only when loading from transformers/diffusers
+    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
+        """
+        Update arch_config with source_model_dict
+        """
+        arch_config = self.arch_config
+
+        for key, value in source_model_dict.items():
+            setattr(arch_config, key, value)
+
+        if hasattr(arch_config, "__post_init__"):
+            arch_config.__post_init__()
+
+    def update_model_config(self, source_model_dict: dict[str, Any]) -> None:
+        assert (
+            "arch_config" not in source_model_dict
+        ), "Source model config shouldn't contain arch_config."
+
+        valid_fields = {f.name for f in fields(self)}
+
+        for key, value in source_model_dict.items():
+            if key in valid_fields:
+                setattr(self, key, value)
+            else:
+                logger.warning(
+                    "%s does not contain field '%s'!", type(self).__name__, key
+                )
+                raise AttributeError(f"Invalid field: {key}")
+
+        if hasattr(self, "__post_init__"):
+            self.__post_init__()

python/sglang/multimodal_gen/configs/models/bridges/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
+    MOVADualTowerConfig,
+)
+
+__all__ = ["MOVADualTowerConfig"]

python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for MOVA dual tower bridge model."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def _is_conditioner_block(name: str, module) -> bool:
+    """Check if module is a ConditionalCrossAttentionBlock."""
+    return "ConditionalCrossAttentionBlock" in type(module).__name__
+
+
+@dataclass
+class MOVADualTowerArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_conditioner_block]
+    )
+
+    # Model architecture parameters
+    visual_layers: int = 40
+    audio_layers: int = 30
+    visual_hidden_dim: int = 5120
+    audio_hidden_dim: int = 1536
+    audio_fps: float = 50.0
+    head_dim: int = 128
+    interaction_strategy: str = "full"
+    apply_cross_rope: bool = True
+    apply_first_frame_bias_in_rope: bool = False
+    trainable_condition_scale: bool = False
+    pooled_adaln: bool = False
+    eps: float = 1e-6
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.visual_hidden_dim
+        self.num_attention_heads = self.visual_hidden_dim // self.head_dim
+
+
+@dataclass
+class MOVADualTowerConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVADualTowerArchConfig)

python/sglang/multimodal_gen/configs/models/dits/__init__.py

@@ -0,0 +1,29 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
+from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
+from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
+from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
+from sglang.multimodal_gen.configs.models.dits.lingbot_world import (
+    LingBotWorldVideoConfig,
+)
+from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
+from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
+from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
+    StableDiffusion3TransformerConfig,
+)
+from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig
+
+__all__ = [
+    "Cosmos3VideoConfig",
+    "HeliosConfig",
+    "HunyuanVideoConfig",
+    "Ideogram4DiTConfig",
+    "LingBotWorldVideoConfig",
+    "WanVideoConfig",
+    "Hunyuan3DDiTConfig",
+    "MOVAAudioConfig",
+    "MOVAVideoConfig",
+    "StableDiffusion3TransformerConfig",
+]

python/sglang/multimodal_gen/configs/models/dits/base.py

@@ -0,0 +1,90 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Any
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class DiTArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=list)
+    _compile_conditions: list = field(default_factory=list)
+
+    # convert weights name from HF-format to SGLang-dit-format
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # convert weights name from misc-format to HF-format
+    # usually applicable if the LoRA is trained with official repo implementation
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN,
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.AITER,
+            AttentionBackendEnum.AITER_SAGE,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.SAGE_ATTN_3,
+            AttentionBackendEnum.LASER_ATTN,
+            AttentionBackendEnum.BLOCK_SPARSE_ATTN,
+            AttentionBackendEnum.RAIN_FUSION_ATTN,
+        }
+    )
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+    exclude_lora_layers: list[str] = field(default_factory=list)
+    boundary_ratio: float | None = None
+
+    def __post_init__(self) -> None:
+        if not self._compile_conditions:
+            self._compile_conditions = self._fsdp_shard_conditions.copy()
+
+
+@dataclass
+class DiTConfig(ModelConfig):
+    arch_config: DiTArchConfig = field(default_factory=DiTArchConfig)
+
+    # sglang-diffusion DiT-specific parameters
+    prefix: str = ""
+    quant_config: QuantizationConfig | None = None
+    torch_compile_mode: str = "max-autotune-no-cudagraphs"
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for DiTConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=DiTConfig.prefix,
+            help="Prefix for the DiT model",
+        )
+
+        parser.add_argument(
+            f"--{prefix}.quant-config",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.quant_config",
+            default=None,
+            help="Quantization configuration for the DiT model",
+        )
+        parser.add_argument(
+            f"--{prefix}.torch-compile-mode",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.torch_compile_mode",
+            default=DiTConfig.torch_compile_mode,
+            help="torch.compile mode for the DiT model",
+        )
+
+        return parser

python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 video DiT architecture configuration."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in
+
+
+def is_layers(n: str, m) -> bool:
+    return is_module_list_entry_in(n, ("layers", "gen_layers"))
+
+
+def _build_cosmos3_param_names_mapping() -> dict:
+    """Map diffusers-format Cosmos3 weights to the sglang model namespace.
+
+    Source keys (diffusers transformer ckpt) → target keys (sglang model):
+        embed_tokens.weight                                    -> language_model.embed_tokens.weight
+        layers.X.input_layernorm.weight                        -> language_model.layers.X.input_layernorm.weight
+        layers.X.input_layernorm_moe_gen.weight                -> gen_layers.X.input_layernorm.weight
+        layers.X.self_attn.{to_q,to_k,to_v}.weight            -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0)
+        layers.X.self_attn.{add_q,add_k,add_v}_proj.weight    -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0)
+        layers.X.mlp.{gate,up}_proj.weight                     -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0)
+        layers.X.mlp_moe_gen.{gate,up}_proj.weight             -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0)
+        norm_moe_gen.weight                                    -> norm_moe_gen.weight
+        time_embedder.linear_{1,2}.weight                      -> (pass-through)
+        proj_in.weight, proj_out.weight                        -> (pass-through)
+        vae2llm.weight                                         -> proj_in.weight  (FP8 ckpt alias)
+        llm2vae.weight                                         -> proj_out.weight (FP8 ckpt alias)
+
+    GEN patterns (`*_moe_gen`, `add_*`, `to_add_out`, `norm_added_*`) must
+    precede the UND catch-all so the catch-all can't claim GEN keys.
+    `norm.weight` and `lm_head.weight` are inherited from Qwen3-VL
+    pretraining and not used at inference; both are skipped.
+    Audio and action keys (``audio_proj_*``, ``action_proj_*``, modality
+    embeds) pass through unchanged.
+    """
+    return {
+        # Inherited from Qwen3-VL pretraining; unused at diffusion inference.
+        r"^lm_head\.weight$": "",
+        r"^norm\.weight$": "",
+        # Top-level norms / embeddings.
+        r"^norm_moe_gen\.(.*)$": r"norm_moe_gen.\1",
+        r"^embed_tokens\.(.*)$": r"language_model.embed_tokens.\1",
+        # FP8 checkpoint aliases for the latent projection layers.
+        r"^vae2llm\.(.*)$": r"proj_in.\1",
+        r"^llm2vae\.(.*)$": r"proj_out.\1",
+        # GEN pathway: per-layer (must run before the UND catch-all below).
+        # Q/K/V merge into MergedColumnParallelLinear to_qkv (concat order: Q, K, V).
+        r"^layers\.(\d+)\.self_attn\.add_q_proj\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            0,
+            3,
+        ),
+        r"^layers\.(\d+)\.self_attn\.add_k_proj\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            1,
+            3,
+        ),
+        r"^layers\.(\d+)\.self_attn\.add_v_proj\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            2,
+            3,
+        ),
+        r"^layers\.(\d+)\.self_attn\.to_add_out\.(.*)$": r"gen_layers.\1.cross_attention.to_out.\2",
+        r"^layers\.(\d+)\.self_attn\.norm_added_q\.(.*)$": r"gen_layers.\1.cross_attention.norm_q.\2",
+        r"^layers\.(\d+)\.self_attn\.norm_added_k\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2",
+        r"^layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2",
+        r"^layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2",
+        # GEN MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
+        # Must precede the mlp_moe_gen catch-all below.
+        r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": (
+            r"gen_layers.\1.mlp.gate_up_proj.\2",
+            0,
+            2,
+        ),
+        r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": (
+            r"gen_layers.\1.mlp.gate_up_proj.\2",
+            1,
+            2,
+        ),
+        r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2",
+        # UND pathway: Q/K/V merge into to_qkv; remaining attention keys
+        # (to_out, norm_q, norm_k) and layernorms pass through the catch-all.
+        r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            0,
+            3,
+        ),
+        r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            1,
+            3,
+        ),
+        r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            2,
+            3,
+        ),
+        # UND MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
+        # Must precede the layers catch-all below.
+        r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$": (
+            r"language_model.layers.\1.mlp.gate_up_proj.\2",
+            0,
+            2,
+        ),
+        r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$": (
+            r"language_model.layers.\1.mlp.gate_up_proj.\2",
+            1,
+            2,
+        ),
+        # UND pathway: layernorms + remaining attention/mlp keys pass through
+        # under language_model.layers namespace.
+        r"^layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2",
+    }
+
+
+@dataclass
+class Cosmos3VideoArchConfig(DiTArchConfig):
+    """Architecture config for Cosmos3 Omni Transformer.
+
+    Cosmos3 uses a dual-pathway design:
+    - Understanding (UND): Causal self-attention for text tokens
+    - Generation (GEN): Cross-attention from visual to cached UND K/V
+
+    Field names mirror the diffusers ``transformer/config.json`` so values
+    flow through ``update_model_arch`` without translation. Defaults are
+    Qwen3-8B-Instruct-derived and overridden by the checkpoint at load time.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layers])
+
+    # Transformer architecture
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8  # GQA
+    head_dim: int = 128
+    intermediate_size: int = 12288
+
+    # Latent space configuration
+    latent_patch_size: int = 2
+    latent_channel: int = 48
+    out_channels: int = 48
+
+    # RoPE configuration (Qwen3-VL 3D mRoPE: temporal, height, width)
+    mrope_section: tuple[int, int, int] = (24, 20, 20)
+    rope_theta: float = 5000000.0
+    # Populated from rope_scaling in the diffusers config when present.
+    rope_scaling: dict | None = None
+
+    # Temporal configuration
+    base_fps: float = 24.0
+    temporal_compression_factor: int = 4
+    unified_3d_mrope_temporal_modality_margin: int = 15000
+
+    # Audio (sound) modality
+    sound_gen: bool = False
+    sound_dim: int = 64
+    sound_latent_fps: float = 25.0
+    temporal_compression_factor_sound: int = 1
+
+    # Action modality
+    action_gen: bool = False
+    action_dim: int = 64
+    num_embodiment_domains: int = 32
+
+    # Timestep embedding
+    timestep_scale: float = 0.001
+    frequency_embedding_size: int = 256
+
+    # Vocab size (Qwen3-VL tokenizer)
+    vocab_size: int = 151936
+
+    # RMSNorm epsilon
+    rms_norm_eps: float = 1e-6
+
+    # Weight mapping from checkpoint to model
+    param_names_mapping: dict = field(
+        default_factory=_build_cosmos3_param_names_mapping
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+    # FP8 checkpoint quantization_config.ignore uses checkpoint module names;
+    # translate them to model names so is_layer_excluded matches correctly.
+    quant_ignore_remap: dict = field(
+        default_factory=lambda: {"vae2llm": "proj_in", "llm2vae": "proj_out"}
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Diffusers configs nest the mrope sizes under `rope_scaling`; lift it.
+        if isinstance(self.rope_scaling, dict) and "mrope_section" in self.rope_scaling:
+            self.mrope_section = tuple(self.rope_scaling["mrope_section"])
+        self.in_channels = self.latent_channel
+        self.num_channels_latents = self.out_channels
+        # Patch latent dimension: (patch_size^2) * latent_channel
+        self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel
+
+
+@dataclass
+class Cosmos3VideoConfig(DiTConfig):
+    """DiT config wrapper for Cosmos3 Video model."""
+
+    arch_config: DiTArchConfig = field(default_factory=Cosmos3VideoArchConfig)
+    prefix: str = "Cosmos3"

python/sglang/multimodal_gen/configs/models/dits/ernie_image.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_layer
+
+
+@dataclass
+class ErnieImageArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 128
+    out_channels: int = 128
+    num_layers: int = 36
+    attention_head_dim: int = 128
+    num_attention_heads: int = 32
+    ffn_hidden_size: int = 12288
+    text_in_dim: int = 3072
+    rope_theta: int = 256
+    rope_axes_dim: Tuple[int, int, int] = (32, 48, 48)
+    eps: float = 1e-6
+    qk_layernorm: bool = True
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"(.*)\.mlp\.gate_proj\.(.*)": (r"\1.mlp.gate_up_proj.\2", 0, 2),
+            r"(.*)\.mlp\.up_proj\.(.*)": (r"\1.mlp.gate_up_proj.\2", 1, 2),
+        }
+    )
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class ErnieImageDitConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=ErnieImageArchConfig)
+    prefix: str = "ernieimage"

python/sglang/multimodal_gen/configs/models/dits/flux.py

@@ -0,0 +1,116 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class FluxArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = True
+    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
+
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: [
+            "time_guidance_embed.timestep_embedder.linear_1",
+            "time_guidance_embed.timestep_embedder.linear_2",
+            "time_guidance_embed.guidance_embedder.linear_1",
+            "time_guidance_embed.guidance_embedder.linear_2",
+        ]
+    )
+
+    # nunchaku checkpoint uses different weight names; map to sglang flux layout
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # HF diffusers format: strip leading "transformer." prefix
+            r"^transformer\.(\w*)\.(.*)$": r"\1.\2",
+            # FLUX2-nvfp4 format: double blocks - image attention QKV (packed, fused)
+            r"^double_blocks\.(\d+)\.img_attn\.qkv\.(.*)$": r"transformer_blocks.\1.attn.to_qkv.\2",
+            r"^double_blocks\.(\d+)\.img_attn\.proj\.(.*)$": r"transformer_blocks.\1.attn.to_out.0.\2",
+            r"^double_blocks\.(\d+)\.img_attn\.norm\.query_norm\.(.*)$": r"transformer_blocks.\1.attn.norm_q.\2",
+            r"^double_blocks\.(\d+)\.img_attn\.norm\.key_norm\.(.*)$": r"transformer_blocks.\1.attn.norm_k.\2",
+            # FLUX2-nvfp4 format: double blocks - text/context attention QKV (packed, fused)
+            r"^double_blocks\.(\d+)\.txt_attn\.qkv\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
+            r"^double_blocks\.(\d+)\.txt_attn\.proj\.(.*)$": r"transformer_blocks.\1.attn.to_add_out.\2",
+            r"^double_blocks\.(\d+)\.txt_attn\.norm\.query_norm\.(.*)$": r"transformer_blocks.\1.attn.norm_added_q.\2",
+            r"^double_blocks\.(\d+)\.txt_attn\.norm\.key_norm\.(.*)$": r"transformer_blocks.\1.attn.norm_added_k.\2",
+            # FLUX2-nvfp4  format: double blocks - image MLP
+            r"^double_blocks\.(\d+)\.img_mlp\.0\.(.*)$": r"transformer_blocks.\1.ff.linear_in.\2",
+            r"^double_blocks\.(\d+)\.img_mlp\.2\.(.*)$": r"transformer_blocks.\1.ff.linear_out.\2",
+            # FLUX2-nvfp4  format: double blocks - text/context MLP
+            r"^double_blocks\.(\d+)\.txt_mlp\.0\.(.*)$": r"transformer_blocks.\1.ff_context.linear_in.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.2\.(.*)$": r"transformer_blocks.\1.ff_context.linear_out.\2",
+            # FLUX2-nvfp4  format: single blocks
+            r"^single_blocks\.(\d+)\.linear1\.(.*)$": r"single_transformer_blocks.\1.attn.to_qkv_mlp_proj.\2",
+            r"^single_blocks\.(\d+)\.linear2\.(.*)$": r"single_transformer_blocks.\1.attn.to_out.\2",
+            r"^single_blocks\.(\d+)\.norm\.query_norm\.(.*)$": r"single_transformer_blocks.\1.attn.norm_q.\2",
+            r"^single_blocks\.(\d+)\.norm\.key_norm\.(.*)$": r"single_transformer_blocks.\1.attn.norm_k.\2",
+            # FLUX2-nvfp4  format: non-block input/output projections
+            r"^img_in\.(.*)$": r"x_embedder.\1",
+            r"^txt_in\.(.*)$": r"context_embedder.\1",
+            r"^time_in\.in_layer\.(.*)$": r"time_guidance_embed.timestep_embedder.linear_1.\1",
+            r"^time_in\.out_layer\.(.*)$": r"time_guidance_embed.timestep_embedder.linear_2.\1",
+            r"^guidance_in\.in_layer\.(.*)$": r"time_guidance_embed.guidance_embedder.linear_1.\1",
+            r"^guidance_in\.out_layer\.(.*)$": r"time_guidance_embed.guidance_embedder.linear_2.\1",
+            r"^double_stream_modulation_img\.lin\.(.*)$": r"double_stream_modulation_img.linear.\1",
+            r"^double_stream_modulation_txt\.lin\.(.*)$": r"double_stream_modulation_txt.linear.\1",
+            r"^single_stream_modulation\.lin\.(.*)$": r"single_stream_modulation.linear.\1",
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"norm_out.linear.\1",
+            r"^final_layer\.linear\.(.*)$": r"proj_out.\1",
+            # FLUX2-nvfp4 format: RMSNorm uses "scale" parameter; rename to "weight" (model uses .weight)
+            r"^(.*)\.scale$": r"\1.weight",
+            # transformer_blocks nunchaku format (raw export - before internal conversion)
+            r"^transformer_blocks\.(\d+)\.mlp_fc1\.(.*)$": r"transformer_blocks.\1.ff.net.0.proj.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_fc2\.(.*)$": r"transformer_blocks.\1.ff.net.2.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_context_fc1\.(.*)$": r"transformer_blocks.\1.ff_context.net.0.proj.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_context_fc2\.(.*)$": r"transformer_blocks.\1.ff_context.net.2.\2",
+            # nunchaku packed QKV → fused to_qkv / to_added_qkv (matches use_fused_qkv in model)
+            r"^transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_qkv.\2",
+            r"^transformer_blocks\.(\d+)\.qkv_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
+            r"^transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"transformer_blocks.\1.attn.to_out.0.\2",
+            r"^transformer_blocks\.(\d+)\.out_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_add_out.\2",
+            r"^transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"transformer_blocks.\1.attn.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"transformer_blocks.\1.attn.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.norm_added_q\.(.*)$": r"transformer_blocks.\1.attn.norm_added_q.\2",
+            r"^transformer_blocks\.(\d+)\.norm_added_k\.(.*)$": r"transformer_blocks.\1.attn.norm_added_k.\2",
+            # nunchaku format (already converted): add_qkv_proj → fused to_added_qkv
+            r"^transformer_blocks\.(\d+)\.attn\.add_qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
+            # single_transformer_blocks nunchaku format (raw export - before internal conversion)
+            r"^single_transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_qkv_mlp_proj.\2",
+            r"^single_transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_out.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"single_transformer_blocks.\1.attn.norm_q.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"single_transformer_blocks.\1.attn.norm_k.\2",
+            # nunchaku quantization parameter name conversions (apply to all blocks)
+            r"^(.*)\.smooth_orig$": r"\1.smooth_factor_orig",
+            r"^(.*)\.smooth$": r"\1.smooth_factor",
+            r"^(.*)\.lora_down$": r"\1.proj_down",
+            r"^(.*)\.lora_up$": r"\1.proj_up",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class FluxConfig(DiTConfig):
+
+    arch_config: DiTArchConfig = field(default_factory=FluxArchConfig)
+
+    prefix: str = "Flux"

python/sglang/multimodal_gen/configs/models/dits/glmimage.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class GlmImageArchConfig(DiTArchConfig):
+    patch_size: int = 2
+    in_channels: int = 16
+    out_channels: int | None = 16
+    num_layers: int = 30
+    attention_head_dim: int = 128
+    num_attention_heads: int = 32
+    condition_dim: int = 256
+    prior_vq_quantizer_codebook_size: int = 16384
+    text_embed_dim: int = 1472
+    time_embed_dim: int = 512
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # LoRA mappings
+            r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class GlmImageDitConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=GlmImageArchConfig)
+
+    prefix: str = "glmimage"

python/sglang/multimodal_gen/configs/models/dits/helios.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_block
+
+
+@dataclass
+class HeliosArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Patch embeddings
+            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+            # Condition embedder: text
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            # Condition embedder: time
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            # Blocks: self-attention (keep attn1. prefix, drop .0. from to_out)
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.attn1.to_out.\2",
+            # Blocks: cross-attention output (drop .0. from to_out)
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            # Blocks: feed-forward
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            # Blocks: cross-attn residual norm
+            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.\2",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 226
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 13824
+    num_layers: int = 40
+    cross_attn_norm: bool = True
+    qk_norm: str = "rms_norm_across_heads"
+    eps: float = 1e-6
+    added_kv_proj_dim: int | None = None
+    rope_max_seq_len: int = 1024
+    pos_embed_seq_len: int | None = None
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    # Helios-specific
+    rope_dim: tuple[int, int, int] = (44, 42, 42)
+    rope_theta: float = 10000.0
+    guidance_cross_attn: bool = True
+    zero_history_timestep: bool = True
+    has_multi_term_memory_patch: bool = True
+    is_amplify_history: bool = False
+    history_scale_mode: str = "per_head"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class HeliosConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig)
+
+    prefix: str = "Helios"

python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class Hunyuan3DDiTArchConfig(DiTArchConfig):
+    """Architecture config for Hunyuan3D DiT (Flux-style for Hunyuan3D-2.0)."""
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Strip leading "model." prefix used by some exports
+            r"^model\.(.*)$": r"\1",
+            # MLP linear renames (double-stream blocks)
+            r"^(double_blocks\.\d+\.img_mlp)\.0\.(.*)$": r"\1.fc_in.\2",
+            r"^(double_blocks\.\d+\.img_mlp)\.2\.(.*)$": r"\1.fc_out.\2",
+            r"^(double_blocks\.\d+\.txt_mlp)\.0\.(.*)$": r"\1.fc_in.\2",
+            r"^(double_blocks\.\d+\.txt_mlp)\.2\.(.*)$": r"\1.fc_out.\2",
+            # Double-stream attention: fuse split Q/K/V into fused qkv for both txt_attn and img_attn
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.(?:to_q|q_proj|query)\.(.*)$": (
+                r"\1.qkv.\2",
+                0,
+                3,
+            ),
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.(?:to_k|k_proj|key)\.(.*)$": (
+                r"\1.qkv.\2",
+                1,
+                3,
+            ),
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.(?:to_v|v_proj|value)\.(.*)$": (
+                r"\1.qkv.\2",
+                2,
+                3,
+            ),
+            # Double-stream out projection (image/text): to_out[.0].{weight,bias} and
+            # txt_attn.to_add_out[.0].{weight,bias} -> proj.{weight,bias}
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.to_out(?:\.0)?\.(weight|bias)$": r"\1.proj.\2",
+            r"^(double_blocks\.\d+\.txt_attn)\.to_add_out(?:\.0)?\.(weight|bias)$": r"\1.proj.\2",
+            # Double-stream Q/K norm aliases and convert HF 'weight' to internal 'scale'
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.norm_q\.(.*)$": r"\1.norm.query_norm.\2",
+            r"^(double_blocks\.\d+\.(?:txt_attn|img_attn))\.norm_k\.(.*)$": r"\1.norm.key_norm.\2",
+            r"^(.*norm\.query_norm)\.weight$": r"\1.scale",
+            r"^(.*norm\.key_norm)\.weight$": r"\1.scale",
+            # Single-stream blocks: pack Q/K/V and MLP into linear1 ([Q, K, V, MLP]) and map out-proj to linear2
+            # Apply to both single_blocks.* and single_transformer_blocks.* exports
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.(?:to_q|q_proj|query)\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                0,
+                4,
+            ),
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.(?:to_k|k_proj|key)\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                1,
+                4,
+            ),
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.(?:to_v|v_proj|value)\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                2,
+                4,
+            ),
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.(?:proj_mlp|mlp_fc1)\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                3,
+                4,
+            ),
+            # Single-stream out projection variants -> linear2 (only weight/bias)
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.(?:proj_out|out_proj)(?:\.0)?\.(weight|bias)$": r"single_blocks.\1.linear2.\2",
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.to_out(?:\.0)?\.(weight|bias)$": r"single_blocks.\1.linear2.\2",
+            # Single-stream Q/K norm aliases
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.norm_q\.(.*)$": r"single_blocks.\1.norm.query_norm.\2",
+            r"^(?:single_blocks|single_transformer_blocks)\.(\d+)\.attn\.norm_k\.(.*)$": r"single_blocks.\1.norm.key_norm.\2",
+        }
+    )
+
+    in_channels: int = 64
+    hidden_size: int = 1024
+    num_attention_heads: int = 16
+    num_layers: int = 16
+    num_single_layers: int = 32
+    mlp_ratio: float = 4.0
+    context_in_dim: int = 1536
+    axes_dim: tuple[int, ...] = (64,)
+    theta: int = 10000
+    qkv_bias: bool = True
+    guidance_embed: bool = False
+    time_factor: float = 1000.0
+
+    def __post_init__(self) -> None:
+        if self.num_channels_latents == 0:
+            self.num_channels_latents = self.in_channels
+        super().__post_init__()
+
+
+@dataclass
+class Hunyuan3DDiTConfig(DiTConfig):
+    """DiT configuration for Hunyuan3D shape generation (Flux-style)."""
+
+    arch_config: Hunyuan3DDiTArchConfig = field(default_factory=Hunyuan3DDiTArchConfig)
+    subfolder: str = "hunyuan3d-dit-v2-0"

python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py

@@ -0,0 +1,174 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_double_block,
+    is_refiner_block,
+    is_single_block,
+    is_txt_in,
+)
+
+
+@dataclass
+class HunyuanVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_single_block, is_refiner_block]
+    )
+
+    _compile_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_single_block, is_txt_in]
+    )
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1",
+            r"^context_embedder\.proj_in\.(.*)$": r"txt_in.input_embedder.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$": r"txt_in.refiner_blocks.\1.norm1.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$": r"txt_in.refiner_blocks.\1.norm2.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
+                r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
+                0,
+                3,
+            ),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
+                r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
+                1,
+                3,
+            ),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
+                r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
+                2,
+                3,
+            ),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+            # 3. x_embedder mapping:
+            r"^x_embedder\.proj\.(.*)$": r"img_in.proj.\1",
+            # 4. Top-level time_text_embed mappings:
+            r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"time_in.mlp.fc_in.\1",
+            r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"time_in.mlp.fc_out.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$": r"guidance_in.mlp.fc_in.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$": r"guidance_in.mlp.fc_out.\1",
+            r"^time_text_embed\.text_embedder\.linear_1\.(.*)$": r"vector_in.fc_in.\1",
+            r"^time_text_embed\.text_embedder\.linear_2\.(.*)$": r"vector_in.fc_out.\1",
+            # 5. transformer_blocks mapping:
+            r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$": r"double_blocks.\1.img_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$": r"double_blocks.\1.txt_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
+                r"double_blocks.\1.img_attn_qkv.\2",
+                0,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
+                r"double_blocks.\1.img_attn_qkv.\2",
+                1,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
+                r"double_blocks.\1.img_attn_qkv.\2",
+                2,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": (
+                r"double_blocks.\1.txt_attn_qkv.\2",
+                0,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": (
+                r"double_blocks.\1.txt_attn_qkv.\2",
+                1,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": (
+                r"double_blocks.\1.txt_attn_qkv.\2",
+                2,
+                3,
+            ),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"double_blocks.\1.img_attn_proj.\2",
+            # Corrected: merge attn.to_add_out into the main projection.
+            r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2",
+            # 6. single_transformer_blocks mapping:
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"single_blocks.\1.q_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"single_blocks.\1.k_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                0,
+                4,
+            ),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                1,
+                4,
+            ),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                2,
+                4,
+            ),
+            r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": (
+                r"single_blocks.\1.linear1.\2",
+                3,
+                4,
+            ),
+            # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
+            r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$": r"single_blocks.\1.linear2.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$": r"single_blocks.\1.modulation.linear.\2",
+            # 7. Final layers mapping:
+            r"^norm_out\.linear\.(.*)$": r"final_layer.adaLN_modulation.linear.\1",
+            r"^proj_out\.(.*)$": r"final_layer.linear.\1",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: int = 2
+    patch_size_t: int = 1
+    in_channels: int = 16
+    out_channels: int = 16
+    num_attention_heads: int = 24
+    attention_head_dim: int = 128
+    mlp_ratio: float = 4.0
+    num_layers: int = 20
+    num_single_layers: int = 40
+    num_refiner_layers: int = 2
+    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
+    guidance_embeds: bool = False
+    dtype: torch.dtype | None = None
+    text_embed_dim: int = 4096
+    pooled_projection_dim: int = 768
+    rope_theta: int = 256
+    qk_norm: str = "rms_norm"
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"]
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
+        self.num_channels_latents: int = self.in_channels
+
+
+@dataclass
+class HunyuanVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig)
+
+    prefix: str = "Hunyuan"

python/sglang/multimodal_gen/configs/models/dits/ideogram.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_layer
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class Ideogram4DiTArchConfig(DiTArchConfig):
+    adaln_dim: int = 512
+    attention_head_dim: int = 256
+    in_channels: int = 128
+    intermediate_size: int = 12288
+    llm_features_dim: int = 53248
+    mrope_section: tuple[int, int, int] | list[int] = (24, 20, 20)
+    norm_eps: float = 1e-5
+    num_attention_heads: int = 18
+    num_layers: int = 34
+    rope_theta: int = 5_000_000
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.TORCH_SDPA,
+        }
+    )
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.in_channels
+
+
+@dataclass
+class Ideogram4DiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=Ideogram4DiTArchConfig)
+    prefix: str = "ideogram4"

python/sglang/multimodal_gen/configs/models/dits/joy_echo.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.ltx_2 import (
+    LTX2ArchConfig,
+    LTX2Config,
+)
+
+
+@dataclass
+class JoyEchoArchConfig(LTX2ArchConfig):
+    """JoyEcho DiT architecture config (LTX-2.3 AV base)."""
+
+    caption_proj_before_connector: bool = True
+    cross_attention_adaln: bool = True
+    apply_gated_attention: bool = True
+
+
+@dataclass
+class JoyEchoConfig(LTX2Config):
+    arch_config: JoyEchoArchConfig = field(default_factory=JoyEchoArchConfig)
+    prefix: str = "JoyEcho"

python/sglang/multimodal_gen/configs/models/dits/joy_image.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_double_blocks
+
+
+@dataclass
+class JoyImageArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_double_blocks]
+    )
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Condition embedder mappings
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            # Double blocks mappings
+            r"^double_blocks\.(\d+)\.attn\.(.*)$": r"double_blocks.\1.\2",
+            r"^double_blocks\.(\d+)\.img_mlp\.net\.0\.proj\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.img_mlp\.net\.2\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.net\.0\.proj\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.net\.2\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2",
+            r"^double_blocks\.(\d+)\.img_attn_qkv\.(.*)$": r"double_blocks.\1.img_attn_qkv.\2",
+            r"^double_blocks\.(\d+)\.txt_attn_qkv\.(.*)$": r"double_blocks.\1.txt_attn_qkv.\2",
+            r"^double_blocks\.(\d+)\.img_attn_proj\.(.*)$": r"double_blocks.\1.img_attn_proj.\2",
+            r"^double_blocks\.(\d+)\.txt_attn_proj\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2",
+            r"^double_blocks\.(\d+)\.img_mod\.(.*)$": r"double_blocks.\1.img_mod.\2",
+            r"^double_blocks\.(\d+)\.txt_mod\.(.*)$": r"double_blocks.\1.txt_mod.\2",
+            r"^double_blocks\.(\d+)\.img_attn_q_norm\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2",
+            r"^double_blocks\.(\d+)\.img_attn_k_norm\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2",
+            r"^double_blocks\.(\d+)\.txt_attn_q_norm\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2",
+            r"^double_blocks\.(\d+)\.txt_attn_k_norm\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Model architecture parameters
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    mm_double_blocks_depth: int = 40
+    freq_dim: int = 256
+    text_states_dim: int = 4096
+    mlp_width_ratio: float = 4.0
+    rope_theta: int = 10000
+    rope_dim_list: list[int] = field(default_factory=lambda: [16, 56, 56])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class JoyImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=JoyImageArchConfig)
+    prefix: str = "JoyImage"

python/sglang/multimodal_gen/configs/models/dits/krea2.py

@@ -0,0 +1,77 @@
+# Krea-2 (K2) single-stream MMDiT architecture config.
+#
+# Parameter names follow the released K2 checkpoint, so the MMDiT safetensors load
+# without remapping (identity `param_names_mapping`).
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class Krea2ArchConfig(DiTArchConfig):
+    features: int = 6144  # hidden dim
+    tdim: int = 256  # timestep embedding dim
+    txtdim: int = 2560  # text-encoder hidden dim (Qwen3-VL-4B hidden_size)
+    heads: int = 48
+    kvheads: int = 12  # GQA 4:1
+    multiplier: int = 4  # SwiGLU expansion multiplier
+    layers: int = 28
+    patch: int = 2
+    channels: int = 16  # VAE latent channels
+    bias: bool = False
+    theta: float = 1e3  # RoPE theta
+    txtlayers: int = 12  # number of text-encoder hidden-state layers fused by txtfusion
+    txtheads: int = 20
+    txtkvheads: int = 20
+
+    # 3-axis RoPE split over head_dim=128: [global, h, w] = (32, 48, 48).
+    axes_dims: Tuple[int, int, int] = (32, 48, 48)
+
+    # Joint (text+image) sequence is padded to a multiple of this many tokens.
+    seq_multiple_of: int = 256
+
+    # Packed patch-token width (channels * patch**2); used by the VAE unpack path.
+    in_channels: int = 64
+
+    # BaseDiT-required instance attrs (overwritten in __post_init__).
+    hidden_size: int = 6144
+    num_attention_heads: int = 48
+    num_channels_latents: int = 16
+
+    # Module/parameter names match the released checkpoint, so weights load with an
+    # identity mapping.
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # Diffusers LoRA checkpoints prefix every DiT key with the pipeline component name
+    # (e.g. transformer.transformer_blocks.0.attn.to_q.lora_A). Strip that prefix so the
+    # LoRA keys line up with this model's module names. Only the LoRA path uses this; the
+    # main checkpoint still loads with the identity param_names_mapping above.
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {r"^transformer\.": ""}
+    )
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.hidden_size = self.features
+        self.num_attention_heads = self.heads
+        self.num_channels_latents = self.channels
+        assert self.features % self.heads == 0
+        assert (
+            sum(self.axes_dims) == self.features // self.heads
+        ), f"sum(axes_dims)={sum(self.axes_dims)} != head_dim={self.features // self.heads}"
+
+    @property
+    def head_dim(self) -> int:
+        return self.features // self.heads
+
+    @property
+    def in_features_packed(self) -> int:
+        """Patch-embed input width: channels * patch**2."""
+        return self.channels * self.patch**2
+
+
+@dataclass
+class Krea2DitConfig(DiTConfig):
+    arch_config: Krea2ArchConfig = field(default_factory=Krea2ArchConfig)
+    prefix: str = "k2"

python/sglang/multimodal_gen/configs/models/dits/lingbot_world.py

@@ -0,0 +1,91 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# Adapted from: https://github.com/Robbyant/lingbot-world
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class LingBotWorldArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+            r"^patch_embedding_wancamctrl\.(.*)$": r"patch_embedding_wancamctrl.proj.\1",
+            r"^c2ws_hidden_states_layer1\.(.*)$": r"c2ws_mlp.fc_in.\1",
+            r"^c2ws_hidden_states_layer2\.(.*)$": r"c2ws_mlp.fc_out.\1",
+            r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.cam_injector_layer1\.(.*)$": r"blocks.\1.cam_conditioner.cam_injector.fc_in.\2",
+            r"^blocks\.(\d+)\.cam_injector_layer2\.(.*)$": r"blocks.\1.cam_conditioner.cam_injector.fc_out.\2",
+            r"^blocks\.(\d+)\.cam_scale_layer\.(.*)$": r"blocks.\1.cam_conditioner.cam_scale_layer.\2",
+            r"^blocks\.(\d+)\.cam_shift_layer\.(.*)$": r"blocks.\1.cam_conditioner.cam_shift_layer.\2",
+            r"^head\.modulation$": r"scale_shift_table",
+            r"^head\.head\.(.*)$": r"proj_out.\1",
+        }
+    )
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 512
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 36
+    out_channels: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 13824
+    num_layers: int = 40
+    cross_attn_norm: bool = True
+    qk_norm: str = "rms_norm_across_heads"
+    eps: float = 1e-6
+    image_dim: int | None = None
+    added_kv_proj_dim: int | None = None
+    rope_max_seq_len: int = 1024
+    pos_embed_seq_len: int | None = None
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+    boundary_ratio: float | None = None
+    local_attn_size: int = -1
+    sink_size: int = 3
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 45
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class LingBotWorldVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=LingBotWorldArchConfig)
+
+    prefix: str = "Wan"

python/sglang/multimodal_gen/configs/models/dits/ltx_2.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from enum import Enum
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
+
+
+class LTXModelType(Enum):
+    """
+    Model type enum mirroring upstream `LTXModelType`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/model.py::LTXModelType`
+    """
+
+    AudioVideo = "ltx av model"
+    VideoOnly = "ltx video only model"
+    AudioOnly = "ltx audio only model"
+
+    def is_video_enabled(self) -> bool:
+        return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
+
+    def is_audio_enabled(self) -> bool:
+        return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
+
+
+class LTX2RopeType(str, Enum):
+    """
+    Minimal RoPE type enum mirroring LTX-2 upstream `LTXRopeType`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/rope.py::LTXRopeType`
+    """
+
+    INTERLEAVED = "interleaved"
+    SPLIT = "split"
+
+
+class LTX2AttentionFunction(str, Enum):
+    """
+    Placeholder enum for upstream `AttentionFunction.DEFAULT`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/attention.py`
+    """
+
+    DEFAULT = "default"
+
+
+@dataclass
+class LTX2ArchConfig(DiTArchConfig):
+    """Architecture configuration for LTX-2 Video Transformer."""
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_transformer_blocks]
+    )
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Parameter name mappings from HuggingFace checkpoint keys to SGLang module names.
+            # We use upstream variable names (patchify_proj, adaln_single) but HF uses different keys.
+            #
+            # HF key -> SGLang key (upstream naming)
+            r"^model\.diffusion_model\.(.*)$": r"\1",
+            r"^proj_in\.(.*)$": r"patchify_proj.\1",
+            r"^time_embed\.(.*)$": r"adaln_single.\1",
+            r"^audio_proj_in\.(.*)$": r"audio_patchify_proj.\1",
+            r"^audio_time_embed\.(.*)$": r"audio_adaln_single.\1",
+            # FeedForward
+            r"(.*)ff\.net\.0\.proj\.(.*)$": r"\1ff.proj_in.\2",
+            r"(.*)ff\.net\.2\.(.*)$": r"\1ff.proj_out.\2",
+            # Attention Norms
+            r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2",
+            r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2",
+            # Scale Shift Tables (Global)
+            r"^av_cross_attn_video_scale_shift\.(.*)$": r"av_ca_video_scale_shift_adaln_single.\1",
+            r"^av_cross_attn_audio_scale_shift\.(.*)$": r"av_ca_audio_scale_shift_adaln_single.\1",
+            r"^av_cross_attn_video_a2v_gate\.(.*)$": r"av_ca_a2v_gate_adaln_single.\1",
+            r"^av_cross_attn_audio_v2a_gate\.(.*)$": r"av_ca_v2a_gate_adaln_single.\1",
+            # Scale Shift Tables (Block Level)
+            # HF: scale_shift_table_a2v_ca_video -> SGLang: video_a2v_cross_attn_scale_shift_table
+            r"(.*)scale_shift_table_a2v_ca_video": r"\1video_a2v_cross_attn_scale_shift_table",
+            r"(.*)scale_shift_table_a2v_ca_audio": r"\1audio_a2v_cross_attn_scale_shift_table",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Reverse mapping: SGLang module names -> HF checkpoint keys (for saving).
+            r"^patchify_proj\.(.*)$": r"proj_in.\1",
+            r"^adaln_single\.(.*)$": r"time_embed.\1",
+            r"^audio_patchify_proj\.(.*)$": r"audio_proj_in.\1",
+            r"^audio_adaln_single\.(.*)$": r"audio_time_embed.\1",
+            # FeedForward
+            r"(.*)ff\.proj_in\.(.*)$": r"\1ff.net.0.proj.\2",
+            r"(.*)ff\.proj_out\.(.*)$": r"\1ff.net.2.\2",
+            # Attention Norms
+            r"(.*)\.q_norm\.(.*)$": r"\1.norm_q.\2",
+            r"(.*)\.k_norm\.(.*)$": r"\1.norm_k.\2",
+            # Scale Shift Tables (Global)
+            r"^av_ca_video_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_video_scale_shift.\1",
+            r"^av_ca_audio_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_audio_scale_shift.\1",
+            r"^av_ca_a2v_gate_adaln_single\.(.*)$": r"av_cross_attn_video_a2v_gate.\1",
+            r"^av_ca_v2a_gate_adaln_single\.(.*)$": r"av_cross_attn_audio_v2a_gate.\1",
+            # Scale Shift Tables (Block Level)
+            # SGLang: video_a2v_cross_attn_scale_shift_table -> HF: scale_shift_table_a2v_ca_video
+            r"(.*)video_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_video",
+            r"(.*)audio_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_audio",
+        }
+    )
+
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # LoRA parameter name mappings from official repo format to HF format.
+            # This is applied before param_names_mapping when loading LoRA adapters.
+            # Will be populated if LoRA adapters use different naming conventions.
+        }
+    )
+
+    # Model type and attention configuration
+    model_type: LTXModelType = LTXModelType.AudioVideo
+    attention_type: LTX2AttentionFunction = LTX2AttentionFunction.DEFAULT
+    rope_type: LTX2RopeType = LTX2RopeType.INTERLEAVED
+    double_precision_rope: bool = False
+    quantize_video_rope_coords_to_hidden_dtype: bool = False
+    apply_gated_attention: bool = False
+    cross_attention_adaln: bool = False
+    caption_proj_before_connector: bool = False
+
+    # Video parameters
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 128
+    out_channels: int = 128
+    num_layers: int = 48
+    cross_attention_dim: int = 4096
+    norm_eps: float = 1e-6
+    caption_channels: int = 3840
+    positional_embedding_theta: float = 10000.0
+    positional_embedding_max_pos: list[int] | None = None
+    timestep_scale_multiplier: int = 1000
+    use_middle_indices_grid: bool = True
+
+    # Audio parameters
+    audio_num_attention_heads: int = 32
+    audio_attention_head_dim: int = 64
+    audio_in_channels: int = 128
+    audio_out_channels: int = 128
+    audio_cross_attention_dim: int = 2048
+    audio_positional_embedding_max_pos: list[int] | None = None
+    av_ca_timestep_scale_multiplier: int = 1
+
+    # 2.3 connector-related fields may show up in transformer/config.json.
+    connector_attention_head_dim: int = 128
+    connector_num_attention_heads: int = 30
+    connector_num_layers: int = 2
+    audio_connector_attention_head_dim: int = 128
+    audio_connector_num_attention_heads: int = 30
+    audio_connector_num_layers: int = 2
+
+    # SGLang-specific parameters
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 512
+    enable_packed_qkv_input_a2a: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Video derived values
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+        if self.positional_embedding_max_pos is None:
+            self.positional_embedding_max_pos = [20, 2048, 2048]
+
+        # Audio derived values
+        self.audio_hidden_size = (
+            self.audio_num_attention_heads * self.audio_attention_head_dim
+        )
+        if self.audio_positional_embedding_max_pos is None:
+            self.audio_positional_embedding_max_pos = [20]
+
+
+@dataclass
+class LTX2Config(DiTConfig):
+    """Configuration for LTX-2 Video Transformer."""
+
+    arch_config: LTX2ArchConfig = field(default_factory=LTX2ArchConfig)
+
+    prefix: str = "ltx2"
+    torch_compile_mode: str = "default"

python/sglang/multimodal_gen/configs/models/dits/mova_audio.py

@@ -0,0 +1,64 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_block
+
+
+@dataclass
+class MOVAAudioArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2",
+            r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1",
+            r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1",
+            r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1",
+        }
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    dim: int = 1536
+    in_dim: int = 128
+    ffn_dim: int = 6144
+    out_dim: int = 128
+    text_dim: int = 4096
+    freq_dim: int = 256
+    eps: float = 1e-6
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_heads: int = 12
+    num_layers: int = 30
+    has_image_input: bool = False
+    has_image_pos_emb: bool = False
+    has_ref_conv: bool = False
+    add_control_adapter: bool = False
+    in_dim_control_adapter: int = 24
+    separated_timestep: bool = False
+    require_vae_embedding: bool = False
+    require_clip_embedding: bool = False
+    fuse_vae_embedding_in_latents: bool = False
+    vae_type: str = "dac"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.num_channels_latents = self.out_dim
+        assert (
+            not self.has_image_input
+        ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)."
+
+
+@dataclass
+class MOVAAudioConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVAAudioArchConfig)
+    prefix: str = "mova_audio"

python/sglang/multimodal_gen/configs/models/dits/mova_video.py

@@ -0,0 +1,63 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_block
+
+
+@dataclass
+class MOVAVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2",
+            r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1",
+            r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1",
+            r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1",
+        }
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    dim: int = 5120
+    in_dim: int = 16
+    ffn_dim: int = 13824
+    out_dim: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    eps: float = 1e-6
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_heads: int = 40
+    num_layers: int = 40
+    has_image_input: bool = False
+    has_image_pos_emb: bool = False
+    has_ref_conv: bool = False
+    add_control_adapter: bool = False
+    in_dim_control_adapter: int = 24
+    separated_timestep: bool = False
+    require_vae_embedding: bool = True
+    require_clip_embedding: bool = True
+    fuse_vae_embedding_in_latents: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.num_channels_latents = self.out_dim
+        assert (
+            not self.has_image_input
+        ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)."
+
+
+@dataclass
+class MOVAVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVAVideoArchConfig)
+    prefix: str = "mova_video"

python/sglang/multimodal_gen/configs/models/dits/qwenimage.py

@@ -0,0 +1,66 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
+
+
+@dataclass
+class QwenImageArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = False
+    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
+    zero_cond_t: bool = False
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_block])
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # LoRA mappings
+            r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
+            # SVDquant mappings
+            r"(.*)\.add_qkv_proj\.(.+)$": r"\1.to_added_qkv.\2",
+            r"(transformer_blocks\.\d+\.(img_mlp|txt_mlp)\..*\.(smooth_factor_orig|wcscales))$": r"\1",
+            r".*\.wtscale$": r"",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class QwenImageEditPlus_2511_ArchConfig(QwenImageArchConfig):
+    zero_cond_t: bool = True
+
+
+@dataclass
+class QwenImageDitConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=QwenImageArchConfig)
+
+    prefix: str = "qwenimage"
+
+
+@dataclass
+class QwenImageEditPlus_2511_DitConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(
+        default_factory=QwenImageEditPlus_2511_ArchConfig
+    )
+
+    prefix: str = "qwenimageedit"

python/sglang/multimodal_gen/configs/models/dits/sana.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Architecture and model configuration for SANA DiT (Diffusion Transformer).
+#
+# SANA uses a linear-attention-based transformer that replaces standard
+# quadratic self-attention with ReLU-based linear attention, enabling
+# efficient high-resolution image synthesis. Cross-attention (standard SDPA)
+# is used for text conditioning via Gemma2 embeddings.
+#
+# Defaults below correspond to the SANA-1.6B / 1024px variant.
+# For 4.8B, override num_layers=36, num_attention_heads=64, etc.
+#
+# Reference: https://arxiv.org/abs/2410.10629
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class SanaArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 32
+    out_channels: int = 32
+    num_layers: int = 20
+    attention_head_dim: int = 32
+    num_attention_heads: int = 70
+    num_cross_attention_heads: int = 20
+    cross_attention_head_dim: int = 112
+    cross_attention_dim: int = 2240
+    caption_channels: int = 2304
+
+    mlp_ratio: float = 2.5
+    # "rms_norm_across_heads" applies RMSNorm over the full (num_heads * head_dim)
+
+    qk_norm: str = "rms_norm_across_heads"
+    norm_elementwise_affine: bool = False
+    norm_eps: float = 1e-6
+    sample_size: int = 32
+    guidance_embeds: bool = False
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # self linear-attn: merge q/k/v into to_qkv (concat order q, k, v)
+            r"^(transformer_blocks\.\d+\.attn1)\.to_q\.(.*)$": (r"\1.to_qkv.\2", 0, 3),
+            r"^(transformer_blocks\.\d+\.attn1)\.to_k\.(.*)$": (r"\1.to_qkv.\2", 1, 3),
+            r"^(transformer_blocks\.\d+\.attn1)\.to_v\.(.*)$": (r"\1.to_qkv.\2", 2, 3),
+            # cross-attn: merge k/v into to_kv (q stays separate)
+            r"^(transformer_blocks\.\d+\.attn2)\.to_k\.(.*)$": (r"\1.to_kv.\2", 0, 2),
+            r"^(transformer_blocks\.\d+\.attn2)\.to_v\.(.*)$": (r"\1.to_kv.\2", 1, 2),
+            r"^transformer\.(.*)$": r"\1",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaArchConfig)
+    prefix: str = "Sana"

python/sglang/multimodal_gen/configs/models/dits/sana_wm.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
+
+
+@dataclass
+class SanaWMArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_transformer_blocks]
+    )
+
+    # --- Core dims (upstream: depth=20, hidden=2240, heads=20, linear_head_dim=112) ---
+    patch_size: int = 1
+    in_channels: int = 128  # LTX-2 VAE latent channels
+    out_channels: int = 128
+    num_layers: int = 20
+
+    # Patch embedder uses (1, patch_size, patch_size) — temporal patch is always 1.
+    patch_size_t: int = 1
+
+    num_attention_heads: int = 20
+    attention_head_dim: int = 112  # = linear_head_dim
+    linear_head_dim: int = 112
+
+    # --- Cross-attention (text conditioning) ---
+    # In upstream, cross-attn uses num_heads (=20) with head_dim = hidden/num_heads = 112.
+    num_cross_attention_heads: int = 20
+    cross_attention_head_dim: int = 112
+    cross_attention_dim: int = 2240  # query dim used inside MultiHeadCrossAttention
+    cross_norm: bool = True
+
+    # Gemma-2-2b-it hidden size (input to y_embedder.y_proj).
+    caption_channels: int = 2304
+    model_max_length: int = 300
+    y_norm: bool = True
+    y_norm_scale_factor: float = 0.01
+    y_norm_eps: float = 1e-5
+
+    mlp_ratio: float = 3.0
+    qk_norm: bool = True
+    norm_eps: float = 1e-6
+    timestep_norm_scale_factor: float = 1.0
+
+    # --- Hybrid GDN/Softmax attention ---
+    # softmax_every_n=4 => blocks where (i+1)%4 == 0 use softmax main branch,
+    # i.e. block indices {3, 7, 11, 15, 19}.
+    softmax_every_n: int = 4
+
+    # --- GDN ShortConvolution params ---
+    conv_kernel_size: int = 4
+    k_conv_only: bool = True
+    chunk_gdn_chunk_size: int = 21
+    update_rule: str = "torch_chunk"  # main branch update rule
+    cam_update_rule: str = "torch_chunk"  # camera branch update rule
+    # main GDN scan backend: "auto" uses the SANA-WM Triton fast path on
+    # supported CUDA inference runs, otherwise falls back to the torch scan.
+    gdn_backend: str = "auto"
+
+    # --- Camera conditioning ---
+    cam_attn_compress: int = 1  # cam_dim == in_dim
+    init_cam_from_base: bool = True
+    use_chunk_plucker_post_attn: bool = True
+    use_chunk_plucker_input: bool = False
+    chunk_plucker_channels: int = 48  # 8 orig frames × 6D Plücker
+    chunk_plucker_post_attn_blocks: int = 20
+
+    chunk_split_strategy: str = "first_chunk_plus_one"
+    chunk_size: int = 10
+    # Upstream currently forwards chunk metadata through the softmax blocks but
+    # does not apply a chunk-causal mask there. Keep this disabled by default
+    # for checkpoint-output parity; it can be enabled for experiments.
+    use_chunked_softmax_attention: bool = False
+
+    # --- Temporal FFN (GLUMBConvTemp) ---
+    ffn_type: str = "GLUMBConvTemp"
+    t_kernel_size: int = 3
+    mlp_acts: tuple = field(default_factory=lambda: ("silu", "silu", None))
+
+    # --- Position embedding ---
+    pos_embed_type: str = "wan_rope"
+
+    # --- VAE coupling (LTX-2) ---
+    vae_temporal_stride: int = 8  # original-frames per latent frame
+    vae_spatial_stride: int = 32  # pixels per latent token (per spatial axis)
+
+    sample_size: int = 32  # legacy, unused
+    guidance_embeds: bool = False
+    class_dropout_prob: float = 0.0
+
+    # the released checkpoints store raw upstream parameter names; streaming
+    # also keeps an unused all-zero pos_embed while the native model uses RoPE
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            "^pos_embed$": "",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaWMConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaWMArchConfig)
+    prefix: str = "SanaWM"

python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
+
+
+@dataclass
+class SanaWMRefinerArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_transformer_blocks]
+    )
+
+    # Core dims
+    in_channels: int = 128
+    out_channels: int = 128
+    patch_size: int = 1
+    patch_size_t: int = 1
+    num_layers: int = 28
+    num_attention_heads: int = 32
+    attention_head_dim: int = 64
+    cross_attention_dim: int = 4096
+    caption_channels: int = 4096
+
+    qk_norm: bool = True
+    norm_eps: float = 1e-6
+    apply_gated_attention: bool = False
+
+    timestep_scale_multiplier: float = 1000.0
+    rope_type: str = "interleaved"
+
+    # RoPE coord generation
+    sampling_rate: int = 16000
+    hop_length: int = 160
+    scale_factors: tuple = (8, 32, 32)
+    base_num_frames: int = 20
+    base_height: int = 2048
+    base_width: int = 2048
+    causal_offset: int = 1
+
+    # Map Diffusers-style param keys to sglang's LTX-2 primitive naming.
+    # The refiner reuses LTX2Attention / LTX2FeedForward, so it inherits the
+    # same naming differences vs Diffusers:
+    #   * ff.net.0.proj / ff.net.2 -> proj_in / proj_out  (LTX2FeedForward)
+    #   * norm_q / norm_k          -> q_norm / k_norm     (LTX2Attention)
+    # Keep this aligned with LTX2ArchConfig.param_names_mapping in ltx_2.py.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^(transformer_blocks\.\d+\.ff)\.net\.0\.proj\.(.*)$": r"\1.proj_in.\2",
+            r"^(transformer_blocks\.\d+\.ff)\.net\.2\.(.*)$": r"\1.proj_out.\2",
+            r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2",
+            r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaWMRefinerConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaWMRefinerArchConfig)
+    prefix: str = "SanaWMRefiner"

python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+"""StableDiffusion3 Transformer model configuration"""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class StableDiffusion3TransformerArchConfig(DiTArchConfig):
+    """Architecture configuration for StableDiffusion3 Transformer, applicable to SD3-medium, SD3.5-medium, SD3.5-large."""
+
+    sample_size: int = 128
+    patch_size: int = 2
+    in_channels: int = 16
+    out_channels: int = 16
+    num_layers: int = 18
+    attention_head_dim: int = 64
+    num_attention_heads: int = 18
+    cross_attention_dim: int = 4096
+    joint_attention_dim: int = 4096
+    caption_projection_dim: int = 1152
+    pooled_projection_dim: int = 2048
+    pos_embed_max_size: int = 96
+    dual_attention_layers: tuple[int, ...] = ()
+    qk_norm: str | None = None
+
+    _class_name: str = "SD3Transformer2DModel"
+
+
+@dataclass
+class StableDiffusion3TransformerConfig(DiTConfig):
+    """Configuration for StableDiffusion3 Transformer model."""
+
+    arch_config: StableDiffusion3TransformerArchConfig = field(
+        default_factory=StableDiffusion3TransformerArchConfig
+    )

python/sglang/multimodal_gen/configs/models/dits/wanvideo.py

@@ -0,0 +1,125 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_block
+
+
+@dataclass
+class WanVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": r"condition_embedder.image_embedder.ff.fc_in.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": r"condition_embedder.image_embedder.ff.fc_out.\1",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn1\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.local_attn.proj_l.\2",
+            r"^blocks\.(\d+)\.attn2\.norm_added_q\.(.*)$": "",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.proj\.(.*)$": r"patch_embedding.\1",
+            r"^condition_embedder\.text_embedder\.fc_in\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1",
+            r"^condition_embedder\.text_embedder\.fc_out\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1",
+            r"^condition_embedder\.time_embedder\.mlp\.fc_in\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1",
+            r"^condition_embedder\.time_embedder\.mlp\.fc_out\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1",
+            r"^condition_embedder\.time_modulation\.linear\.(.*)$": r"condition_embedder.time_proj.\1",
+            r"^condition_embedder\.image_embedder\.ff\.fc_in\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1",
+            r"^condition_embedder\.image_embedder\.ff\.fc_out\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1",
+            r"^blocks\.(\d+)\.to_q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
+            r"^blocks\.(\d+)\.to_k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
+            r"^blocks\.(\d+)\.to_v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
+            r"^blocks\.(\d+)\.to_out\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
+            r"^blocks\.(\d+)\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2",
+            r"^blocks\.(\d+)\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn1\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_op.local_attn.proj_l.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
+            r"^blocks\.(\d+)\.ffn\.fc_in\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2",
+            r"^blocks\.(\d+)\.ffn\.fc_out\.(.*)$": r"blocks.\1.ffn.net.2.\2",
+            r"^blocks\.(\d+)\.self_attn_residual_norm\.norm\.(.*)$": r"blocks.\1.norm2.\2",
+        }
+    )
+
+    # Some LoRA adapters use the original official layer names instead of hf layer names,
+    # so apply this before the param_names_mapping
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+        }
+    )
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len = 512
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 13824
+    num_layers: int = 40
+    cross_attn_norm: bool = True
+    qk_norm: str = "rms_norm_across_heads"
+    eps: float = 1e-6
+    image_dim: int | None = None
+    added_kv_proj_dim: int | None = None
+    rope_max_seq_len: int = 1024
+    pos_embed_seq_len: int | None = None
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    # Wan MoE
+    boundary_ratio: float | None = None
+
+    # Causal Wan
+    local_attn_size: int = (
+        -1
+    )  # Window size for temporal local attention (-1 indicates global attention)
+    sink_size: int = (
+        0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+    )
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 21
+    attention_type: str = "original"
+    sla_topk: float = 0.1
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class WanVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)
+
+    prefix: str = "Wan"

python/sglang/multimodal_gen/configs/models/dits/zimage.py

@@ -0,0 +1,101 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_zimage_layer
+
+
+@dataclass
+class ZImageArchConfig(DiTArchConfig):
+    all_patch_size: Tuple[int, ...] = (2,)
+    all_f_patch_size: Tuple[int, ...] = (1,)
+    in_channels: int = 16
+    out_channels: int | None = None
+    dim: int = 3840
+    num_layers: int = 30
+    n_refiner_layers: int = 2
+    num_attention_heads: int = 30
+    n_kv_heads: int = 30
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    cap_feat_dim: int = 2560
+    rope_theta: float = 256.0
+    t_scale: float = 1000.0
+    axes_dims: Tuple[int, int, int] = (32, 48, 48)
+    axes_lens: Tuple[int, int, int] = (1024, 512, 512)
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_layer])
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".feed_forward.w13", ".feed_forward.w1", "gate"),
+            (".feed_forward.w13", ".feed_forward.w3", "up"),
+        ]
+    )
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3),
+            r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3),
+            r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3),
+            r"(.*)\.attention\.to_q\.weight_scale_inv$": (
+                r"\1.attention.to_qkv.weight_scale_inv",
+                0,
+                3,
+            ),
+            r"(.*)\.attention\.to_k\.weight_scale_inv$": (
+                r"\1.attention.to_qkv.weight_scale_inv",
+                1,
+                3,
+            ),
+            r"(.*)\.attention\.to_v\.weight_scale_inv$": (
+                r"\1.attention.to_qkv.weight_scale_inv",
+                2,
+                3,
+            ),
+            r"(.*)\.attention\.to_q\.(lora_A|lora_B)$": (
+                r"\1.attention.to_qkv.\2",
+                0,
+                3,
+            ),
+            r"(.*)\.attention\.to_k\.(lora_A|lora_B)$": (
+                r"\1.attention.to_qkv.\2",
+                1,
+                3,
+            ),
+            r"(.*)\.attention\.to_v\.(lora_A|lora_B)$": (
+                r"\1.attention.to_qkv.\2",
+                2,
+                3,
+            ),
+            r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
+            r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
+            r"(.*)\.feed_forward\.w1\.(lora_A|lora_B)$": (
+                r"\1.feed_forward.w13.\2",
+                0,
+                2,
+            ),
+            r"(.*)\.feed_forward\.w3\.(lora_A|lora_B)$": (
+                r"\1.feed_forward.w13.\2",
+                1,
+                2,
+            ),
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.num_channels_latents = self.in_channels
+        self.hidden_size = self.dim
+
+
+@dataclass
+class ZImageDitConfig(DiTConfig):
+    arch_config: ZImageArchConfig = field(default_factory=ZImageArchConfig)
+
+    prefix: str = "zimage"

python/sglang/multimodal_gen/configs/models/encoders/__init__.py

@@ -0,0 +1,45 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    BaseEncoderOutput,
+    EncoderConfig,
+    ImageEncoderConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.encoders.clip import (
+    CLIPTextConfig,
+    CLIPVisionConfig,
+)
+from sglang.multimodal_gen.configs.models.encoders.flux_2 import (
+    FLUX_2_SYSTEM_MESSAGE,
+    Flux2MistralTextConfig,
+    build_flux2_text_messages,
+)
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
+from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
+from sglang.multimodal_gen.configs.models.encoders.ideogram import (
+    Ideogram4TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
+from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
+from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig
+from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
+
+__all__ = [
+    "EncoderConfig",
+    "TextEncoderConfig",
+    "ImageEncoderConfig",
+    "BaseEncoderOutput",
+    "CLIPTextConfig",
+    "CLIPVisionConfig",
+    "FLUX_2_SYSTEM_MESSAGE",
+    "Flux2MistralTextConfig",
+    "build_flux2_text_messages",
+    "LlamaConfig",
+    "Qwen3TextConfig",
+    "Qwen3VLConfig",
+    "T5Config",
+    "Gemma2Config",
+    "Gemma3Config",
+    "Ideogram4TextEncoderConfig",
+]

python/sglang/multimodal_gen/configs/models/encoders/base.py

@@ -0,0 +1,92 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class EncoderArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
+    architectures: list[str] = field(default_factory=lambda: [])
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.SAGE_ATTN_3,
+        }
+    )
+    output_hidden_states: bool = False
+    use_return_dict: bool = True
+
+
+@dataclass
+class TextEncoderArchConfig(EncoderArchConfig):
+    vocab_size: int = 0
+    hidden_size: int = 0
+    num_hidden_layers: int = 0
+    num_attention_heads: int = 0
+    pad_token_id: int = 0
+    eos_token_id: int = 0
+    text_len: int = 0
+    hidden_state_skip_layer: int = 0
+    decoder_start_token_id: int = 0
+    output_past: bool = True
+    scalable_attention: bool = True
+    tie_word_embeddings: bool = False
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names
+    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class ImageEncoderArchConfig(EncoderArchConfig):
+    pass
+
+
+@dataclass
+class BaseEncoderOutput:
+    last_hidden_state: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    attention_mask: torch.Tensor | None = None
+
+
+@dataclass
+class EncoderConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
+
+    prefix: str = ""
+    quant_config: QuantizationConfig | None = None
+    lora_config: Any | None = None
+
+    # Parallel folding: during the encoding stage the whole DiT replica is idle,
+    # so TP-shard the encoder across those otherwise-unused GPUs instead of
+    # running it on a single rank
+    parallel_folding_mode: str | None = None
+
+
+@dataclass
+class TextEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
+
+
+@dataclass
+class ImageEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig)

python/sglang/multimodal_gen/configs/models/encoders/clip.py

@@ -0,0 +1,97 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    ImageEncoderArchConfig,
+    ImageEncoderConfig,
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embeddings,
+    is_layer,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class CLIPTextArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 49408
+    hidden_size: int = 512
+    intermediate_size: int = 2048
+    projection_dim: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 8
+    max_position_embeddings: int = 77
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+    dropout: float = 0.0
+    attention_dropout: float = 0.0
+    initializer_range: float = 0.02
+    initializer_factor: float = 1.0
+    pad_token_id: int = 1
+    bos_token_id: int = 49406
+    eos_token_id: int = 49407
+    text_len: int = 77
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.TORCH_SDPA,  # Force TORCH_SDPA to support attention_mask
+        }
+    )
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embeddings]
+    )
+
+
+@dataclass
+class CLIPVisionArchConfig(ImageEncoderArchConfig):
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    projection_dim: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    num_channels: int = 3
+    image_size: int = 224
+    patch_size: int = 32
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+    dropout: float = 0.0
+    attention_dropout: float = 0.0
+    initializer_range: float = 0.02
+    initializer_factor: float = 1.0
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+    )
+
+
+@dataclass
+class CLIPTextConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=CLIPTextArchConfig)
+
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
+    prefix: str = "clip"
+
+
+@dataclass
+class CLIPVisionConfig(ImageEncoderConfig):
+    arch_config: ImageEncoderArchConfig = field(default_factory=CLIPVisionArchConfig)
+
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
+    prefix: str = "clip"

python/sglang/multimodal_gen/configs/models/encoders/flux_2.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FLUX.2 Mistral text encoder configuration and prompt formatting."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import is_layer
+
+FLUX_2_SYSTEM_MESSAGE = (
+    "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\n"
+    "attribution and actions without speculation."
+)
+
+
+def build_flux2_text_messages(prompts: list[str]) -> list[list[dict]]:
+    cleaned_prompts = [prompt.replace("[IMG]", "") for prompt in prompts]
+    return [
+        [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": FLUX_2_SYSTEM_MESSAGE}],
+            },
+            {"role": "user", "content": [{"type": "text", "text": prompt}]},
+        ]
+        for prompt in cleaned_prompts
+    ]
+
+
+@dataclass
+class Flux2MistralTextArchConfig(TextEncoderArchConfig):
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 512,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Flux2MistralTextConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=Flux2MistralTextArchConfig
+    )
+    prefix: str = "flux_2_mistral"

python/sglang/multimodal_gen/configs/models/encoders/gemma2.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Text encoder configuration for Gemma2 2B, used by SANA for text conditioning.
+#
+# SANA uses the hidden states from Gemma2 (not logits) as the conditioning
+# signal for cross-attention in the DiT. The encoder output dimension (2304)
+# is projected to the DiT's inner_dim via caption_projection.
+#
+# Defaults match google/gemma-2-2b-it (the model used in SANA HF checkpoints).
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class Gemma2ArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 256000
+    hidden_size: int = 2304
+    intermediate_size: int = 9216
+    num_hidden_layers: int = 26
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    hidden_act: str = "gelu_pytorch_tanh"
+    hidden_activation: str = "gelu_pytorch_tanh"
+    max_position_embeddings: int = 8192
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    eos_token_id: int = 1
+    bos_token_id: int = 2
+    tie_word_embeddings: bool = True
+    rope_theta: float = 10000.0
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+
+    # Gemma2 alternates between global and sliding-window attention
+    # on odd/even layers, respectively.
+    sliding_window: int = 4096
+
+    # query_pre_attn_scalar replaces the standard 1/sqrt(head_dim) scaling.
+    query_pre_attn_scalar: int = 256
+
+    # Softcapping bounds raw attention logits via tanh(logits/cap)*cap.
+    # NOTE: SDPA does not natively support softcapping; the runtime model
+    # currently skips this (see Gemma2Attention.forward). Quality impact
+    # is minimal for short text-encoder sequences but should be revisited
+    # for longer context.
+    attn_logit_softcapping: float = 50.0
+    final_logit_softcapping: float = 30.0
+
+    text_len: int = 300
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", "0"),
+            (".gate_up_proj", ".up_proj", "1"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+
+@dataclass
+class Gemma2Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=Gemma2ArchConfig)
+    prefix: str = "gemma_2"

python/sglang/multimodal_gen/configs/models/encoders/gemma_3.py

@@ -0,0 +1,74 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class Gemma3ArchConfig(TextEncoderArchConfig):
+    """Minimal Gemma text-encoder config for tokenizer kwargs.
+
+    Note: runtime will load the actual `text_encoder/` module from the model repo
+    (e.g. Gemma3Model) via transformers; this config mainly controls tokenization.
+    """
+
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "gelu_pytorch_tanh"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = True
+    rope_theta: float = 10000.0
+    rope_scaling: dict | None = None
+    rope_local_base_freq: float = 10000.0
+    sliding_window: int = 4096
+    layer_types: list[str] = field(default_factory=list)
+    query_pre_attn_scalar: int | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 1024
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", "0"),  # type: ignore
+            (".gate_up_proj", ".up_proj", "1"),  # type: ignore
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+
+@dataclass
+class Gemma3Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=Gemma3ArchConfig)
+
+    prefix: str = "gemma_3"

python/sglang/multimodal_gen/configs/models/encoders/ideogram.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.qwen3vl import (
+    Qwen3VLArchConfig,
+    Qwen3VLConfig,
+)
+
+
+@dataclass
+class Ideogram4TextEncoderConfig(Qwen3VLConfig):
+    """Use the local Ideogram text_encoder as a language-only Qwen3-VL encoder."""
+
+    def update_model_arch(self, source_model_dict):
+        super().update_model_arch(source_model_dict)
+        self.post_diffusers_config_update()
+
+    def post_diffusers_config_update(self):
+        self.arch_config.architectures = ["IdeogramQwen3VLTextEncoder"]
+        quant_config = getattr(self.arch_config, "quantization_config", None)
+        if isinstance(quant_config, dict):
+            quant_method = quant_config.get("quant_method")
+            load_in_4bit = quant_config.get("load_in_4bit", False)
+        else:
+            quant_method = getattr(quant_config, "quant_method", None)
+            load_in_4bit = getattr(quant_config, "load_in_4bit", False)
+        quant_method_name = str(quant_method).lower()
+        use_bitsandbytes = "bitsandbytes" in quant_method_name and load_in_4bit
+        self.arch_config.ideogram_bnb_4bit_weight_only = use_bitsandbytes
+        self.arch_config.ideogram_fp8_weight_only = not use_bitsandbytes
+        self.arch_config.requires_gpu_resident_text_encoder = use_bitsandbytes
+
+    def finalize_model_arch(self):
+        self.post_diffusers_config_update()
+
+    arch_config: Qwen3VLArchConfig = field(default_factory=Qwen3VLArchConfig)

python/sglang/multimodal_gen/configs/models/encoders/llama.py

@@ -0,0 +1,62 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class LlamaArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = False
+    rope_theta: float = 10000.0
+    rope_scaling: float | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 256
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
+            (".gate_up_proj", ".up_proj", 1),  # type: ignore
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+
+@dataclass
+class LlamaConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig)
+
+    prefix: str = "llama"

python/sglang/multimodal_gen/configs/models/encoders/mistral3.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Mistral3 text encoder configuration for SGLang diffusion models."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class Mistral3EncoderArchConfig(TextEncoderArchConfig):
+    """Mistral3 text encoder architecture config for ErnieImage.
+
+    Uses Mistral3Model (vision-language model) as text encoder,
+    extracting the second-to-last hidden state layer.
+    """
+
+    vocab_size: int = 131072
+    hidden_size: int = 3072
+    intermediate_size: int = 9216
+    num_hidden_layers: int = 26
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 262144
+    rms_norm_eps: float = 1e-5
+    pad_token_id: int = 11
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    tie_word_embeddings: bool = True
+    head_dim: int = 128
+    hidden_state_skip_layer: int = 2  # Use second-to-last hidden state
+    text_len: int = 0
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+    )
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+    def __post_init__(self):
+        # Let the parent populate tokenizer_kwargs["max_length"] = self.text_len
+        super().__post_init__()
+
+
+@dataclass
+class Mistral3EncoderConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=Mistral3EncoderArchConfig
+    )

python/sglang/multimodal_gen/configs/models/encoders/qwen3.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3 text encoder configuration for SGLang diffusion models."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class Qwen3TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for Qwen3 text encoder.
+
+    Qwen3 is similar to LLaMA but with QK-Norm (RMSNorm on Q and K before attention).
+    """
+
+    vocab_size: int = 151936
+    hidden_size: int = 2560
+    intermediate_size: int = 9728
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 40960
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 151643
+    bos_token_id: int = 151643
+    eos_token_id: int = 151645
+    tie_word_embeddings: bool = True
+    rope_theta: float = 1000000.0
+    rope_scaling: dict | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int = 128
+    text_len: int = 512
+    output_hidden_states: bool = True  # Klein needs hidden states from layers 9, 18, 27
+
+    # Stacked params for weight loading with tensor parallelism
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+    )
+
+    # FSDP sharding conditions for CPU offload
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Qwen3TextConfig(TextEncoderConfig):
+    """Top-level config for Qwen3 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3TextArchConfig)
+    prefix: str = "qwen3"

python/sglang/multimodal_gen/configs/models/encoders/qwen3vl.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class Qwen3VLArchConfig(TextEncoderArchConfig):
+    """Architecture configuration for Qwen3-VL text encoder.
+
+    Qwen3-VL-8B-Instruct is used by JoyImage model.
+    Architecture is similar to Qwen2.5-VL but with Qwen3 improvements.
+    """
+
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = -1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = False
+    rope_theta: float = 10000.0
+    rope_scaling: float | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 2048
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+    # JoyImage specific settings
+    text_token_max_length: int = 2048
+    prompt_template_encode_start_idx = {
+        "image": 34,
+        "video": 91,
+    }
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.tokenizer_kwargs = {
+            "padding": True,
+            "truncation": True,
+            "max_length": self.text_len
+            + self.prompt_template_encode_start_idx["image"],
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Qwen3VLConfig(TextEncoderConfig):
+    """Configuration for Qwen3-VL text encoder.
+
+    Used by JoyImage model.
+    """
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3VLArchConfig)

python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py

@@ -0,0 +1,66 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_embed_tokens,
+    is_final_norm,
+    is_layer,
+)
+
+
+@dataclass
+class QwenImageArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = -1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = False
+    rope_theta: float = 10000.0
+    rope_scaling: float | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 512
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+    vision_token_id: int = 151654
+    image_token_id: int = 151655
+    video_token_id: int = 151656
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
+            (".gate_up_proj", ".up_proj", 1),  # type: ignore
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_layer, is_embed_tokens, is_final_norm]
+    )
+
+
+@dataclass
+class Qwen2_5VLConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=QwenImageArchConfig)
+    # prefix: str = "qwen_image"

python/sglang/multimodal_gen/configs/models/encoders/t5.py

@@ -0,0 +1,86 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import argparse
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.fsdp import (
+    is_final_layer_norm,
+    is_shared,
+    is_t5_block,
+)
+
+
+@dataclass
+class T5ArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32128
+    d_model: int = 512
+    d_kv: int = 64
+    d_ff: int = 2048
+    num_layers: int = 6
+    num_decoder_layers: int | None = None
+    num_heads: int = 8
+    relative_attention_num_buckets: int = 32
+    relative_attention_max_distance: int = 128
+    dropout_rate: float = 0.1
+    layer_norm_epsilon: float = 1e-6
+    initializer_factor: float = 1.0
+    feed_forward_proj: str = "relu"
+    dense_act_fn: str = ""
+    is_gated_act: bool = False
+    is_encoder_decoder: bool = True
+    use_cache: bool = True
+    pad_token_id: int = 0
+    eos_token_id: int = 1
+    classifier_dropout: float = 0.0
+    text_len: int = 512
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q", "q"),
+            (".qkv_proj", ".k", "k"),
+            (".qkv_proj", ".v", "v"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [
+            is_t5_block,
+            is_shared,
+            is_final_layer_norm,
+        ]
+    )
+
+    # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
+    def __post_init__(self):
+        super().__post_init__()
+        act_info = self.feed_forward_proj.split("-")
+        self.dense_act_fn: str = act_info[-1]
+        self.is_gated_act: bool = act_info[0] == "gated"
+        if self.feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"
+
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class T5Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig)
+
+    prefix: str = "t5"
+
+    @staticmethod
+    def add_cli_args(
+        parser: argparse.ArgumentParser, prefix: str = "t5-config"
+    ) -> argparse.ArgumentParser:
+        return parser

python/sglang/multimodal_gen/configs/models/fsdp.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+def is_module_list_entry(name: str, container_name: str) -> bool:
+    # Match only direct block entries, not their inner submodules.
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-2] == container_name and parts[-1].isdigit()
+
+
+def is_module_list_entry_in(name: str, container_names: tuple[str, ...]) -> bool:
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-2] in container_names and parts[-1].isdigit()
+
+
+def is_layer(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "layers")
+
+
+def is_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "blocks")
+
+
+def is_t5_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "block")
+
+
+def is_transformer_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "transformer_blocks")
+
+
+def is_double_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "double_blocks")
+
+
+def is_single_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "single_blocks")
+
+
+def is_refiner_block(name: str, module: object) -> bool:
+    return is_module_list_entry(name, "refiner_blocks")
+
+
+def is_blocks_or_double_blocks(name: str, module: object) -> bool:
+    return is_module_list_entry_in(name, ("blocks", "double_blocks"))
+
+
+def is_blocks_or_transformer_blocks(name: str, module: object) -> bool:
+    return is_module_list_entry_in(name, ("blocks", "transformer_blocks"))
+
+
+def is_zimage_layer(name: str, module: object) -> bool:
+    last_part = name.split(".")[-1]
+    # Preserve Z-Image's finer historical FSDP granularity for perf.
+    return last_part.isdigit() and (
+        "layers" in name or "noise_refiner" in name or "context_refiner" in name
+    )
+
+
+def is_embed_tokens(name: str, module: object) -> bool:
+    return name.endswith("embed_tokens")
+
+
+def is_embeddings(name: str, module: object) -> bool:
+    return name.endswith("embeddings")
+
+
+def is_final_norm(name: str, module: object) -> bool:
+    return name.endswith("norm")
+
+
+def is_shared(name: str, module: object) -> bool:
+    return name.endswith("shared")
+
+
+def is_final_layer_norm(name: str, module: object) -> bool:
+    return name.endswith("final_layer_norm")
+
+
+def is_txt_in(name: str, module: object) -> bool:
+    return name.split(".")[-1] == "txt_in"

python/sglang/multimodal_gen/configs/models/vaes/__init__.py

@@ -0,0 +1,17 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
+from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig
+from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
+    StableDiffusion3VAEConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig
+
+__all__ = [
+    "DacVAEConfig",
+    "HunyuanVAEConfig",
+    "StableDiffusion3VAEConfig",
+    "WanVAEConfig",
+    "Hunyuan3DVAEConfig",
+]

python/sglang/multimodal_gen/configs/models/vaes/base.py

@@ -0,0 +1,234 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import argparse
+import dataclasses
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+from sglang.multimodal_gen.utils import StoreBoolean
+
+AUTO_PARALLEL_DECODE_MODE = "auto"
+SPATIAL_SHARD_PARALLEL_DECODE_MODES = ("spatial_shard", "spatial")
+
+
+@lru_cache(maxsize=8)
+def is_spatial_shard_parallel_decode_mode(mode: str) -> bool:
+    return mode in SPATIAL_SHARD_PARALLEL_DECODE_MODES
+
+
+@lru_cache(maxsize=8)
+def is_auto_parallel_decode_mode(mode: str) -> bool:
+    return mode == AUTO_PARALLEL_DECODE_MODE
+
+
+@lru_cache(maxsize=128)
+def _should_use_auto_spatial_shard_parallel_decode(
+    z_shape: tuple[int, ...],
+    world_size: int,
+    min_latent_elements_per_rank: int,
+) -> bool:
+    if world_size <= 1 or z_shape[-2] < world_size:
+        return False
+    latent_elements_per_rank = (
+        z_shape[0] * z_shape[-3] * z_shape[-2] * z_shape[-1]
+    ) // world_size
+    return latent_elements_per_rank >= min_latent_elements_per_rank
+
+
+def should_use_spatial_shard_parallel_decode(
+    config: Any, z: torch.Tensor | None = None, world_size: int = 1
+) -> bool:
+    if not config.use_parallel_decode:
+        return False
+
+    if is_spatial_shard_parallel_decode_mode(config.parallel_decode_mode):
+        return True
+
+    if not is_auto_parallel_decode_mode(config.parallel_decode_mode):
+        return False
+
+    if not config.auto_parallel_decode_prefers_spatial_shard():
+        return False
+
+    if z is None:
+        return True
+
+    return config.should_use_auto_spatial_shard_parallel_decode(z, world_size)
+
+
+@dataclass
+class VAEArchConfig(ArchConfig):
+    scaling_factor: float | torch.Tensor = 0
+
+    temporal_compression_ratio: int = 4
+    # or vae_scale_factor?
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class VAEConfig(ModelConfig):
+    arch_config: VAEArchConfig = field(default_factory=VAEArchConfig)
+
+    # sglang-diffusion VAE-specific parameters
+    load_encoder: bool = True
+    load_decoder: bool = True
+
+    tile_sample_min_height: int = 256
+    tile_sample_min_width: int = 256
+    tile_sample_min_num_frames: int = 16
+    tile_sample_stride_height: int = 192
+    tile_sample_stride_width: int = 192
+    tile_sample_stride_num_frames: int = 12
+    blend_num_frames: int = 0
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = True
+    use_parallel_tiling: bool = True
+    use_temporal_scaling_frames: bool = True
+    use_parallel_decode: bool = True
+    parallel_decode_mode: str = AUTO_PARALLEL_DECODE_MODE
+    auto_parallel_decode_min_latent_elements_per_rank: int = 4096
+
+    def __post_init__(self):
+        self.blend_num_frames = (
+            self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+        )
+
+    def post_init(self):
+        pass
+
+    def auto_parallel_decode_prefers_spatial_shard(self) -> bool:
+        return False
+
+    def should_use_auto_spatial_shard_parallel_decode(
+        self, z: torch.Tensor, world_size: int
+    ) -> bool:
+        return _should_use_auto_spatial_shard_parallel_decode(
+            tuple(z.shape),
+            world_size,
+            self.auto_parallel_decode_min_latent_elements_per_rank,
+        )
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any:
+        """Add CLI arguments for VAEConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.load-encoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_encoder",
+            default=None,
+            help="Whether to load the VAE encoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.load-decoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_decoder",
+            default=None,
+            help="Whether to load the VAE decoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_height",
+            default=None,
+            help="Minimum height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_width",
+            default=None,
+            help="Minimum width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames",
+            default=None,
+            help="Minimum number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height",
+            default=None,
+            help="Stride height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width",
+            default=None,
+            help="Stride width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames",
+            default=None,
+            help="Stride number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.blend-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.blend_num_frames",
+            default=None,
+            help="Number of frames to blend for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_tiling",
+            default=None,
+            help="Whether to use tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-temporal-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_temporal_tiling",
+            default=None,
+            help="Whether to use temporal tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-parallel-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_parallel_tiling",
+            default=None,
+            help="Whether to use parallel tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-parallel-decode",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_parallel_decode",
+            default=None,
+            help="Whether to use parallel decode for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.parallel-decode-mode",
+            choices=("tiled", "patch", "spatial_shard", "spatial", "auto"),
+            dest=f"{prefix.replace('-', '_')}.parallel_decode_mode",
+            default=None,
+            help="Parallel decode mode for VAE",
+        )
+
+        return parser
+
+    def get_vae_scale_factor(self):
+        return 2 ** (len(self.arch_config.block_out_channels) - 1)
+
+    def encode_sample_mode(self):
+        return "argmax"
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig":
+        kwargs = {}
+        for attr in dataclasses.fields(cls):
+            value = getattr(args, attr.name, None)
+            if value is not None:
+                kwargs[attr.name] = value
+        return cls(**kwargs)