minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

docs: simplify diffusion new model guide (#30109)

Browse Source
This commit is contained in:
Mick
2026-07-17 19:39:38 +08:00
committed by GitHub
parent 681c223570
commit 85ac56c823
1 changed files with 239 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/sglang-diffusion/support_new_models.mdx
+239 -510
View File
@@ -1,285 +1,156 @@
---
title: "How to Support New Diffusion Models"
metatags:
description: "This document explains how to add support for new diffusion models in SGLang Diffusion."
title: "Support New Diffusion Models"
description: "A concise implementation guide for adding diffusion model families to SGLang-Diffusion."
---
This document explains how to add support for new diffusion models in SGLang Diffusion.
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.
## Architecture Overview
## Read the Code in This Order
SGLang Diffusion is engineered for both performance and flexibility, built upon a pipeline architecture. This
design allows developers to construct pipelines for various diffusion models while keeping the core generation
loop standardized for optimization.
The files are split by runtime responsibility. For a new model, read the
request path first:
At its core, the architecture revolves around two key concepts, as highlighted in our [blog post](https://lmsys.org/blog/2025-11-07-sglang-diffusion/#architecture):
1. `registry.py` chooses the model family, sampling params, and pipeline config.
2. `configs/pipeline_configs/{model}.py` defines model-specific denoising and
decoding behavior.
3. `runtime/pipelines/{model}.py` wires modules into stages.
4. `runtime/pipelines_core/stages/` runs the shared stage logic.
5. `runtime/models/` contains native model components only when the architecture
cannot be reused.
- **`ComposedPipeline`**: This class orchestrates a series of `PipelineStage`s to define the complete generation process for a specific model. It acts as the main entry point for a model and manages the data flow between the different stages of the diffusion process.
- **`PipelineStage`**: Each stage is a modular component that encapsulates a function within the diffusion process. Examples include prompt encoding, the denoising loop, or VAE decoding.
That is the dependency direction. Avoid making a model PR that requires readers
to jump between folders in a different order.
### Two Pipeline Styles
`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.
SGLang Diffusion supports two pipeline composition styles. Both are valid; choose the one that best fits your model.
## Start With the Smallest Change
#### Style A: Hybrid Monolithic Pipeline (Recommended Default)
Before adding files, decide which path fits the model.
The recommended default for most new models. Uses a three-stage structure:
| 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. |
```
BeforeDenoisingStage (model-specific) → DenoisingStage (standard) → DecodingStage (standard)
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 |
| --- | --- | --- |
| Registry | Always, unless extending an already registered family | `python/sglang/multimodal_gen/registry.py` |
| 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 | Pre-processing is too custom for the standard stages | `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:
1. `registry.py`
2. `configs/sample/{model}.py`
3. `configs/pipeline_configs/{model}.py`
4. `runtime/pipelines/{model}.py`
5. `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 use one of these two shapes.
| Shape | Use when | Layout |
| --- | --- | --- |
| Standard stages | Text/image encoding, latent prep, denoising, and decoding match existing helpers | `add_standard_t2i_stages()`, `add_standard_ti2i_stages()`, or a similar helper |
| Model-specific pre-processing | The reference pipeline has custom captioning, image conditioning, latent packing, or timestep preparation | `{Model}BeforeDenoisingStage -> DenoisingStage -> DecodingStage` |
Prefer standard stages when possible. Use a model-specific
`BeforeDenoisingStage` when trying to force the model into shared stages would
create many conditionals.
## Implement the Pieces
### 1. Sampling Params
Create request parameters only for values users can set at runtime.
```python
# 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
```
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
</colgroup>
<thead>
<tr>
<th>Stage</th>
<th>Ownership</th>
<th>Responsibility</th>
</tr>
</thead>
<tbody>
<tr>
<td><code>&#123;Model&#125;BeforeDenoisingStage</code></td>
<td>Model-specific</td>
<td>All pre-processing: input validation, text/image encoding, latent preparation, timestep computation</td>
</tr>
<tr>
<td><code>DenoisingStage</code></td>
<td>Framework-standard</td>
<td>The denoising loop (DiT/UNet forward passes), shared across all models</td>
</tr>
<tr>
<td><code>DecodingStage</code></td>
<td>Framework-standard</td>
<td>VAE decoding from latent space to pixel space, shared across all models</td>
</tr>
</tbody>
</table>
### 2. Pipeline Config
**Why recommended?** Modern diffusion models often 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. Convenience methods like `add_standard_t2i_stages()` and `add_standard_ti2i_stages()` make this very concise.
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.
- **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.
#### How to Choose
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "50%"}} />
<col style={{width: "50%"}} />
</colgroup>
<thead>
<tr>
<th>Situation</th>
<th>Recommended Style</th>
</tr>
</thead>
<tbody>
<tr>
<td>Model has unique/complex pre-processing (VLM captioning, AR token generation, custom latent packing, etc.)</td>
<td><strong>Hybrid</strong> — consolidate into a BeforeDenoisingStage</td>
</tr>
<tr>
<td>Model fits neatly into standard text-to-image or text+image-to-image pattern</td>
<td><strong>Modular</strong> — use <code>add_standard_t2i_stages()</code> / <code>add_standard_ti2i_stages()</code></td>
</tr>
<tr>
<td>Porting a Diffusers pipeline with many custom steps</td>
<td><strong>Hybrid</strong> — copy the <code>__call__</code> logic into a single stage</td>
</tr>
<tr>
<td>Adding a variant of an existing model that shares most logic</td>
<td><strong>Modular</strong> — reuse existing stages, customize via PipelineConfig callbacks</td>
</tr>
<tr>
<td>A specific pre-processing step needs special parallelism or profiling isolation</td>
<td><strong>Modular</strong> — extract that step as a dedicated stage</td>
</tr>
</tbody>
</table>
## Key Components for Implementation
To add support for a new diffusion model, you will need to define or configure the following components:
1. **`PipelineConfig`**: A dataclass holding static configurations for your model pipeline — precision settings, model architecture parameters, and callback methods used by the standard `DenoisingStage` and `DecodingStage`. Each model has its own subclass.
2. **`SamplingParams`**: A dataclass defining runtime generation parameters — `prompt`, `negative_prompt`, `guidance_scale`, `num_inference_steps`, `seed`, `height`, `width`, etc.
3. **Pre-processing stage(s)**: Either a single model-specific `{Model}BeforeDenoisingStage` (Hybrid style) or a combination of standard stages (Modular style). See [Two Pipeline Styles](#two-pipeline-styles) above.
4. **`ComposedPipeline`**: A class that wires together your pre-processing stage(s) with the standard `DenoisingStage` and `DecodingStage`. See base definitions:
- [`ComposedPipelineBase`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py)
- [`PipelineStage`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py)
- [Central registry](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py)
5. **Modules (model components)**: Each pipeline references modules loaded from the model repository (e.g., Diffusers `model_index.json`):
- `text_encoder`: Encodes text prompts into embeddings.
- `tokenizer`: Tokenizes raw text input for the text encoder(s).
- `processor`: Preprocesses images and extracts features; often used in image-to-image tasks.
- `image_encoder`: Specialized image feature extractor.
- `dit/transformer`: The core denoising network (DiT/UNet architecture) operating in latent space.
- `scheduler`: Controls the timestep schedule and denoising dynamics.
- `vae`: Variational Autoencoder for encoding/decoding between pixel space and latent space.
## Pipeline Stages Reference
### Core Stages (used by all pipelines)
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "50%"}} />
<col style={{width: "50%"}} />
</colgroup>
<thead>
<tr>
<th>Stage Class</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr>
<td><code>DenoisingStage</code></td>
<td>Executes the main denoising loop, iteratively applying the model (DiT/UNet) to refine the latents.</td>
</tr>
<tr>
<td><code>DecodingStage</code></td>
<td>Decodes the final latent tensor back into pixel space using the VAE.</td>
</tr>
<tr>
<td><code>DmdDenoisingStage</code></td>
<td>A specialized denoising stage for DMD model architectures.</td>
</tr>
<tr>
<td><code>CausalDMDDenoisingStage</code></td>
<td>A specialized causal denoising stage for specific video models.</td>
</tr>
</tbody>
</table>
### Pre-processing Stages (for Modular Composition Style)
The following fine-grained stages can be composed to build the pre-processing portion of a pipeline. They are best suited for models whose pre-processing largely fits the standard patterns. If your model requires significant customization, consider the Hybrid style with a single `BeforeDenoisingStage` instead.
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "50%"}} />
<col style={{width: "50%"}} />
</colgroup>
<thead>
<tr>
<th>Stage Class</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr>
<td><code>InputValidationStage</code></td>
<td>Validates user-provided <code>SamplingParams</code>.</td>
</tr>
<tr>
<td><code>TextEncodingStage</code></td>
<td>Encodes text prompts into embeddings using one or more text encoders.</td>
</tr>
<tr>
<td><code>ImageEncodingStage</code></td>
<td>Encodes input images into embeddings, often used in image-to-image tasks.</td>
</tr>
<tr>
<td><code>ImageVAEEncodingStage</code></td>
<td>Encodes an input image into latent space using the VAE.</td>
</tr>
<tr>
<td><code>TimestepPreparationStage</code></td>
<td>Prepares the scheduler's timesteps for the diffusion process.</td>
</tr>
<tr>
<td><code>LatentPreparationStage</code></td>
<td>Creates the initial noisy latent tensor that will be denoised.</td>
</tr>
</tbody>
</table>
## Implementation Guide
### Step 1: Obtain and Study the Reference Implementation
Before writing any code, obtain the model's original implementation or Diffusers pipeline code:
- 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 to 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.
2. Read the Diffusers pipeline's `__call__` method to understand:
- How text prompts are encoded
- How latents are prepared (shape, dtype, scaling)
- How timesteps/sigmas are computed
- What conditioning kwargs the DiT expects
- How the denoising loop works
- How VAE decoding is done
### 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 need substantial structural changes or when no architecturally similar implementation exists.
- **Compare against existing pipelines** (Flux, Wan, Qwen-Image, GLM-Image, HunyuanVideo, LTX, etc.). If the new model shares most of its structure with an existing one, prefer adding a new config variant or reusing existing stages.
- **Check existing stages** in `runtime/pipelines_core/stages/` and `stages/model_specific_stages/`.
- **Check existing model components** — many models share VAEs (e.g., `AutoencoderKL`), text encoders (CLIP, T5), and schedulers. Reuse these directly.
### Step 3: Implement Model Components
Adapt the model's core components:
- **DiT/Transformer**: Implement in [`runtime/models/dits/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/)
- **Encoders**: Implement in [`runtime/models/encoders/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/encoders/)
- **VAEs**: Implement in [`runtime/models/vaes/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/vaes/)
- **Schedulers**: Implement in [`runtime/models/schedulers/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/schedulers/) if needed
Use SGLang's fused kernels where possible (see `LayerNormScaleShift`, `RMSNormScaleShift`, `apply_qk_norm`, etc.).
**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 implementations:
- **Wan model** (`runtime/models/dits/wanvideo.py`) — Full TP + SP: `ColumnParallelLinear`/`RowParallelLinear` for attention, sequence dimension sharding via `get_sp_world_size()`
- **Qwen-Image model** (`runtime/models/dits/qwen_image.py`) — SP via `USPAttention` (Ulysses + Ring Attention)
### Step 4: Create Configs
- **DiT Config**: `configs/models/dits/{model_name}.py`
- **VAE Config**: `configs/models/vaes/{model_name}.py`
- **SamplingParams**: `configs/sample/{model_name}.py`
### Step 5: Create PipelineConfig
The `PipelineConfig` provides callbacks that the standard `DenoisingStage` and `DecodingStage` use:
`PipelineConfig` is where shared denoising and decoding stages get model-specific
callbacks.
```python
# python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py
from dataclasses import dataclass, field
@dataclass
class MyModelPipelineConfig(ImagePipelineConfig):
task_type: ModelTaskType = ModelTaskType.T2I
vae_precision: str = "bf16"
should_use_guidance: bool = True
dit_config: DiTConfig = field(default_factory=MyModelDitConfig)
dit_config: DiTConfig = field(default_factory=MyModelDiTConfig)
vae_config: VAEConfig = field(default_factory=MyModelVAEConfig)
def get_freqs_cis(self, batch, device, rotary_emb, dtype):
"""Prepare rotary position embeddings for the DiT."""
...
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],
@@ -287,315 +158,173 @@ class MyModelPipelineConfig(ImagePipelineConfig):
}
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,
}
def get_decode_scale_and_shift(self):
"""Return (scale, shift) for latent denormalization before VAE decode."""
...
```
### Step 6: Implement Pre-processing
Make these kwargs match the denoising module's `forward()` signature exactly.
Choose based on your model's needs (see [How to Choose](#how-to-choose)):
### 3. Pipeline Wiring
#### Option A: BeforeDenoisingStage (Hybrid Style)
Create a single stage that handles all pre-processing. Best when the model has custom/complex pre-processing logic.
```python
# python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/my_model.py
class MyModelBeforeDenoisingStage(PipelineStage):
"""Monolithic pre-processing stage for MyModel.
Consolidates: input validation, text/image encoding, latent
preparation, and timestep computation.
"""
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
@torch.no_grad()
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
device = get_local_torch_device()
# 1. Encode prompt (model-specific logic)
prompt_embeds, negative_prompt_embeds = self._encode_prompt(...)
# 2. Prepare latents
latents = self._prepare_latents(...)
# 3. Prepare timesteps
timesteps, sigmas = self._prepare_timesteps(...)
# 4. Populate batch for DenoisingStage
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.generator = generator
batch.raw_latent_shape = latents.shape
return batch
```
#### Option B: Standard Stages (Modular Style)
Skip creating a custom stage entirely — configure via `PipelineConfig` callbacks and use framework helpers. Best when the model fits standard patterns.
(This option has no separate stage file; the pipeline class in Step 7 calls `add_standard_t2i_stages()` directly.)
**Key batch fields that `DenoisingStage` expects** (regardless of which option you choose):
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
</colgroup>
<thead>
<tr>
<th>Field</th>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr>
<td><code>batch.latents</code></td>
<td><code>torch.Tensor</code></td>
<td>Initial noisy latent tensor</td>
</tr>
<tr>
<td><code>batch.timesteps</code></td>
<td><code>torch.Tensor</code></td>
<td>Timestep schedule</td>
</tr>
<tr>
<td><code>batch.num_inference_steps</code></td>
<td><code>int</code></td>
<td>Number of denoising steps</td>
</tr>
<tr>
<td><code>batch.sigmas</code></td>
<td><code>list[float]</code></td>
<td>Sigma schedule (must be a Python list, not numpy)</td>
</tr>
<tr>
<td><code>batch.prompt_embeds</code></td>
<td><code>list[torch.Tensor]</code></td>
<td>Positive prompt embeddings (wrapped in a list)</td>
</tr>
<tr>
<td><code>batch.negative_prompt_embeds</code></td>
<td><code>list[torch.Tensor]</code></td>
<td>Negative prompt embeddings (wrapped in a list)</td>
</tr>
<tr>
<td><code>batch.generator</code></td>
<td><code>torch.Generator</code></td>
<td>RNG generator for reproducibility</td>
</tr>
<tr>
<td><code>batch.raw_latent_shape</code></td>
<td><code>tuple</code></td>
<td>Original latent shape before any packing</td>
</tr>
</tbody>
</table>
### Step 7: Define the Pipeline Class
#### Hybrid Style
Use the standard helper when the model fits it.
```python
# python/sglang/multimodal_gen/runtime/pipelines/my_model.py
class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "MyModelPipeline" # Must match model_index.json _class_name
pipeline_name = "MyModelPipeline"
_required_config_modules = [
"text_encoder", "tokenizer", "vae", "transformer", "scheduler",
"text_encoder",
"tokenizer",
"transformer",
"scheduler",
"vae",
]
def create_pipeline_stages(self, server_args: ServerArgs):
self.add_standard_t2i_stages()
EntryClass = [MyModelPipeline]
```
Use a model-specific pre-processing stage when the reference pipeline cannot be
cleanly expressed by standard helpers.
```python
class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "MyModelPipeline"
_required_config_modules = [
"text_encoder",
"tokenizer",
"transformer",
"scheduler",
"vae",
]
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"),
),
vae=self.get_module("vae"),
)
)
# 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()
EntryClass = [MyModelPipeline]
```
#### Modular Style
### 4. Optional Before-Denoising Stage
A `BeforeDenoisingStage` should populate the batch fields consumed by
`DenoisingStage`.
```python
# python/sglang/multimodal_gen/runtime/pipelines/my_model.py
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)
class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "MyModelPipeline"
_required_config_modules = [
"text_encoder", "tokenizer", "vae", "transformer", "scheduler",
]
def create_pipeline_stages(self, server_args: ServerArgs):
# All pre-processing + denoising + decoding in one call
self.add_standard_t2i_stages(
prepare_extra_timestep_kwargs=[prepare_mu], # model-specific hooks
)
EntryClass = [MyModelPipeline]
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
```
### Step 8: Register the Model
Required fields for `DenoisingStage`:
Register your configs in [`registry.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py):
| 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. Denoising Module
Add a file under `runtime/models/dits/` only when the architecture is new. Reuse
existing encoders, VAEs, schedulers, normalization layers, and fused kernels
whenever possible.
For multi-GPU serving, add TP/SP support after the single-GPU path is correct.
Useful references:
- `runtime/models/dits/wanvideo.py` for TP plus SP.
- `runtime/models/dits/qwen_image.py` for USP attention.
### 6. Registry
Register the family once the sampling params and pipeline config exist.
```python
register_configs(
model_family="my_model",
sampling_param_cls=MyModelSamplingParams,
pipeline_config_cls=MyModelPipelineConfig,
hf_model_paths=["org/my-model-name"],
hf_model_paths=["org/my-model"],
)
```
The `EntryClass` in your pipeline file is automatically discovered by the registry — no additional registration needed for the pipeline class itself.
The pipeline file is discovered through its `EntryClass`; do not add a second
pipeline registry unless the existing registry requires it.
### Step 9: Verify Output Quality
## Verify the Port
After implementation, verify that the generated output is not noise. A noisy or garbled output is the most common sign of an incorrect implementation. Common causes include:
Use one deterministic prompt and seed while comparing with the reference
implementation.
- Incorrect latent scale/shift factors
- Wrong timestep/sigma schedule (order, dtype, or value range)
- Mismatched conditioning kwargs
- Rotary embedding style mismatch (`is_neox_style`)
1. Run a single-GPU smoke test and check that the output contains coherent
content.
2. Compare latent scale and shift, timestep order, sigma values, and conditioning
kwargs against Diffusers or the official implementation.
3. Verify VAE decode and post-processing separately from denoising.
4. If the model supports LoRA, CFG parallelism, TP, SP, or disaggregation, test
each feature explicitly.
5. Add or update docs, examples, or the compatibility matrix when users need a
new launch command.
Debug by comparing intermediate tensor values against the Diffusers reference pipeline with the same seed.
Common failure points:
## Reference Implementations
- 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`.
### Hybrid Style
## PR Checklist
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "25%"}} />
<col style={{width: "25%"}} />
<col style={{width: "25%"}} />
<col style={{width: "25%"}} />
</colgroup>
<thead>
<tr>
<th>Model</th>
<th>Pipeline</th>
<th>BeforeDenoisingStage</th>
<th>PipelineConfig</th>
</tr>
</thead>
<tbody>
<tr>
<td>GLM-Image</td>
<td><code>runtime/pipelines/glm_image.py</code></td>
<td><code>stages/model_specific_stages/glm_image.py</code></td>
<td><code>configs/pipeline_configs/glm_image.py</code></td>
</tr>
<tr>
<td>Qwen-Image-Layered</td>
<td><code>runtime/pipelines/qwen_image.py</code></td>
<td><code>stages/model_specific_stages/qwen_image_layered.py</code></td>
<td><code>configs/pipeline_configs/qwen_image.py</code></td>
</tr>
</tbody>
</table>
### Modular Style
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
</colgroup>
<thead>
<tr>
<th>Model</th>
<th>Pipeline</th>
<th>Notes</th>
</tr>
</thead>
<tbody>
<tr>
<td>Qwen-Image (T2I)</td>
<td><code>runtime/pipelines/qwen_image.py</code></td>
<td>Uses <code>add_standard_t2i_stages()</code></td>
</tr>
<tr>
<td>Qwen-Image-Edit</td>
<td><code>runtime/pipelines/qwen_image.py</code></td>
<td>Uses <code>add_standard_ti2i_stages()</code></td>
</tr>
<tr>
<td>Flux</td>
<td><code>runtime/pipelines/flux.py</code></td>
<td>Uses <code>add_standard_t2i_stages()</code> with custom <code>prepare_mu</code></td>
</tr>
<tr>
<td>Wan</td>
<td><code>runtime/pipelines/wan_pipeline.py</code></td>
<td>Uses <code>add_standard_ti2v_stages()</code></td>
</tr>
</tbody>
</table>
## Checklist
Before submitting your implementation, verify:
**Common (both styles):**
- [ ] **Pipeline file** 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`
- [ ] **Model configs** (DiT, VAE) at `configs/models/dits/` and `configs/models/vaes/`
- [ ] **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`, etc.) match the DiT's `forward()` signature
- [ ] Uses framework-standard `DenoisingStage` and `DecodingStage` (not custom denoising loops)
- [ ] **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 batch fields required by `DenoisingStage`
- [ ] 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.
- [ ] Tested multi-GPU paths if the PR claims TP, SP, CFG parallelism, or
distributed serving support.