Support New Diffusion Models - SGLang Documentation

Use this guide as a triage flow for finding the smallest change that can support a model. Most new model work should touch a small set of files, even though the runtime is split into separate folders.

Read the Code in This Order

The files are split by runtime responsibility. For a new model, read the request path first:

configs/pipeline_configs/{model}.py defines model-specific denoising and decoding behavior.
runtime/pipelines/{model}.py wires modules into stages.
runtime/pipelines_core/stages/ runs the shared stage logic.
runtime/models/ contains native model components only when the architecture cannot be reused.

That is the dependency direction. Avoid making a model PR that requires readers to jump between folders in a different order. runtime/models/ owns modeling code: checkpoint-defined neural modules, architecture wrappers, and weight-loading or forward-path details that are intrinsic to one model family. Reusable serving infrastructure belongs in SGLang-Diffusion runtime folders such as runtime/cache/, runtime/distributed/, runtime/utils/, or shared pipeline stages. This includes cache managers, graph runners, process-group transport, request utilities, and common action-policy helpers. Model packages may call these helpers. Keep ownership in shared runtime folders unless the code is truly architecture-specific.

Place helpers with their owners

Use the narrowest existing owner before adding a utility module:

Scope	Location
Used by one file, or specific to one operation	A private helper in the consuming file; inline short one-off expressions
Shared within a domain	A purpose-named module in that domain, such as `runtime/realtime/video.py` or `runtime/layers/attention/mask_strategy.py`
Shared across domains, without model or pipeline semantics	An existing focused module under `runtime/utils/`, such as `argparse.py`, `process.py`, or `precision.py`

Do not create a top-level utils/ package or grow a catch-all utils.py or common.py. Split large mixed-responsibility files along ownership boundaries, not arbitrary line counts. A helper folder is warranted only when several cohesive modules need it, not for a single function or hypothetical reuse. Model code must not import pipeline stages. Put contracts shared by models and stages in a lower-level domain module; for example, realtime cache keys belong under runtime/realtime/. Keep GPU initialization, monkey patches, and model loading out of generic utility imports. When moving internal helpers, update all callers, tests, and cookbook examples together. Preserve documented registration and serving entry points; do not add re-export chains just to retain obsolete internal utility paths.

Out-of-Tree Models and Pipelines

An installed package can register native component models and a pipeline without modifying SGLang-Diffusion. Register them in the package’s __init__.py:

from sglang.multimodal_gen.registry import register_pipeline
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry

from .configs import CustomPipelineConfig, CustomSamplingParams
from .pipeline import CustomPipeline


ModelRegistry.register_model(
    "CustomTransformer2DModel",
    "custom_diffusion.models:CustomTransformer2DModel",
)
register_pipeline(
    CustomPipeline,
    sampling_param_cls=CustomSamplingParams,
    pipeline_config_cls=CustomPipelineConfig,
    hf_model_paths=["my-org/custom-diffusion-model"],
    model_detectors=[lambda value: "custom-diffusion" in value.lower()],
)

Install the package in the server environment and set the same environment variable used by SRT plugins:

pip install -e /path/to/custom-diffusion
SGLANG_EXTERNAL_MODEL_PACKAGE=custom_diffusion \
  sglang serve --model-path my-org/custom-diffusion-model

Notes:

The string form of register_model keeps component imports lazy.
hf_model_paths also supports checkpoints without model_index.json. Other Diffusers checkpoints can select the pipeline through _class_name.
For a standalone safetensors file, pass --pipeline CustomPipeline.
Set the environment variable before startup. Each process imports the package once. Use overwrite=True only to intentionally replace a built-in pipeline.

Start With the Smallest Change

Before adding files, decide which path fits the model.

Situation	What to do
A new checkpoint uses an existing native family	Add the Hugging Face path and, if needed, a small `SamplingParams` or `PipelineConfig` variant. Reuse the existing pipeline and modules.
The model has a new native DiT/UNet architecture	Add a native SGLang pipeline and the missing model components. Keep denoising and decoding on the shared stages unless measured behavior requires model-specific logic.
The model is long-tail or you only need compatibility first	Prefer the Diffusers backend for compatibility-first support. Add native support later if performance or deployment needs justify it.

Do not add a folder just to mirror the Diffusers repository layout. Add a new file only when an existing pipeline, stage, module, config, or sampler cannot express the behavior clearly.

Minimal File Map

The source tree is split by runtime responsibility. That split is useful for optimization. Keep new model PRs focused on the files required by model behavior.

Area	Add or edit when	Typical file
Runtime parameters	The request schema differs from existing models	`configs/sample/{model}.py`
Pipeline config	Denoising, decoding, precision, position encoding, or CFG hooks differ	`configs/pipeline_configs/{model}.py`
Pipeline wiring	The model needs a new stage layout or module list	`runtime/pipelines/{model}.py`
DiT/UNet module	The denoising network is new	`runtime/models/dits/{model}.py`
dVLA or policy module	The checkpoint defines a new action-policy architecture	`runtime/models/{family}/modeling_*.py` or a task-specific model subfolder
Shared runtime infrastructure	Cache, CUDA graph, distributed transfer, request utilities, or action-policy helpers can be reused by future models	`runtime/cache/`, `runtime/distributed/`, `runtime/utils/`, `runtime/pipelines_core/stages/`
Model component config	A model component has static architecture config	`configs/models/dits/{model}.py`, `configs/models/vaes/{model}.py`
Model-specific stage	A single stage’s runtime semantics cannot be expressed by a native stage or a narrow subclass of one	`runtime/pipelines_core/stages/model_specific_stages/{model}.py`
Encoder, VAE, scheduler	No existing implementation can be reused	`runtime/models/encoders/`, `runtime/models/vaes/`, `runtime/models/schedulers/`

For a new native architecture, the common minimum is:

configs/sample/{model}.py
configs/pipeline_configs/{model}.py
runtime/pipelines/{model}.py
runtime/models/dits/{model}.py

Every extra file should map to model behavior that existing code cannot express clearly. For dVLA or other non-image diffusion policies, keep the same ownership rule. The policy network, VLM/action expert modules, checkpoint mapping, and model-specific forward code belong under runtime/models/. Prefix caches, request-local contexts, denoising graph runners, OpenPI-compatible transport, and prefix/action process-group utilities should be shared SGLang-Diffusion runtime infrastructure when they are useful beyond the first model.

Read the Reference First

Use the model’s Diffusers pipeline, official implementation, or model_index.json as the source of truth. Write down:

Which modules must be loaded: tokenizer, text encoder, image encoder, transformer, scheduler, VAE, processor, and any extra adapters.
The prompt and image encoding flow.
Latent shape, packing, scale, shift, dtype, and device rules.
Timestep and sigma schedule.
The exact forward() kwargs expected by the denoising network.
VAE decode rules and output post-processing.

If the new model is close to Flux, Qwen-Image, GLM-Image, Wan, HunyuanVideo, or LTX, extend that implementation before starting from an empty file.

Choose a Pipeline Shape

SGLang-Diffusion uses ComposedPipelineBase to wire stages together. Most native pipelines should choose the least invasive stage shape that preserves the runtime semantics.

Shape	Use when	Layout
Native stages	Text/image encoding, latent prep, timestep prep, denoising, and decoding match existing helpers	`add_standard_t2i_stages()`, `add_standard_ti2i_stages()`, or a similar helper
Native-stage subclass	One stage has model-specific details, but the stage boundary and batch contract still match an existing native stage	`{Model}TextEncodingStage(TextEncodingStage) -> LatentPreparationStage -> TimestepPreparationStage -> DenoisingStage -> DecodingStage`
Custom single-purpose stage	One step has a different state owner or batch-field lifecycle and cannot cleanly inherit from a native stage	Native stages with one `{Model}{Purpose}Stage` inserted or substituted
Aggregated custom stage	Several preparation steps are inseparable in the reference pipeline and cannot be split without fragile duplicate state	`{Model}BeforeDenoisingStage -> DenoisingStage -> DecodingStage`

Prefer this order:

Use native stages directly. This keeps the model on shared code paths for offload, component readiness, profiling, disaggregation, batching, and future stage-level optimizations.
Subclass the narrowest native stage. If only prompt processing differs, inherit from TextEncodingStage. If only latent setup, timestep setup, denoising, or decode differs, inherit from that specific native stage. Preserve the existing input/output fields whenever possible.
Add a custom single-purpose stage only when no native stage contract fits. Keep the stage owner narrow: one stage should own one coherent transformation, such as a custom condition assembly step or a model-specific policy/action bridge.
Use an aggregated BeforeDenoisingStage only as a last resort. This is the least preferred shape because it hides multiple runtime responsibilities in one stage, increases code size and review cost, and bypasses shared hooks for offload, profiling, disaggregation, batching, and future stage-level optimizations.

Implement the Pieces

1. Sampling Params

Create request parameters only for values users can set at runtime.

# python/sglang/multimodal_gen/configs/sample/my_model.py
from dataclasses import dataclass

from sglang.multimodal_gen.configs.sample.sampling_params import ImageSamplingParams


@dataclass
class MyModelSamplingParams(ImageSamplingParams):
    guidance_scale: float = 4.0
    num_inference_steps: int = 28

2. Pipeline Config

PipelineConfig is where shared denoising and decoding stages get model-specific callbacks.

# python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py
from dataclasses import dataclass, field


@dataclass
class MyModelPipelineConfig(ImagePipelineConfig):
    task_type: ModelTaskType = ModelTaskType.T2I
    should_use_guidance: bool = True
    dit_config: DiTConfig = field(default_factory=MyModelDiTConfig)
    vae_config: VAEConfig = field(default_factory=MyModelVAEConfig)

    def prepare_pos_cond_kwargs(self, batch, latent_model_input, t, **kwargs):
        return {
            "hidden_states": latent_model_input,
            "encoder_hidden_states": batch.prompt_embeds[0],
            "timestep": t,
        }

    def prepare_neg_cond_kwargs(self, batch, latent_model_input, t, **kwargs):
        return {
            "hidden_states": latent_model_input,
            "encoder_hidden_states": batch.negative_prompt_embeds[0],
            "timestep": t,
        }

Make these kwargs match the denoising module’s forward() signature exactly.

3. Pipeline Wiring

Use the standard helper when the model fits it.

# python/sglang/multimodal_gen/runtime/pipelines/my_model.py
class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "MyModelPipeline"

    _required_config_modules = [
        "text_encoder",
        "tokenizer",
        "transformer",
        "scheduler",
        "vae",
    ]

    def create_pipeline_stages(self, server_args: ServerArgs):
        self.add_standard_t2i_stages()


EntryClass = [MyModelPipeline]

When a standard helper is not enough, first check whether only one native stage needs model-specific behavior. In that case, subclass that stage and keep the rest of the pipeline standard. For example, custom tokenization or prompt-window logic should usually inherit from TextEncodingStage directly.

class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "MyModelPipeline"

    _required_config_modules = [
        "text_encoder",
        "tokenizer",
        "transformer",
        "scheduler",
        "vae",
    ]

    def create_pipeline_stages(self, server_args: ServerArgs):
        self.add_stage(InputValidationStage())
        self.add_stage_factory(
            RoleType.ENCODER,
            lambda: MyModelTextEncodingStage(
                text_encoder=self.get_module("text_encoder"),
                tokenizer=self.get_module("tokenizer"),
            ),
            "my_model_text_encoding_stage",
        )
        self.add_standard_latent_preparation_stage()
        self.add_standard_timestep_preparation_stage()
        self.add_standard_denoising_stage()
        self.add_standard_decoding_stage()


EntryClass = [MyModelPipeline]

Use a custom single-purpose stage only when the reference pipeline has one step that cannot be represented cleanly by a hook or native-stage subclass. Keep the custom stage narrow and reuse native stages before and after it.

class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "MyModelPipeline"

    _required_config_modules = [
        "text_encoder",
        "tokenizer",
        "transformer",
        "scheduler",
        "vae",
    ]

    def create_pipeline_stages(self, server_args: ServerArgs):
        self.add_stage(InputValidationStage())
        self.add_standard_text_encoding_stage()
        self.add_stage_factory(
            RoleType.ENCODER,
            lambda: MyModelConditioningStage(
                scheduler=self.get_module("scheduler"),
                vae=self.get_module("vae"),
            ),
            "my_model_conditioning_stage",
        )
        self.add_standard_latent_preparation_stage()
        self.add_standard_timestep_preparation_stage()
        self.add_standard_denoising_stage()
        self.add_standard_decoding_stage()


EntryClass = [MyModelPipeline]

Use an aggregated BeforeDenoisingStage only when the reference pipeline couples several preparation steps so tightly that splitting them would require fragile duplicate state or extra synchronization. Do not start with this shape.

class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "MyModelPipeline"

    _required_config_modules = [
        "text_encoder",
        "tokenizer",
        "transformer",
        "scheduler",
        "vae",
    ]

    def create_pipeline_stages(self, server_args: ServerArgs):
        self.add_stage(InputValidationStage())
        self.add_stage(
            MyModelBeforeDenoisingStage(
                text_encoder=self.get_module("text_encoder"),
                tokenizer=self.get_module("tokenizer"),
                scheduler=self.get_module("scheduler"),
                vae=self.get_module("vae"),
            )
        )
        self.add_standard_denoising_stage()
        self.add_standard_decoding_stage()


EntryClass = [MyModelPipeline]

Reuse component loaders

Most components should keep the default loader for their role (transformer, text encoder, VAE, scheduler, or tokenizer). For an auxiliary module that accepts model_cls(**config) and loads an unchanged state dict, select PlainStateDictComponentLoader instead of adding a model-specific loader class:

from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
    PlainStateDictComponentLoader,
)


class MyModelPipeline(ComposedPipelineBase):
    _required_config_modules = ["transformer", "queryformer", "text_projection"]
    component_loaders = {
        "queryformer": PlainStateDictComponentLoader,
        "text_projection": PlainStateDictComponentLoader,
    }

    # Wire the stages as in the examples above.

The mapping uses exact pipeline component names, including aliases declared in _extra_config_module_map. It is local to this pipeline; omitted components keep their existing dispatch. An explicitly selected loader fails on errors instead of falling back to a different implementation. Register each module through EntryClass in its model file, or ModelRegistry.register_model for an out-of-tree package. The shared loader:

Resolves the registered class from config.json’s _class_name, falling back to the architecture in model_index.json, and passes non-metadata config fields to its constructor.
Loads a single safetensors file or an indexed sharded checkpoint using the shared checkpoint selector, with strict key and shape checks.
Honors exact overrides such as --component-weights-paths.queryformer PATH and --component-precisions.queryformer bf16. Precision defaults to the pipeline’s dit_precision.
Leaves eval mode and final CPU/GPU placement to the shared loading lifecycle. It does not implement TP/FSDP weight sharding, quantization, or direct GPU loading. Explicit quantization and direct-GPU-loading overrides are rejected.

Keep a specialized loader when the component needs a config-object constructor, weight remapping or fusion, sharding, or a non-standard checkpoint format. Such a loader can also be selected through component_loaders; registering a model alone does not opt it into the plain state-dict protocol.

4. Last-Resort Before-Denoising Stage

A BeforeDenoisingStage is not a catch-all replacement for the native stages. Use it when the model has custom latent packing, conditioning assembly, timestep preparation, or request-local state that does not fit LatentPreparationStage or TimestepPreparationStage, and only after checking whether the work can be a native-stage subclass or a custom single-purpose stage. If the difference is prompt handling, subclass TextEncodingStage instead. A proper BeforeDenoisingStage should populate the batch fields consumed by DenoisingStage.

class MyModelBeforeDenoisingStage(PipelineStage):
    @torch.no_grad()
    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
        prompt_embeds, negative_prompt_embeds = self._encode_prompt(batch)
        latents = self._prepare_latents(batch)
        timesteps, sigmas = self._prepare_timesteps(batch)

        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.tolist()
        batch.raw_latent_shape = latents.shape
        return batch

Required fields for DenoisingStage:

Field	Notes
`batch.latents`	Initial latent tensor, including any packing required by the model.
`batch.timesteps`	Timestep tensor in the exact order used by the reference pipeline.
`batch.sigmas`	Python list when the scheduler expects sigma values.
`batch.prompt_embeds`	Positive embeddings, wrapped in a list.
`batch.negative_prompt_embeds`	Negative embeddings, wrapped in a list when CFG is used.
`batch.num_inference_steps`	Number of denoising iterations.
`batch.raw_latent_shape`	Original latent shape before packing, if decode needs it.

5. Distributed and memory integration

Single-GPU parity is only the first milestone. Complete native support also requires:

Encoder and DiT TP/SP: use native parallel projections and sharded weight loading for TP, and USPAttention for SP. Handle masks, RoPE, padding, and output gathering without falling back to a replicated full model. TP and SP must work together.
VAE parallel decode: subclass ParallelTiledVAE, or reuse an existing native base with the same contract. Support tiled and spatial_shard decode through DecodingStage and the shared decode group. Reuse runtime/layers/parallel_conv.py and runtime/models/vaes/parallel/diffusers_spatial.py where applicable.
Layerwise offload: every loaded neural module must inherit LayerwiseOffloadableModuleMixin and list all repeated block paths in layer_names. Set layerwise_offload_dit_group_enabled = False for non-DiT modules. Component CPU offload is not a substitute.

See wanvideo.py and qwen_image.py for DiT TP/SP, gemma_3.py for encoder TP and offload, and autoencoder_kl_qwenimage.py or ltx_2_vae.py for VAE decode. The Diffusers backend is compatibility-first and does not need to meet this native integration contract.

6. Registry

Define a register() function in configs/pipeline_configs/{model}.py. The runtime auto-discovers it on startup and calls it to register the sampling params and pipeline config.

# python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py
def register():
    from sglang.multimodal_gen.registry import register_configs

    register_configs(
        sampling_param_cls=MyModelSamplingParams,
        pipeline_config_cls=MyModelPipelineConfig,
        hf_model_paths=["org/my-model"],
        model_detectors=[lambda hf_id: "my-model" in hf_id.lower()],
    )

model_detectors matches a model path or model_index.json _class_name when the Hugging Face path varies; see wan.py or qwen_image21.py for real examples. The pipeline file is discovered through its EntryClass; do not add a second pipeline registry unless the existing registry requires it.

Verify the Port

Use one deterministic prompt and seed while comparing with the reference implementation.

Run a single-GPU smoke test and check that the output contains coherent content.
Compare latent scale and shift, timestep order, sigma values, and conditioning kwargs against Diffusers or the official implementation.
Compare VAE decode separately, including tiled and multi-GPU spatial_shard.
Run encoder and DiT TP, SP, combined TP x SP, and --layerwise-offload-components all; compare with the single-GPU resident baseline.
If the model supports LoRA, CFG parallelism, or disaggregation, test each feature explicitly.
Add or update the cookbook, examples, and Supported Models catalog when users need a new launch command.

Common failure points:

Wrong latent scale or shift.
Reversed or dtype-mismatched timesteps.
Missing negative embeddings when CFG is enabled.
Conditioning kwarg names mismatched with the DiT forward().
Rotary embedding shape or style mismatch.
Decoding packed latents without restoring raw_latent_shape.

PR Checklist

Reused an existing family, stage, module, scheduler, or VAE wherever possible.
Kept the new-model touch surface small and justified any extra files.
Added SamplingParams, PipelineConfig, pipeline wiring, DiT module, and registry entry when native support is needed.
Confirmed pipeline_name matches the Diffusers model_index.json _class_name when applicable.
Confirmed _required_config_modules matches the model repo.
Verified image or video quality against a reference output.
Completed the distributed and memory integration checks above.
Tested CFG parallelism and distributed serving paths when they apply.

Declare multiple task types

Keep PipelineConfig.task_type as the default task. To let one configured pipeline serve several tasks, declare its capabilities as a class variable:

from dataclasses import dataclass
from typing import ClassVar

from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.configs.task_type import ModelTaskType


@dataclass
class MultiTaskPipelineConfig(PipelineConfig):
    task_type: ModelTaskType = ModelTaskType.T2V
    supported_task_types: ClassVar[tuple[ModelTaskType, ...]] = (
        ModelTaskType.T2I,
        ModelTaskType.I2I,
        ModelTaskType.T2V,
        ModelTaskType.I2V,
    )

The declaration must be nonempty, contain no duplicates, and include the configured default. Existing pipelines that omit it retain the singleton (task_type,). Override get_supported_task_types() if capabilities depend on the loaded checkpoint or configuration, preserving these invariants. Flat CLI/Python task_type arguments select the request. To change the configured default, use pipeline_config={"task_type": ModelTaskType.I2V} (or its serialized configuration), choosing a supported task. Capabilities describe implemented execution paths; declaring a task does not load additional components or implement conditioning and decoding for it. Requests carry SamplingParams.task_type, accepting an enum or its name (case-insensitive). Admission resolves it to an enum and derives the output data_type. Explicit tasks must belong to the supported set and match the HTTP endpoint’s output type. Without an explicit task, admission filters by output type and conditioning inputs, prefers a compatible default, then selects a unique compatible alternative. Python/CLI requests also prefer the default’s output type among alternatives. Ambiguous requests must specify task_type. No request changes the shared configuration. Use the admitted batch.task_type in model stages when behavior depends on the request. Shared stages use get_request_task_type(batch, pipeline_config) to preserve compatibility with direct callers that have no resolved task. Audit model-specific frame adjustment, reference loading, latent construction, and decoding for assumptions about the server default. Declare the union of components used by all supported tasks in stage residency metadata. Synthetic warmup uses the default task; provide the pipeline’s warmup hook when that task needs model-specific reference inputs. Dynamic batching compares the admitted task along with the other sampling fields. I2I and I2V require image input; T2I and T2V reject it. Existing TI2I and TI2V retain optional image conditioning. V2V and F2V require video input; F2V denotes video-prefix conditioning. Models own the details of prefix length and temporal alignment. Pipelines with explicit capabilities reject video input for tasks that do not accept it; legacy singleton pipelines retain their model-owned video checks. See request task selection for the HTTP interface. Test task selection, invalid inputs, discovery, serialization, batching, and alternating tasks on one running server.

​Read the Code in This Order

​Place helpers with their owners

​Out-of-Tree Models and Pipelines

​Start With the Smallest Change

​Minimal File Map

​Read the Reference First

​Choose a Pipeline Shape

​Implement the Pieces

​1. Sampling Params

​2. Pipeline Config

​3. Pipeline Wiring

​Reuse component loaders

​4. Last-Resort Before-Denoising Stage

​5. Distributed and memory integration

​6. Registry

​Verify the Port

​PR Checklist

​Declare multiple task types