[None][perf] Reordering torch.compile and cache-dit call

[BrianLi23]

Summary

• Move cache-acceleration enablement (TeaCache / Cache-DiT) out of each pipeline's post_load_weights into PipelineLoader.load, so it runs after torch.compile; Cache-DiT captures block-module references at enable time, so enabling it first left it driving stale eager blocks and torch.compile contributed nothing.
• Centralize the call in the loader for all pipelines (Flux, Flux2, Wan T2V/I2V, LTX-2), leaving TeaCache coefficient resolution/validation in post_load_weights.
• Skip cache_dit's inspect-based forward-pattern and output-count checks when transformer blocks are torch._dynamo.OptimizedModule wrappers, since compiled (*args, **kwargs) forwards can't match the declared patterns.
• Update TeaCache unit tests to invoke _setup_cache_acceleration explicitly, mirroring the new loader-driven ordering.

Test Coverage

Model	Mean Before (s)	Mean After (s)	Std Before (s)	Std After (s)
Flux2	3.4827	3.2402	0.0183	0.0178
Qwen-Image	1.7481	1.5432	0.0093	0.0048
WAN2.2	15.8647	11.8930	0.0278	0.0112
LTX2	2.7206	2.7227	0.0148	0.0027

Note: the Qwen-Image results were measured with the Cache-DiT support from #16339 applied on top of this change, since Qwen-Image caching is not yet in main.

Flux2 Before

Flux2 After

Qwen-Image Before

Qwen-Image After

WAN2.2 Before

wan22-unpatched.mp4

WAN2.2 After
https://github.com/user-attachments/assets/56498c2b-f5e8-4eed-883e-b563257d6d05

Follow up

• LTX2 shows ~2.03s denoise for pre/post changes. Cache-DiT works fine in both runs (6/20 steps cached), but torch.compile seems to be broken on LTX2 as transformer_ltx2.py:1033 (v_kv_list[idx]) guards on the static block index which causes dynamo to hit the 128-recompile limit which falls back to eager. Follow-up: fix the LTX2 index guard so compile survives; until then LTX2 gets zero compile benefit even without caching.

• Reapply these changes on top of [None][feat] Qwen-Image TeaCache/Cache-DiT Support #16339 to reflect changes on Qwen Image.

Summary by CodeRabbit

• Bug Fixes

  • Improved TeaCache and Cache-DiT compatibility with compiled transformer blocks.
  • Ensured cache acceleration is initialized in the correct order after model compilation.
  • Improved validation handling for compiled blocks and multi-transformer pipelines.

• Tests

  • Updated TeaCache coverage for LTX-2 and Wan 2.2 video pipelines.

[BrianLi23]

/bot help

[github-actions]

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[BrianLi23]

cc: @zhenhuaw-me

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Cache acceleration flow

Layer / File(s)	Summary
Defer cache setup from pipeline post-load hooks tensorrt_llm/_torch/visual_gen/models/{flux,ltx2,wan}/*	Pipeline post-load methods retain TeaCache coefficient registration or validation while removing direct cache acceleration setup.
Enable acceleration after compilation tensorrt_llm/_torch/visual_gen/pipeline_loader.py	PipelineLoader.load() invokes _setup_cache_acceleration() after torch.compile() when a transformer is present.
Adapt compiled-block validation tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py, tests/unittest/_torch/visual_gen/test_teacache.py	FLUX Cache-DiT detects compiled blocks and disables incompatible forward-pattern checks; TeaCache tests explicitly invoke the relocated setup step.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PipelineLoader
  participant TorchCompile
  participant Pipeline
  participant CacheDiTAdapter
  PipelineLoader->>TorchCompile: compile pipeline when enabled
  TorchCompile-->>PipelineLoader: return compiled transformer blocks
  PipelineLoader->>Pipeline: _setup_cache_acceleration()
  Pipeline->>CacheDiTAdapter: configure compiled blocks
  CacheDiTAdapter-->>Pipeline: disable forward-pattern checks

Loading

Suggested reviewers: o-stoner

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 12.50% 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
Title check	✅ Passed	The title is concise and accurately reflects the main change: reordering torch.compile with cache-DiT/TeaCache setup.
Description check	✅ Passed	The description includes a clear summary and test coverage table; the missing PR checklist section is non-critical.
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.

✨ 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

🤖 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 `@tests/unittest/_torch/visual_gen/test_teacache.py`:
- Line 365: Add loader-level test coverage for the compiled Cache-DiT path by
exercising PipelineLoader.load() rather than only
BasePipeline._setup_cache_acceleration(). Assert that torch compilation occurs
before cache acceleration setup, or add an equivalent compiled Cache-DiT case
while preserving the existing direct setup tests.

🪄 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: cbb8b9cb-1439-40e5-9bdd-801e5bfbeb41

📥 Commits

Reviewing files that changed from the base of the PR and between a2595c0 and d562d09.

📒 Files selected for processing (8)

• tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py
• tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py
• tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py
• tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py
• tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py
• tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tests/unittest/_torch/visual_gen/test_teacache.py

[zhenhuaw-me]

@luyiyun1021 could help to take a look?

[zhenhuaw-me]

LGTM. @luyiyun1021 @o-stoner could you also take a look? Thanks!

[zhenhuaw-me]

/bot run

[tensorrt-cicd]

PR_Github #59821 [ run ] triggered by Bot. Commit: d562d09 Link to invocation

[tensorrt-cicd]

PR_Github #59821 [ run ] completed with state FAILURE. Commit: d562d09
/LLM/main/L0_MergeRequest_PR pipeline #48230 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[luyiyun1021]

The new SUPPORTED_CACHE_BACKENDS gate logs a warning and returns when the configured backend isn't in the pipeline's tuple. Before this PR, LTX2TwoStagesPipeline + cache_backend=cache_dit inherited LTX2Pipeline.post_load_weights's cache_dit branch and raised ValueError("Cache-DiT: no enabler registered for pipeline class 'LTX2TwoStagesPipeline'..."). Now the run proceeds at full denoise cost with one log warning. Do you think this works fine? @zhenhuaw-me

[luyiyun1021]

Leave LTX enable cache dit with torch compile as a followup

[zhenhuaw-me]

The new SUPPORTED_CACHE_BACKENDS gate logs a warning and returns when the configured backend isn't in the pipeline's tuple. Before this PR, LTX2TwoStagesPipeline + cache_backend=cache_dit inherited LTX2Pipeline.post_load_weights's cache_dit branch and raised ValueError("Cache-DiT: no enabler registered for pipeline class 'LTX2TwoStagesPipeline'..."). Now the run proceeds at full denoise cost with one log warning. Do you think this works fine? @zhenhuaw-me

I think we should explicitly raise exceptions

[luyiyun1021]

Suggested test additions (all in tests/unittest/_torch/visual_gen/test_cache_dit.py

1. Compile-on integration test (GPU) — in TestCacheDiTRealPipelineForward, add a Wan 2.1 1.3B variant with torch_compile.enable=True. Assert after PipelineLoader.load: the blocks cache_dit holds are torch._dynamo.OptimizedModule, and after one forward, dynamo counters show compiled frames actually ran.
   Why: nothing currently locks the enable-after-compile ordering — every existing cache_dit test runs with compile off, so reverting the order keeps CI green, and the regression is perf-only (~25% on Wan) and invisible to correctness assertions. The counters check is needed because OptimizedModule blocks can still silently fall back to eager (LTX-2 does today).

2. Flux enabler check-flags unit test (CPU) — new test class. Call enable_cache_dit_for_flux with blocks wrapped by real torch.compile (lazy, no actual compilation) vs. plain eager blocks, patch BlockAdapter to capture kwargs, and assert check_forward_pattern/check_num_outputs are False for compiled and True for eager.
   Why: compiled Flux + cache_dit is a supported config (benchmarked in this PR) with zero test coverage; if _has_compiled_blocks breaks or the flags are inverted, the server crashes at startup on cache_dit's inspect assert while CI stays green.

3. Registry ↔ declaration consistency test (CPU) — new test class. Assert bidirectionally that every pipeline class in CUSTOM_CACHE_DIT_ENABLERS declares "cache_dit" in SUPPORTED_CACHE_BACKENDS, and vice versa.
   Why: the new opt-in gate turns a missing declaration into a log warning instead of an error, so a pipeline wired to an enabler but missing the declaration silently loses cache acceleration. The Qwen-Image reapply ([None][feat] Qwen-Image TeaCache/Cache-DiT Support #16339) is the first realistic case.

[o-stoner]

Left a comment on perhaps enforcing explicit exception raising rather than logging, but otherwise LGTM, thanks!

[chang-l]

/bot run

[tensorrt-cicd]

PR_Github #61405 [ run ] triggered by Bot. Commit: 19c8e02 Link to invocation

[tensorrt-cicd]

PR_Github #61405 [ run ] completed with state FAILURE. Commit: 19c8e02
/LLM/main/L0_MergeRequest_PR pipeline #49632 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chang-l]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61417 [ run ] triggered by Bot. Commit: 341edf4 Link to invocation

[chang-l]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61432 [ run ] triggered by Bot. Commit: e3d43bd Link to invocation

[tensorrt-cicd]

PR_Github #61417 [ run ] completed with state ABORTED. Commit: 341edf4

Link to invocation

[luyiyun1021]

LGTM. Thanks.

[tensorrt-cicd]

PR_Github #61432 [ run ] completed with state SUCCESS. Commit: e3d43bd
/LLM/main/L0_MergeRequest_PR pipeline #49657 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chang-l]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61493 [ run ] triggered by Bot. Commit: ed0d8b6 Link to invocation