[None][feat] Add HunyuanDiT text-to-image pipeline with Ulysses parallelism

[pkisfaludi-nv]

@coderabbitai summary

Summary

• New pipeline: HunyuanDiTPipeline for Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers (and v1.1, v1.0), following the existing BasePipeline pattern
• Bilingual text encoding: BERT-based CLIP encoder (text_encoder) + MT5EncoderModel (text_encoder_2), both tokenizers, concatenated for cross-attention
• Ulysses sequence parallelism: custom attention processor injects all_to_all_4d around self-attention; cross-attention uses standard SDPA (text replicated); image_rotary_emb (2D RoPE) sliced per rank; all_gather after the final block
• Resolution binning: aspect-ratio snapping to training bucket resolutions via log-space distance
• Config: examples/visual_gen/serve/configs/hunyuandit.yml
• Docs: supported models table + feature matrix (14-column, Cache-DiT column) + footnote in visual-generation.md

Architecture notes

HunyuanDiT uses a diffusers HunyuanDiT2DModel backbone with U-Net skip connections. Because the model uses diffusers' HunyuanAttnProcessor2_0 rather than TRT-LLM's Attention module, Ulysses is implemented via:

1. A custom HunyuanDiTUlyssesAttnProcessor replacing the self-attention processor in each block
2. A patched forward() (via types.MethodType) that shards the latent sequence after patch-embed and gathers before the final norm/proj

Constraint: num_attention_heads=16 must be divisible by ulysses_size.

Files changed

tensorrt_llm/_torch/visual_gen/models/hunyuandit/   (new)
  __init__.py
  defaults.py
  transformer_hunyuandit.py
  pipeline_hunyuandit.py
tensorrt_llm/_torch/visual_gen/models/__init__.py   (modified)
tensorrt_llm/_torch/visual_gen/pipeline_registry.py (modified)
examples/visual_gen/serve/configs/hunyuandit.yml    (new)
docs/source/models/visual-generation.md             (modified)

Test plan

• Import check: python -c "from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY; assert 'HunyuanDiTPipeline' in PIPELINE_REGISTRY"
• Single-GPU inference: trtllm-serve visual_gen --model Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers
• Multi-GPU Ulysses: torchrun --nproc-per-node=2 ... ulysses_size: 2
• trtllm-serve /v1/images/generations endpoint

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (HunyuanDiT CI tests not yet in test lists; will be added in a follow-up PR)

• No API changes introduced — label not needed.

• No new external dependencies beyond diffusers (already a project dependency).

• CODEOWNERS not affected.

• Documentation updated: docs/source/models/visual-generation.md updated with model table entry, feature matrix row, and footnote.

• No significant design change; tava diagram not affected.

• Reviewers assigned automatically via CODEOWNERS.

• Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds HunyuanDiT text-to-image generation to TensorRT-LLM's VisualGen module, including dual-language prompt conditioning, resolution binning, and optional distributed sequence parallelism via Ulysses, with supporting documentation and checkpoint detection.

Changes

HunyuanDiT VisualGen Implementation

Layer / File(s)	Summary
Module structure and exports tensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.py, tensorrt_llm/_torch/visual_gen/models/__init__.py	Creates the hunyuandit module package with SPDX headers and exports HunyuanDiTPipeline and HunyuanDiT2DModelWrapper for public consumption; updates parent models __init__.py to expose the new pipeline.
Configuration and generation defaults examples/visual_gen/serve/configs/hunyuandit.yml, tensorrt_llm/_torch/visual_gen/models/hunyuandit/defaults.py	Defines YAML configuration with VANILLA attention backend and parallel config (cfg_size=1, ulysses_size=1); provides default generation parameters (1024×1024 resolution, 50 steps, 7.5 guidance, 77 max tokens) and extra parameter specs for negative prompts and resolution binning toggle.
Transformer wrapper with Ulysses parallelism tensorrt_llm/_torch/visual_gen/models/hunyuandit/transformer_hunyuandit.py	Implements HunyuanDiT2DModelWrapper wrapping the diffusers transformer with conditional Ulysses sequence parallelism: HunyuanDiTUlyssesAttnProcessor injects all-to-all sharding+gathering for self-attention via SDPA, HunyuanDiT2DModelUlysses patches forward to shard/gather image token sequences across ranks with adjusted RoPE frequencies, wrapper validates divisibility constraints, conditionally applies Ulysses patch, loads weights with logging, and converts to inference dtype.
Pipeline generation flow tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py	Implements HunyuanDiTPipeline with resolution aspect-ratio binning to snap dimensions to training buckets, loads dual text encoders (BERT for ≤77 tokens, MT5 for ≤256 tokens), VAE, and DDPM scheduler, supports classifier-free guidance via negative-prompt encoding, prepares latents from VAE downsampling, computes 2D RoPE embeddings, executes denoising loop with transformer noise prediction and CFG, and returns generated uint8 images with rank-0 timing logs.
Documentation and checkpoint detection docs/source/models/visual-generation.md, tensorrt_llm/_torch/visual_gen/pipeline_registry.py	Adds HunyuanDiT v1.2/v1.1/v1.0 checkpoint IDs to supported models table and feature matrix with footnote explaining bilingual text conditioning and Ulysses parallelism constraints; extends AutoPipeline._detect_from_checkpoint() to recognize and dispatch HunyuanDiT model class names.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested labels

VisualGen

Suggested reviewers

• Shixiaowei02
• zhenhuaw-me
• yibinl-nvidia
• NVShreyas
• chang-l

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and specifically describes the main addition: a new HunyuanDiT text-to-image pipeline with Ulysses parallelism support, which directly aligns with the changeset.
Description check	✅ Passed	The PR description covers the change summary, architecture notes, files, test plan, and checklist, with only minor template-heading differences.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py (1)

507-529: ⚡ Quick win

Unused loop variable i.

The loop counter i is declared but never used in the body. Rename to _ per PEP 8 and Ruff B007.

♻️ Proposed fix

-        for i, t in enumerate(timesteps):
+        for _, t in enumerate(timesteps):
             lat_in = torch.cat([latents] * 2) if do_cfg else latents

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py`
around lines 507 - 529, The for-loop in pipeline_hunyuandit.py currently
declares an unused loop variable "i" (for i, t in enumerate(timesteps)) which
triggers Ruff B007; change the loop to use an underscore for the unused index
(for _, t in enumerate(timesteps)) or drop enumerate if not needed, updating the
loop that constructs lat_in, calls self.transformer(...), performs CFG handling
and calls self.scheduler.step so only used variables remain.

tensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.py (1)

9-12: ⚡ Quick win

Sort __all__ entries alphabetically.

As per coding guidelines, maintain sorted __all__ lists for consistency across the codebase.

📋 Proposed fix

 __all__ = [
-    "HunyuanDiTPipeline",
     "HunyuanDiT2DModelWrapper",
+    "HunyuanDiTPipeline",
 ]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.py` around lines 9
- 12, The __all__ list in the module exports is unsorted; reorder the entries
alphabetically so exports are consistent across the codebase — change the list
containing "HunyuanDiTPipeline" and "HunyuanDiT2DModelWrapper" to alphabetical
order (place "HunyuanDiT2DModelWrapper" before "HunyuanDiTPipeline") while
keeping the same string names and no other modifications.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py`:
- Around line 338-373: The _get_image_rotary_emb function currently returns None
when diffusers' get_2d_rotary_pos_embed is unavailable but its signature is
Tuple[torch.Tensor, torch.Tensor]; change the return annotation to
Optional[Tuple[torch.Tensor, torch.Tensor]] (import Optional from typing) and
remove the type: ignore on the None return so the signature matches callers
(e.g., where image_rotary_emb is checked for None); update any docstring/type
comments accordingly so static checkers understand None is a valid return.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.py`:
- Around line 9-12: The __all__ list in the module exports is unsorted; reorder
the entries alphabetically so exports are consistent across the codebase —
change the list containing "HunyuanDiTPipeline" and "HunyuanDiT2DModelWrapper"
to alphabetical order (place "HunyuanDiT2DModelWrapper" before
"HunyuanDiTPipeline") while keeping the same string names and no other
modifications.

In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py`:
- Around line 507-529: The for-loop in pipeline_hunyuandit.py currently declares
an unused loop variable "i" (for i, t in enumerate(timesteps)) which triggers
Ruff B007; change the loop to use an underscore for the unused index (for _, t
in enumerate(timesteps)) or drop enumerate if not needed, updating the loop that
constructs lat_in, calls self.transformer(...), performs CFG handling and calls
self.scheduler.step so only used variables remain.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9bd57c0d-3159-4d02-abd7-66cbeb72038c

📥 Commits

Reviewing files that changed from the base of the PR and between 81e86a5 and 33b4360.

📒 Files selected for processing (8)

• docs/source/models/visual-generation.md
• examples/visual_gen/serve/configs/hunyuandit.yml
• tensorrt_llm/_torch/visual_gen/models/__init__.py
• tensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.py
• tensorrt_llm/_torch/visual_gen/models/hunyuandit/defaults.py
• tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py
• tensorrt_llm/_torch/visual_gen/models/hunyuandit/transformer_hunyuandit.py
• tensorrt_llm/_torch/visual_gen/pipeline_registry.py