docs: simplify diffusion new model guide (#30109)

2026-07-17 19:39:38 +08:00
parent 681c223570
commit 85ac56c823
1 changed files with 239 additions and 510 deletions
@@ -1,285 +1,156 @@
 ---
-title: "How to Support New Diffusion Models"
+title: "Support New Diffusion Models"
-metatags:
+description: "A concise implementation guide for adding diffusion model families to SGLang-Diffusion."
    description: "This document explains how to add support for new diffusion models in 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
+The files are split by runtime responsibility. For a new model, read the
-design allows developers to construct pipelines for various diffusion models while keeping the core generation
+request path first:
 loop standardized for optimization.
-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.
+That is the dependency direction. Avoid making a model PR that requires readers
-   **`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.
+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. |
-```
+Do not add a folder just to mirror the Diffusers repository layout. Add a new
-BeforeDenoisingStage (model-specific)  →  DenoisingStage (standard)  →  DecodingStage (standard)
+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"}}>
+### 2. Pipeline Config
  <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>
-**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.
+`PipelineConfig` is where shared denoising and decoding stages get model-specific
-
+callbacks.
 #### 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:
 ```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)
+Use the standard helper when the model fits it.
 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
 ```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):
+        batch.prompt_embeds = [prompt_embeds]
-    pipeline_name = "MyModelPipeline"
+        batch.negative_prompt_embeds = [negative_prompt_embeds]
-
+        batch.latents = latents
-    _required_config_modules = [
+        batch.timesteps = timesteps
-        "text_encoder", "tokenizer", "vae", "transformer", "scheduler",
+        batch.num_inference_steps = len(timesteps)
-    ]
+        batch.sigmas = sigmas.tolist()
-
+        batch.raw_latent_shape = latents.shape
-    def create_pipeline_stages(self, server_args: ServerArgs):
+        return batch
        # All pre-processing + denoising + decoding in one call
        self.add_standard_t2i_stages(
            prepare_extra_timestep_kwargs=[prepare_mu],  # model-specific hooks
        )
 EntryClass = [MyModelPipeline]
 ```
-### 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
+1. Run a single-GPU smoke test and check that the output contains coherent
- Wrong timestep/sigma schedule (order, dtype, or value range)
+   content.
- Mismatched conditioning kwargs
+2. Compare latent scale and shift, timestep order, sigma values, and conditioning
- Rotary embedding style mismatch (`is_neox_style`)
+   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"}}>
+- [ ] Reused an existing family, stage, module, scheduler, or VAE wherever
-  <colgroup>
+      possible.
-    <col style={{width: "25%"}} />
+- [ ] Kept the new-model touch surface small and justified any extra files.
-    <col style={{width: "25%"}} />
+- [ ] Added `SamplingParams`, `PipelineConfig`, pipeline wiring, DiT module, and
-    <col style={{width: "25%"}} />
+      registry entry when native support is needed.
-    <col style={{width: "25%"}} />
+- [ ] Confirmed `pipeline_name` matches the Diffusers `model_index.json`
-  </colgroup>
+      `_class_name` when applicable.
-  <thead>
+- [ ] Confirmed `_required_config_modules` matches the model repo.
-    <tr>
+- [ ] Verified image or video quality against a reference output.
-      <th>Model</th>
+- [ ] Tested multi-GPU paths if the PR claims TP, SP, CFG parallelism, or
-      <th>Pipeline</th>
+      distributed serving support.
      <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`