Two Pipeline Styles#

SGLang Diffusion supports two pipeline composition styles. Both are valid; choose the one that best fits your model.

BeforeDenoisingStage (model-specific)  →  DenoisingStage (standard)  →  DecodingStage (standard)

Core Stages (used by all pipelines)#

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.

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/

• Encoders: Implement in runtime/models/encoders/

• VAEs: Implement in runtime/models/vaes/

• Schedulers: Implement in 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/sglang/multimodal_gen/configs/pipeline_configs/my_model.py

@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)
    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],
            "timestep": t,
        }

    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#

Choose based on your model’s needs (see How to Choose):

# 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

Step 7: Define the Pipeline Class#

# python/sglang/multimodal_gen/runtime/pipelines/my_model.py

class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "MyModelPipeline"  # Must match model_index.json _class_name

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

    def create_pipeline_stages(self, server_args: ServerArgs):
        # 1. Monolithic pre-processing (model-specific)
        self.add_stage(
            MyModelBeforeDenoisingStage(
                vae=self.get_module("vae"),
                text_encoder=self.get_module("text_encoder"),
                tokenizer=self.get_module("tokenizer"),
                transformer=self.get_module("transformer"),
                scheduler=self.get_module("scheduler"),
            ),
        )

        # 2. Standard denoising loop (framework-provided)
        self.add_stage(
            DenoisingStage(
                transformer=self.get_module("transformer"),
                scheduler=self.get_module("scheduler"),
            ),
        )

        # 3. Standard VAE decoding (framework-provided)
        self.add_standard_decoding_stage()


EntryClass = [MyModelPipeline]

# python/sglang/multimodal_gen/runtime/pipelines/my_model.py

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]

Step 8: Register the Model#

Register your configs in registry.py:

register_configs(
    model_family="my_model",
    sampling_param_cls=MyModelSamplingParams,
    pipeline_config_cls=MyModelPipelineConfig,
    hf_model_paths=["org/my-model-name"],
)

The EntryClass in your pipeline file is automatically discovered by the registry — no additional registration needed for the pipeline class itself.

Step 9: Verify Output Quality#

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:

• Incorrect latent scale/shift factors

• Wrong timestep/sigma schedule (order, dtype, or value range)

• Mismatched conditioning kwargs

• Rotary embedding style mismatch (is_neox_style)

Debug by comparing intermediate tensor values against the Diffusers reference pipeline with the same seed.

Hybrid Style#

Modular Style#