> ## Documentation Index > Fetch the complete documentation index at: https://docs.sglang.io/llms.txt > Use this file to discover all available pages before exploring further. # Support New Diffusion Models > A concise implementation guide for adding diffusion model families to 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. ## Read the Code in This Order The files are split by runtime responsibility. For a new model, read the request path first: 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. 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. ## 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 | | ----------------------------- | ------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | | 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 | 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: 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 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: 1. **Use native stages directly.** This keeps the model on shared code paths for offload, component readiness, profiling, disaggregation, batching, and future stage-level optimizations. 2. **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. 3. **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. 4. **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 theme={null} # 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 theme={null} # 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 theme={null} # 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. ```python theme={null} 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. ```python theme={null} 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. ```python theme={null} 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] ``` ### 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`. ```python theme={null} 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 Register the family once the sampling params and pipeline config exist. ```python theme={null} register_configs( model_family="my_model", sampling_param_cls=MyModelSamplingParams, pipeline_config_cls=MyModelPipelineConfig, hf_model_paths=["org/my-model"], ) ``` 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. 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. Compare VAE decode separately, including tiled and multi-GPU `spatial_shard`. 4. Run encoder and DiT TP, SP, combined TP x SP, and `--layerwise-offload-components all`; compare with the single-GPU resident baseline. 5. If the model supports LoRA, CFG parallelism, or disaggregation, test each feature explicitly. 6. Add or update docs, examples, or the compatibility matrix 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.