Visual Generation (Beta)

Note

This feature is in beta stage. APIs, supported models, and optimization options are actively evolving and may change in future releases.

Background

Visual generation models based on diffusion transformers (DiT) have become the standard for high-quality image and video synthesis. These models iteratively denoise latent representations through a learned transformer backbone, then decode the final latents with a VAE to produce pixels.

TensorRT-LLM VisualGen provides a unified inference stack for diffusion models, with a pipeline architecture separate from the LLM inference path. Key capabilities include:

• A shared pipeline abstraction covering the denoising loop, guidance strategies, and component loading.

• Pluggable attention backends: PyTorch SDPA (VANILLA), TRT-LLM kernels (TRTLLM), cuDNN fused SDPA (CUDNN), FlashInfer FP16/BF16 dense prefill (FLASHINFER), TRT-LLM CuTe DSL kernels (CUTEDSL, Blackwell-class GPUs), and Flash Attention 4 (FA4).

• Quantization support (dynamic and static) using the ModelOpt configuration format.

• Quantized attention support: see VisualGen Quantized Attention.

• Sparse attention support: see VisualGen Sparse Attention.

• Multi-GPU parallelism (CFG parallel, Ulysses sequence parallel, Tensor parallelism).

• Step caching — two runtime caching backends (TeaCache and Cache-DiT) that skip transformer computation on steps where the step-to-step change is small.

• CPU offloading to reduce peak GPU memory usage.

• trtllm-serve integration with OpenAI-compatible API endpoints for image and video generation.

Supported Models

HuggingFace Model ID	Tasks
black-forest-labs/FLUX.1-dev	Text-to-Image
black-forest-labs/FLUX.2-dev	Text-to-Image
Wan-AI/Wan2.1-T2V-1.3B-Diffusers	Text-to-Video
Wan-AI/Wan2.1-T2V-14B-Diffusers	Text-to-Video
FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers	Text-to-Video (VSA)
Wan-AI/Wan2.1-I2V-14B-480P-Diffusers	Image-to-Video
Wan-AI/Wan2.1-I2V-14B-720P-Diffusers	Image-to-Video
Wan-AI/Wan2.2-T2V-A14B-Diffusers	Text-to-Video
Wan-AI/Wan2.2-I2V-A14B-Diffusers	Image-to-Video
Wan-AI/Wan2.2-TI2V-5B-Diffusers	Text-to-Video, Image-to-Video
FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers	Text-to-Video (3-step distilled)
Lightricks/LTX-2	Text-to-Video (with Audio), Image-to-Video (with Audio)
Qwen/Qwen-Image	Text-to-Image
Qwen/Qwen-Image-2512	Text-to-Image
Qwen/Qwen-Image-Layered	Image-to-Image
Qwen/Qwen-Image-Edit-2511	Image Editing (text+images-to-image)
nvidia/Cosmos3-Nano	Text-to-Image, Text-to-Video, Image-to-Video
nvidia/Cosmos3-Super	Text-to-Image, Text-to-Video, Image-to-Video
nvidia/Cosmos3-Super-Text2Image-4Step	Text-to-Image (DMD2-distilled, fixed 4-step schedule)
nvidia/Cosmos3-Super-Image2Video-4Step	Image-to-Video (DMD2-distilled, fixed 4-step schedule)
nvidia/Cosmos3-Edge	Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native)
hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v	Text-to-Video
hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v	Text-to-Video
MiniMaxAI/MiniMax-H3	Text-to-Video (with Audio), First/Last-Frame-to-Video (with Audio)
zai-org/GLM-Image	Text-to-Image

Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via model_index.json (or modular_model_index.json for modular pipelines); LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The AutoPipeline registry selects the appropriate pipeline class automatically.

Feature Matrix

Model	FP8 blockwise	NVFP4	TeaCache	Cache-DiT	CPU Offloading	CFG Parallelism	Ulysses Parallelism	Parallel VAE	CUDA Graph	torch.compile	trtllm-serve	Attention2D	Ring Attention	Tensor Parallelism	VSA
FLUX.1	Yes	Yes	Yes	Yes	No	No	Yes	No	Yes	Yes	Yes	Yes	Yes	Yes	No
FLUX.2	Yes	Yes	Yes	Yes	No	No	Yes	No	Yes	Yes	Yes	Yes	Yes	Yes	No
Wan 2.1	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	No
Wan 2.1 VSA	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	No	No	Yes	Yes
Wan 2.2	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	Yes	No
FastWan 2.2	Yes	Yes	No	No	No	No	No	No	Yes	Yes	Yes	No	No	No	No
LTX-2	Yes	Yes	Yes	Yes	No	Yes	Yes	No	No	Yes	Yes	Yes	Yes	No	No
MiniMax-H3	Yes	Yes	No	No	No	No	No	No	No	Yes	Yes	No	No	No	No
Qwen-Image	Yes	Yes	Yes	Yes	No	Yes	Yes	No	Yes	Yes	Yes	Yes	Yes	No	No
Qwen-Image-Layered	No	No	No	No	No	No	No	No	Yes	Yes	Yes	No	No	No	No
Qwen-Image-Edit-2511	Yes	Yes	No	No	No	Yes	No	No	Yes	Yes	Yes	No	No	No	No
Cosmos3	Yes	Yes	No	No	Yes	Yes	Yes	Yes	Yes	Yes	Yes	No	No	Yes	No
HunyuanVideo 1.5	Yes	Yes	No	No	No	No	No	No	No	No	Yes	No	No	No	No
GlmImage	Yes	Yes	No	No	No	No	No	No	No	No	Yes	No	No	No	No

Quick Start

Here is a simple example to generate a video with Wan 2.1:

python examples/visual_gen/quickstart_example.py

To learn more about VisualGen, see examples/visual_gen/ for more examples including text-to-image, image-to-video, and batch generation.

Usage with trtllm-serve

The trtllm-serve command automatically detects diffusion models (by the presence of model_index.json) and launches an OpenAI-compatible visual generation server with image and video generation endpoints.

See examples/visual_gen/serve/ for server launch instructions, example configurations, and API usage.

Serving Endpoints

When served via trtllm-serve, the following OpenAI-compatible endpoints are available:

Endpoint	Method	Purpose
/v1/images/generations	POST	Synchronous image generation
/v1/images/edits	POST	Image editing
/v1/videos	POST	Asynchronous video generation
/v1/videos/sync	POST	Synchronous video generation
/v1/videos/generations	POST	Deprecated alias of /v1/videos/sync (kept for back-compat)
/v1/videos/{id}	GET	Video status / metadata
/v1/videos/{id}/content	GET	Download generated video
/v1/videos/{id}	DELETE	Delete generated video
/v1/videos	GET	List all videos

The asynchronous /v1/videos job advances through GET /v1/videos/{id}: queued → generating (model inference) → postprocessing (encode the media and/or write the output file) → completed. The generating → postprocessing transition marks the end of inference; the video is downloadable via /content once completed.

response_format="path" returns the generated file’s server-side path (under TRTLLM_MEDIA_STORAGE_PATH) for co-located clients, enabled by default. Set TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1 to reject such requests with HTTP 400; the same switch also rejects a reference sent with format="path", since both ask the server to trust a local filesystem path. See the serve examples for the full response_format reference.

Reference Inputs

Qwen-Image-Layered supports BF16 image-conditioned layer decomposition through trtllm-serve image-edit routing. It returns one RGBA image per generated layer; set extra_params.save_layers_to_grid to true to pack layers into one saveable image grid. Quantization, cache acceleration, attention-parallel/Sage/VSA backends, and Tensor Parallelism are not enabled for this pipeline yet.

Conditioning references are supplied through the typed fields image_reference, video_reference, and audio_reference. Each field takes a single reference or a list. A reference is MediaRef(content=..., format=...), and format is required.

format	Content	Notes
path	A local file readable by the coordinator process	Bare path or file:// URI.
url	An http(s) URL	Fetched on the coordinator through the SSRF-guarded loader.
base64	Base64 text	A data: URI is also accepted.
bytes	Raw bytes	Python API only.

Every pipeline declares the reference slots and roles it accepts through ref_slot_specs, and a request is validated against that declaration before generation begins. References are resolved to raw bytes on the coordinator, so a worker never needs a filesystem shared with the client.

Most models take a single reference whose role is unambiguous:

from tensorrt_llm import VisualGen
from tensorrt_llm.visual_gen import MediaRef

vg = VisualGen(model="Wan-AI/Wan2.2-TI2V-5B-Diffusers")
params = vg.default_params
params.image_reference = MediaRef(content="start.png", format="path")
output = vg.generate(inputs="the scene comes alive with gentle motion", params=params)

Models that accept the same modality in more than one role need role. Wan 2.1 I2V takes a first frame and an optional last frame:

from tensorrt_llm import VisualGen
from tensorrt_llm.visual_gen import MediaRef

vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers")
params = vg.default_params
params.image_reference = [
    MediaRef(content="start.png", format="path", role="first_frame"),
    MediaRef(content="end.png", format="path", role="last_frame"),
]

FLUX.2 and Qwen-Image-Edit accept a list of reference images on image_reference.

The same fields carry references over trtllm-serve; see examples/visual_gen/serve/ for request examples.

MiniMax-H3 Notes

• Text-to-video (T2VA) and first/last-frame-to-video (FL2VA) are supported. Reference-to-video (Ref2VA) is not enabled yet.

• MiniMax-H3 currently restricts TRTLLM attention to SM100 or SM103. This is a model-specific restriction, not a general VisualGen backend requirement.

• The published MiniMax-H3 checkpoint license restricts use by territory. Obtain legal approval before downloading or running the weights.

Optimizations

Quantization

VisualGen supports both dynamic quantization (on-the-fly at weight-loading time from BF16 checkpoints) and static quantization (loading pre-quantized checkpoints with embedded scales). Both modes use the ModelOpt quantization_config format.

Configure via VisualGenArgs.quant_config (YAML or programmatic):

quant_config:
  quant_algo: FP8        # or FP8_BLOCK_SCALES, NVFP4
  dynamic: true

from tensorrt_llm import VisualGenArgs
args = VisualGenArgs(model="/path/to/model", quant_config={"quant_algo": "FP8", "dynamic": True})

Omit quant_config for BF16/FP16 baseline.

Wan VAE NVFP4

The VAE can represent a substantial fraction of end-to-end latency in distilled video-generation pipelines with few denoising steps. Blackwell Tensor Cores offer up to four times the peak NVFP4 compute throughput of BF16, making VAE Conv3d operators an important optimization target. FP4 VAE replaces eligible native VAE Conv3d operators with NVFP4 weight-and-activation kernels while retaining BF16 outputs.

Linear-layer, attention, and VAE quantization are selected independently. Use quant_config for transformer linear layers, attention_config.quant_attention_config for attention, and vae_config.quant_conv_config for VAE convolutions. NVFP4 VAE execution currently supports native Wan-family VAEs on SM100 and SM103 GPUs. Unsupported pipelines and algorithms fail before pipeline construction. On an unsupported device, an explicit NVFP4 request fails; checkpoint-driven NVFP4 instead uses dequantized BF16 operators.

By default, vae_config.quant_conv_config is unset (None) and VAE convolutions follow the checkpoint metadata. A high-precision checkpoint remains high precision. A packed NVFP4 checkpoint selects NVFP4 execution automatically and reuses any valid calibrated activation scales; layers without one derive it dynamically.

To quantize eligible Wan Conv3d layers from a high-precision checkpoint and derive activation scales dynamically, use the shorthand configuration:

vae_config:
  quant_conv_config:
    quant_algo: NVFP4
    dynamic: true

The quant_conv_config.dynamic shorthand sets both weight and activation modes. Configure them independently when their sources differ. For example, packed NVFP4 weights with rank-local dynamic activation scales use:

vae_config:
  quant_conv_config:
    quant_algo: NVFP4
    config_groups:
      default:
        weights:
          dynamic: false
        input_activations:
          dynamic: true

weights.dynamic: true requires high-precision checkpoint weights, while false requires packed NVFP4 weights. input_activations.dynamic: false requires a valid calibrated checkpoint scale for every selected convolution; true derives rank-local activation scales at runtime. Do not specify both the top-level shorthand and config_groups.

The optional ignore list excludes matching VAE modules. With a high-precision checkpoint, excluded convolutions remain in BF16. With a packed NVFP4 checkpoint, their weights can only be dequantized back to BF16; the original high-precision weights cannot be recovered.

See the Model Optimizer diffusion example for a Wan 2.2 VAE NVFP4 calibration and checkpoint-generation recipe.

Runtime LoRA

VisualGen can preload a local LoRA adapter at startup and fuse its deltas into transformer weights before warmup, CUDA graph capture, and cache acceleration setup. Configure this through VisualGenArgs.runtime_lora_config in Python or YAML:

runtime_lora_config:
  path: /path/to/adapter-or-safetensors
  target_components:
    - transformer

from tensorrt_llm import VisualGenArgs
from tensorrt_llm.visual_gen import RuntimeLoRAConfig

args = VisualGenArgs(
    model="/path/to/model",
    runtime_lora_config=RuntimeLoRAConfig(
        path="/path/to/adapter-or-safetensors",
        target_components=["transformer"],
    ),
)

The loader accepts safetensors adapters that use Comfy/Kohya-style lora_down / lora_up keys or PEFT-style lora_A / lora_B keys. It applies .alpha tensors when present, and reads lora_alpha from a colocated adapter_config.json for PEFT adapters. scale multiplies the resulting alpha/rank factor.

By default, strict=True raises when adapter tensors cannot be matched, have unsupported shapes, or partially apply to the selected transformer component. Set target_components explicitly for pipelines with multiple transformer components. Runtime LoRA is not supported with VisualGen weight quantization, and startup fusion does not support per-request adapter switching.

Quantized Attention

In addition to linear-layer quantization, VisualGen exposes multiple backend-specific quantized-attention recipes that operate inside the attention kernel. They are configured through AttentionConfig.quant_attention_config and can be enabled independently with any linear layer configuration. See VisualGen Quantized Attention for the full recipe table, the V scale-granularity trade-off, and the block-scaled MXFP8 / NVFP4 recipes.

CUDA Graphs

VisualGen CUDA graphs capture transformer forward calls during denoising and replay them for later steps with compatible inputs. See VisualGen CUDA Graphs for capture scope, graph keys, and sparse-attention phase behavior.

Step Caching

Both caching backends are configured through VisualGenArgs.cache_config. The backend is selected by the cache_backend discriminator field.

FastWan 2.2 is a 3-step distilled model; TeaCache and Cache-DiT are not applicable.

TeaCache

TeaCache caches transformer outputs when timestep embeddings change slowly between denoising steps, skipping redundant computation. Enable via VisualGenArgs.cache_config (YAML or programmatic):

cache_config:
  cache_backend: teacache
  teacache_thresh: 0.2

Parameter	Type	Default	Description
teacache_thresh	float	0.2	Accumulated timestep-embedding distance threshold. A step is skipped when the accumulated polynomial-rescaled L1 change stays below this value; higher values cache more aggressively (more speedup, possible quality loss). The example configs use 0.6 for FLUX.1 and 0.2 for all other supported models.
use_ret_steps	bool	false	Enable retention-step caching variant.
coefficients	list[float]	per-model	Polynomial coefficients used by the TeaCache decision function. Set automatically at load time based on the checkpoint.

Wan 2.2 and LTX-2 have no built-in TeaCache coefficient tables. Set cache_config.coefficients explicitly for LTX-2. Wan 2.2 requires both cache_config.coefficients for its high-noise transformer and cache_config.coefficients_2 for its low-noise transformer.

Cache-DiT

Cache-DiT uses residual-difference gating (DBCache) to adaptively skip transformer blocks, with optional TaylorSeer polynomial prediction and step-computation mask (SCM).

Enable via VisualGenArgs.cache_config:

cache_config:
  cache_backend: cache_dit

from tensorrt_llm import VisualGenArgs
from tensorrt_llm.visual_gen import CacheDiTConfig

args = VisualGenArgs(
    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
    cache_config=CacheDiTConfig(
        residual_diff_threshold=0.20,
        max_continuous_cached_steps=4,
    ),
)

Commonly used parameters:

Parameter	Type	Default	Description
Fn_compute_blocks	int	1	Number of leading transformer blocks that are always fully computed at every denoising step (Fn in the Cache-DiT paper).
Bn_compute_blocks	int	0	Number of trailing transformer blocks used for prediction refinement (Bn).
max_warmup_steps	int	4	Initial denoising steps that always run a full forward pass; caching is disabled for this many steps at the start.
max_cached_steps	int	-1	Total cap on cached (skipped) steps across the run; -1 means unlimited.
max_continuous_cached_steps	int	3	Maximum consecutive cached steps before a forced full-compute step is inserted. -1 means unlimited.
residual_diff_threshold	float	0.24	L1-distance threshold for DBCache residual gating. Increase to cache more aggressively (higher speedup, potential quality loss); decrease for more conservative caching.
enable_taylorseer	bool	false	Enable TaylorSeer calibration. Uses Taylor series expansion to approximate hidden states at cached steps, improving output quality over plain residual reuse.
taylorseer_order	int	1	Polynomial order for TaylorSeer (1–4). Only used when enable_taylorseer=true.
scm_steps_mask_policy	str | None	None	Named step-computation mask policy from the cache_dit library ("slow", "medium", "fast", "ultra").
scm_steps_policy	"dynamic" | "static"	"dynamic"	Execution policy for the SCM mask; only active when scm_steps_mask_policy is set.
force_refresh_step_hint	int | None	None	Step index at which a forced full-compute pass is injected (useful for scheduled quality checkpoints).
force_refresh_step_policy	"once" | "repeat"	"once"	Whether force_refresh_step_hint fires only on the first call ("once") or at that interval repeatedly ("repeat").

Wan 2.2 dual-transformer note: Wan 2.2 uses two expert transformers (high-noise and low-noise stacks). All CacheDiTConfig parameters apply to both stacks, except max_warmup_steps and max_cached_steps: the low-noise stack always uses fixed internal caps (max_warmup_steps=2, max_cached_steps=20) regardless of user config.

Video Sparse Attention (VSA)

VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs a block-sparse CuTe kernel over only those blocks. The two outputs are blended with learned gates.

Requirements:

• VSA-fine-tuned checkpoint: FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers. Standard Wan checkpoints do not have the learned VSA gates.

• Blackwell GPU (sm_100+) for the CuTe JIT kernel. Falls back to dense SDPA on older hardware with no accuracy loss.

• CUTEDSL attention backend.

• Not compatible with Ring attention or Attention2D (VSA does not produce per-split LSE). Ulysses is supported.

vsa_sparsity controls the fraction of K/V blocks skipped in the fine branch (0.0 = dense, 0.9 = 90% blocks skipped). Higher sparsity gives more speedup at the cost of some quality.

Python API:

from tensorrt_llm import VisualGenArgs
from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig

args = VisualGenArgs(
    model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers",
    attention_config=AttentionConfig(
        backend="CUTEDSL",
        sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9),
    ),
)

YAML (for use with --visual_gen_args or trtllm-serve):

attention_config:
  backend: CUTEDSL
  sparse_attention_config:
    algorithm: vsa
    vsa_sparsity: 0.90

CPU Offloading

CPU offloading stages move selected Wan and Cosmos3 T2V pipeline components between CPU and GPU to reduce peak GPU memory usage; enable it with cpu_offload_config.enable: true.

Multi-GPU Parallelism

Configured under VisualGenArgs.parallel_config. Modes can be combined:

• CFG Parallelism (cfg_size: 2): Splits positive/negative guidance prompts across GPUs. FLUX uses embedded guidance without a separate negative prompt path; CFG parallelism is not applicable to FLUX or the distilled FastWan 2.2 model.

• Ulysses Parallelism (ulysses_size: N): Splits the sequence dimension across GPUs for longer sequences.

  • Async Ulysses A2A pipeline (async_ulysses: true in parallel_config): Overlaps per-rank V/Q/K projection compute with the cross-rank all-to-all on a dedicated side stream. Requires ulysses_size > 1 and an NVLink-connected GPU domain (uses PyTorch _SymmetricMemory with CUDA IPC for peer pushes; not currently supported across nodes without MNNVL). Currently wired for WAN and LTX-2 self-attention.

• Parallel VAE (parallel_vae_size: N): Shards the final VAE decode along a spatial axis (constraint: parallel_vae_size ≤ world_size; WAN/Cosmos3 only).

• Context Parallel (CP) — Partitions the sequence into shards so that each rank computes partial attention. Requires an LSE-capable attention backend (FA4 or CUTEDSL). CP can be composed with Ulysses, giving a total sequence-parallel (SP) degree = cp_size · ulysses_size. The CP degree depends on the implementation below:

  • Attention2D (attn2d_size: [N, M]): Shards the sequence axis across an N × M device mesh (CP degree = N · M; total SP degree = N · M · ulysses_size).

  • Ring Attention (ring_size: N): Shards the sequence axis across a 1D ring of N ranks, streaming K/V blocks (CP degree = N; total SP degree = N · ulysses_size; mutually exclusive with Attention2D).

• Tensor Parallelism (tp_size: N): Splits attention heads and transformer MLPs across GPUs for faster compute and reduced memory usage.

For multi-node execution, VisualGen relies on the external launcher to terminate the remaining ranks when any rank exits. torchrun provides this behavior. With SLURM, launch the VisualGen rank group with srun --kill-on-bad-exit=1; without this option, surviving ranks can continue running and retain GPU memory after a peer or the coordinator exits.

Developer Guide

Architecture Overview

The VisualGen module lives under tensorrt_llm._torch.visual_gen. At a high level, the inference flow is:

1. Config — User-facing VisualGenArgs (CLI / YAML) is merged with checkpoint metadata into DiffusionModelConfig.

2. Pipeline creation & loading — AutoPipeline detects the model type from model_index.json, instantiates the matching BasePipeline subclass, and loads weights (with optional dynamic quantization) and standard components (VAE, text encoder, tokenizer, scheduler).

3. Execution — DiffusionExecutor coordinates multi-GPU inference via worker processes communicating over ZeroMQ IPC.

Key components:

Component	Location	Role
VisualGen	tensorrt_llm/visual_gen/__init__.py	High-level API: manages workers, generate() / generate_async()
DiffusionExecutor	visual_gen/executor.py	Worker process: loads pipeline, processes requests via ZeroMQ
BasePipeline	visual_gen/pipeline.py	Base class: denoising loop, CFG handling, step caching (TeaCache / Cache-DiT), CUDA graph
AutoPipeline	visual_gen/pipeline_registry.py	Factory: auto-detects model type, selects pipeline class
PipelineLoader	visual_gen/pipeline_loader.py	Resolves checkpoint, loads config/weights, creates pipeline
TeaCacheAccelerator / CacheDiTAccelerator	visual_gen/cache/	Runtime caching backends (TeaCache, Cache-DiT) wrapping the transformer forward
WeightLoader	visual_gen/checkpoints/	Loads transformer weights from safetensors/bin

VisualGen is a parallel inference subsystem within TensorRT-LLM. It shares low-level primitives (Mapping, QuantConfig, Linear, RMSNorm, ZeroMqQueue, TrtllmAttention) but has its own executor, scheduler (diffusers-based), request types, and pipeline architecture separate from the LLM autoregressive decode path.

Implementing a New Diffusion Model

Adding a new model (e.g., a hypothetical “MyDiT”) requires four steps. The framework handles weight loading, parallelism, quantization, and serving automatically once the pipeline is registered.

1. Create the Transformer Module

Create the DiT backbone in tensorrt_llm/_torch/visual_gen/models/mydit/transformer_mydit.py. It should be an nn.Module that:

• Uses existing modules (e.g., Attention with configurable attention backend, Linear for builtin linear ops) wherever possible.

• Implements load_weights(weights: Dict[str, torch.Tensor]) to map checkpoint weight names to module parameters.

2. Create the Pipeline Class

Create a pipeline class extending BasePipeline in tensorrt_llm/_torch/visual_gen/models/mydit/. Override methods for transformer initialization, component loading, and inference. BasePipeline provides the denoising loop, CFG handling, and TeaCache integration — your pipeline only needs to implement model-specific logic. See WanPipeline for a reference implementation.

3. Register the Pipeline

Use the @register_pipeline("MyDiTPipeline") decorator on your pipeline class to register it in the global PIPELINE_REGISTRY. Make sure to export it from models/__init__.py.

4. Update AutoPipeline Detection

In pipeline_registry.py, add detection logic for your model’s _class_name in model_index.json.

After these steps, the framework automatically handles:

• Weight loading with optional dynamic quantization via PipelineLoader

• Multi-GPU execution via DiffusionExecutor

• Cache acceleration (if you call self._setup_cache_acceleration(self.transformer, coefficients=...) in post_load_weights(); supports both TeaCache and Cache-DiT via VisualGenArgs.cache_config)

• Serving via trtllm-serve with the full endpoint set