Sglang Diffusion Add Model

Add a Diffusion Model to SGLang

Use this skill when adding a new diffusion model or pipeline variant to sglang.multimodal_gen.

Three 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.

Style C: Native Task-Contract Pipeline

Use this only when one checkpoint exposes multiple tightly coupled modalities or request profiles that cannot be represented safely by generic image/video sampling fields. MiniMax-H3 is the reference: it selects FL2VA or Ref2VA weights from one root model ID, validates canonical task / conditions / target requests before queueing, packs text/video/audio tokens into one denoise sequence, and returns synchronized video plus audio.

This style still uses ComposedPipelineBase, but owns a model-specific chain under stages/model_specific_stages/<model>/. Keep request validation, media materialization, packed-sequence construction, per-modality encode/decode, and presentation as explicit stages. Do not force coupled state into the standard DenoisingStage / DecodingStage contract just to resemble a simpler model.

Choose this style only with source evidence that the public API, scheduler, or joint latent state needs it. Preserve one canonical request object from API admission through offline generation and server execution so the two entry points cannot silently diverge.

How to Choose

Key principle (standard-denoise styles): For Hybrid and Modular pipelines, the stage(s) before DenoisingStage must produce a Req batch object with all the standard tensor fields that DenoisingStage expects (latents, timesteps, prompt embeds, and model-specific conditioning). Native task-contract pipelines may own a different denoise/decode contract; keep that divergence explicit and covered by request-contract tests.

Key Files and Directories

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 MiniMax-H3, Krea-2, LTX-2/2.3/2.5, HunyuanVideo/FastHunyuan, Wan/FastWan/TurboWan/LingBot World/LingBot Video MoE, MOVA, FLUX/FLUX.2/Klein, LongCat-Image, Z-Image, Qwen-Image/edit/layered, GLM-Image, SD3, Hunyuan3D, Helios, Cosmos3 Nano/Super/Edge/distilled, SANA/SANA-Video/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/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:

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/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):

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):

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/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/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):

Step 7: Define the Pipeline Class

The pipeline class is minimal -- it just wires the stages together.

# 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:

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

Step 10: Decide the ComfyUI Route (Optional)

A model is reachable from ComfyUI two ways. Pick one deliberately — the wrong choice costs several hundred lines of weight-mapping code that buys nothing.

Server route. ComfyUI sends an HTTP request and SGLang runs the whole pipeline. Choose this when the model needs conditioning ComfyUI cannot supply (audio, reference materials, task routing), produces more than one modality, or has its own request contract.

Cost: nothing, if the request fits the existing generate_image / generate_video fields. If the model has extra request fields, pass them through extra_fields — the request schemas accept unknown keys, so the client in apps/ComfyUI_SGLDiffusion/core/server_api.py does not need a per-model change. Add a node in nodes.py only when the inputs are worth surfacing as ComfyUI widgets. SGLDiffusionGenerateH3 is the worked example.

Executor route. ComfyUI's KSampler drives the denoise loop and SGLang replaces the DiT forward, using ComfyUI's own text encoders and VAE. Choose this only when the model denoises a single latent tensor that ComfyUI already knows how to build and decode.

Cost, per model: a runtime/pipelines/comfyui_<model>_pipeline.py that maps ComfyUI's single-file checkpoint layout onto the native module tree (350-690 lines in the existing three), an executor in apps/ComfyUI_SGLDiffusion/executors/ that adapts latent layout and conditioning to Req, and entries in both dicts in core/generator.py.

The deciding question is not model size or modality — it is whether ComfyUI's sampler can drive the model's loop unchanged. If reproducing the conditioning inside ComfyUI would duplicate stages the server already runs, take the server route.