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

tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py

@@ -9,6 +9,7 @@
 from typing import Any, Callable, List, Optional
 
 import cache_dit
+import torch
 import torch.nn as nn
 from cache_dit import (
     BlockAdapter,
@@ -41,6 +42,20 @@ class CacheDiTEnableResult:
     summary_modules: List[nn.Module]
 
 
+def _has_compiled_blocks(*block_lists: Any) -> bool:
+    """True when any transformer block is a torch.compile OptimizedModule wrapper.
+
+    torch.compile'd blocks expose a ``(*args, **kwargs)`` forward signature, so
+    cache_dit's inspect-based forward-pattern checks cannot match them and must
+    be skipped (the underlying blocks still follow the declared pattern).
+    """
+    return any(
+        isinstance(block, torch._dynamo.OptimizedModule)
+        for blocks in block_lists
+        for block in blocks
+    )
+
+
 def _resolved_enable_separate_cfg(cfg: CacheDiTConfig, default: bool) -> bool:
     if cfg.enable_separate_cfg is not None:
         return cfg.enable_separate_cfg
@@ -258,12 +273,18 @@ def enable_cache_dit_for_flux(
         forward_pattern = [ForwardPattern.Pattern_1, ForwardPattern.Pattern_1]
         tag = "FLUX.1"
 
+    compiled_blocks = _has_compiled_blocks(*block_lists)
+    if compiled_blocks:
+        logger.info(
+            f"Cache-DiT: {tag} blocks are torch.compile'd; skipping forward-pattern checks."
+        )
     adapter = BlockAdapter(
         transformer=pipeline.transformer,
         blocks=block_lists,
         forward_pattern=forward_pattern,
         params_modifiers=[modifier],
-        check_forward_pattern=True,
+        check_forward_pattern=not compiled_blocks,
+        check_num_outputs=not compiled_blocks,
     )
 
     logger.info(

tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py

@@ -242,9 +242,9 @@ def post_load_weights(self) -> None:
                 )
             )
 
-            # TeaCache or Cache-DiT
+            # TeaCache or Cache-DiT: resolve coefficients here; the loader enables
+            # cache acceleration after torch.compile (see PipelineLoader.load).
             self._apply_teacache_coefficients(FLUX_TEACACHE_COEFFICIENTS)
-            self._setup_cache_acceleration()
 
     @property
     def default_generation_params(self):

tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py

@@ -336,8 +336,9 @@ def post_load_weights(self) -> None:
                     )
                 )
 
+            # Cache acceleration itself is enabled by the loader after
+            # torch.compile (see PipelineLoader.load).
             self._apply_teacache_coefficients(FLUX2_TEACACHE_COEFFICIENTS)
-            self._setup_cache_acceleration()
 
     @property
     def default_generation_params(self):

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py

@@ -1026,7 +1026,9 @@ def post_load_weights(self) -> None:
         """Finalize after weight loading: TeaCache, Cache-DiT, derived attributes."""
         super().post_load_weights()
 
-        # LTX-2: single transformer (one DiT for video+audio); TeaCache only with explicit coefficients.
+        # LTX-2: single transformer (one DiT for video+audio); TeaCache only with
+        # explicit coefficients. Cache acceleration itself is enabled by the
+        # loader after torch.compile (see PipelineLoader.load).
         if self.transformer is not None and self.pipeline_config.cache_backend == "teacache":
             if self.pipeline_config.teacache.coefficients is None:
                 raise ValueError(
@@ -1037,11 +1039,6 @@ def post_load_weights(self) -> None:
                 "LTXModel",
                 LTX2TeaCacheExtractor(self._compute_ltx2_timestep_embedding),
             )
-            self._setup_cache_acceleration()
-
-        # Cache-DiT
-        if self.transformer is not None and self.pipeline_config.cache_backend == "cache_dit":
-            self._setup_cache_acceleration()
 
         # Compression ratios from native scale factors
         self.vae_spatial_compression_ratio = VIDEO_SCALE_FACTORS.width

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py

@@ -345,16 +345,14 @@ def post_load_weights(self) -> None:
 
             if not self.is_wan22_14b:
                 self._apply_teacache_coefficients(WAN_TEACACHE_COEFFICIENTS)
-                self._setup_cache_acceleration()
-            else:
-                if self.pipeline_config.cache_backend == "cache_dit":
-                    self._setup_cache_acceleration()
 
         if self.transformer_2 is not None:
             if hasattr(self.transformer_2, "post_load_weights"):
                 self.transformer_2.post_load_weights()
 
-        # Wan 2.2 TeaCache after both transformers' post_load_weights (FP8 scales, etc.)
+        # Wan 2.2 TeaCache validation after both transformers' post_load_weights
+        # (FP8 scales, etc.). Cache acceleration itself is enabled by the loader
+        # after torch.compile (see PipelineLoader.load).
         if (
             self.transformer is not None
             and self.transformer_2 is not None
@@ -367,7 +365,6 @@ def post_load_weights(self) -> None:
                     "teacache.coefficients_2 (high-noise and low-noise stage polynomials). "
                     "There is no built-in coefficient table for Wan 2.2."
                 )
-            self._setup_cache_acceleration()
 
     def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None:
         with torch.no_grad():

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py

@@ -356,15 +356,13 @@ def post_load_weights(self) -> None:
 
             if not self.is_wan22_14b:
                 self._apply_teacache_coefficients(WAN_I2V_TEACACHE_COEFFICIENTS)
-                self._setup_cache_acceleration()
-            else:
-                if self.pipeline_config.cache_backend == "cache_dit":
-                    self._setup_cache_acceleration()
 
         if self.transformer_2 is not None:
             if hasattr(self.transformer_2, "post_load_weights"):
                 self.transformer_2.post_load_weights()
 
+        # Wan 2.2 TeaCache validation; cache acceleration itself is enabled by
+        # the loader after torch.compile (see PipelineLoader.load).
         if (
             self.transformer is not None
             and self.transformer_2 is not None
@@ -377,7 +375,6 @@ def post_load_weights(self) -> None:
                     "teacache.coefficients_2 (high-noise and low-noise stage polynomials). "
                     "There is no built-in coefficient table for Wan 2.2."
                 )
-            self._setup_cache_acceleration()
 
     def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None:
         dummy_image = PIL.Image.new("RGB", (width, height))

tensorrt_llm/_torch/visual_gen/pipeline_loader.py

@@ -303,6 +303,15 @@ def load(
         else:
             logger.info("torch.compile disabled by config")
 
+        # Cache acceleration (TeaCache / Cache-DiT) is enabled AFTER torch.compile
+        # on purpose: Cache-DiT captures references to the transformer block
+        # modules at enable time, while torch_compile() replaces the block lists
+        # with compiled copies. If Cache-DiT were enabled first, it would keep
+        # running the stale eager blocks and torch.compile would contribute
+        # nothing.
+        if getattr(pipeline, "transformer", None) is not None:
+            pipeline._setup_cache_acceleration()
+
         if not skip_warmup:
             if config.torch_compile.enable_autotune:
                 with autotune(

tests/unittest/_torch/visual_gen/test_cache_dit.py

@@ -15,6 +15,7 @@
 import logging
 import os
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -157,6 +158,40 @@ def _total_accumulated_cached_steps(stats: dict) -> int:
     return total
 
 
+def _cache_dit_captured_blocks(transformer: torch.nn.Module) -> list[torch.nn.Module]:
+    """Block modules cache_dit captured at enable time.
+
+    cache_dit replaces ``transformer.forward`` with a ``functools.partial``
+    whose closure holds the UnifiedBlocks mapping; each cached-blocks module
+    keeps the block list it captured on ``.transformer_blocks``. Walk the
+    closure graph to collect those blocks so tests can assert on the exact
+    module references cache_dit drives at denoise time.
+    """
+    fwd = transformer.forward
+    stack = [getattr(fwd, "func", fwd)]
+    seen: set[int] = set()
+    captured: list[torch.nn.Module] = []
+    while stack:
+        obj = stack.pop()
+        if id(obj) in seen:
+            continue
+        seen.add(id(obj))
+        closure = getattr(obj, "__closure__", None)
+        if closure:
+            for cell in closure:
+                try:
+                    stack.append(cell.cell_contents)
+                except ValueError:
+                    pass  # empty cell
+        elif isinstance(obj, dict):
+            stack.extend(obj.values())
+        elif isinstance(obj, torch.nn.ModuleList):
+            stack.extend(obj)
+        elif isinstance(obj, torch.nn.Module) and hasattr(obj, "transformer_blocks"):
+            captured.extend(obj.transformer_blocks)
+    return captured
+
+
 # ---------------------------------------------------------------------------
 # 1) Wan 2.2 high/low step split (no mocks)
 # ---------------------------------------------------------------------------
@@ -239,7 +274,12 @@ def _teardown_cache_dit(pipeline: object) -> None:
             pipeline.cache_accelerator = None
 
     @staticmethod
-    def _load_visual_gen_pipeline(checkpoint_dir: str, *, text_encoder_path: str = ""):
+    def _load_visual_gen_pipeline(
+        checkpoint_dir: str,
+        *,
+        text_encoder_path: str = "",
+        enable_torch_compile: bool = False,
+    ):
         from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader
         from tensorrt_llm.visual_gen.args import (
             CacheDiTConfig,
@@ -267,7 +307,7 @@ def _load_visual_gen_pipeline(checkpoint_dir: str, *, text_encoder_path: str = "
                 residual_diff_threshold=0.25,
             ),
             torch_compile_config=TorchCompileConfig(
-                enable=False,
+                enable=enable_torch_compile,
                 enable_autotune=False,
             ),
             compilation_config=CompilationConfig(skip_warmup=True),
@@ -314,6 +354,66 @@ def test_wan_cache_dit_skips_blocks_after_forward(self):
                 if pipeline is not None:
                     self._teardown_cache_dit(pipeline)
 
+    def test_wan_cache_dit_enabled_after_torch_compile(self):
+        """Locks the enable-after-compile ordering from PipelineLoader.load.
+
+        Cache-DiT captures block-module references at enable time. If it were
+        enabled before torch_compile() (which replaces the block lists with
+        compiled copies), it would keep driving the stale eager blocks and
+        torch.compile would contribute nothing — a perf-only regression that
+        is invisible to the correctness assertions of the compile-off tests.
+        """
+        ckpt = _resolve_wan_checkpoint()
+        if ckpt is None:
+            pytest.skip(
+                "Wan 2.1 1.3B not found: set TRTLLM_CACHE_DIT_WAN_CHECKPOINT, "
+                f"install under {_DEFAULT_WAN_CHECKPOINT}, "
+                f"or under $LLM_MODELS_ROOT/{_WAN_SUBPATH}"
+            )
+
+        pipeline = None
+        with _suppress_stdlib_logging_for_cache_dit():
+            try:
+                pipeline = self._load_visual_gen_pipeline(ckpt, enable_torch_compile=True)
+                assert pipeline.cache_accelerator is not None
+                assert pipeline.cache_accelerator.is_enabled()
+
+                # The blocks cache_dit holds must be the torch.compile wrappers,
+                # and exactly the ones currently installed on the transformer —
+                # not the eager blocks that torch_compile() replaced.
+                captured = _cache_dit_captured_blocks(pipeline.transformer)
+                assert captured, "no cache_dit-captured blocks found on the patched forward"
+                assert all(isinstance(b, torch._dynamo.OptimizedModule) for b in captured), (
+                    "cache_dit captured eager blocks; it was enabled before torch.compile"
+                )
+                assert {id(b) for b in captured} == {id(b) for b in pipeline.transformer.blocks}
+
+                # OptimizedModule blocks can still silently fall back to eager
+                # (LTX-2 does today): require dynamo to report compiled frames
+                # after a real forward.
+                counters = torch._dynamo.utils.counters
+                counters.clear()
+                with torch.inference_mode():
+                    pipeline.forward(
+                        prompt="cache dit compile ordering validation",
+                        negative_prompt="",
+                        height=480,
+                        width=832,
+                        num_frames=33,
+                        num_inference_steps=8,
+                        guidance_scale=5.0,
+                        seed=0,
+                        max_sequence_length=256,
+                    )
+                assert counters["frames"]["ok"] > 0, (
+                    "torch.compile produced no compiled frames during forward; "
+                    "compiled blocks fell back to eager"
+                )
+            finally:
+                if pipeline is not None:
+                    self._teardown_cache_dit(pipeline)
+                torch._dynamo.reset()
+
     def test_flux_cache_dit_skips_blocks_after_forward(self):
         ckpt = _resolve_flux_checkpoint()
         if ckpt is None:
@@ -410,3 +510,90 @@ def test_ltx2_cache_dit_skips_blocks_after_forward(self):
             finally:
                 if pipeline is not None:
                     self._teardown_cache_dit(pipeline)
+
+
+# ---------------------------------------------------------------------------
+# 3) FLUX enabler check flags for compiled vs. eager blocks (CPU)
+# ---------------------------------------------------------------------------
+
+
+class _IdentityBlock(torch.nn.Module):
+    def forward(self, hidden_states, encoder_hidden_states):
+        return hidden_states, encoder_hidden_states
+
+
+@requires_cache_dit
+class TestFluxEnablerCompiledBlockCheckFlags:
+    """enable_cache_dit_for_flux must disable cache_dit's inspect-based
+    forward-pattern checks exactly when blocks are torch.compile wrappers.
+
+    Compiled blocks expose a ``(*args, **kwargs)`` forward, so cache_dit's
+    checks would assert at enable time (server startup) if left on; eager
+    blocks must keep them on. torch.compile here is lazy — no compilation
+    happens, so this runs on CPU.
+    """
+
+    @staticmethod
+    def _make_flux_like_pipeline(*, compiled: bool) -> SimpleNamespace:
+        def _blocks() -> torch.nn.ModuleList:
+            blocks = [_IdentityBlock(), _IdentityBlock()]
+            if compiled:
+                blocks = [torch.compile(b) for b in blocks]
+            return torch.nn.ModuleList(blocks)
+
+        return SimpleNamespace(
+            transformer=SimpleNamespace(
+                transformer_blocks=_blocks(),
+                single_transformer_blocks=_blocks(),
+            )
+        )
+
+    def _block_adapter_kwargs(self, *, compiled: bool, is_flux2: bool) -> dict:
+        from tensorrt_llm._torch.visual_gen.cache import cache_dit_enablers
+        from tensorrt_llm.visual_gen.args import CacheDiTConfig
+
+        pipeline = self._make_flux_like_pipeline(compiled=compiled)
+        with (
+            patch.object(cache_dit_enablers, "BlockAdapter") as adapter_cls,
+            patch.object(cache_dit_enablers, "cache_dit"),
+        ):
+            cache_dit_enablers.enable_cache_dit_for_flux(
+                pipeline, CacheDiTConfig(), is_flux2=is_flux2
+            )
+        adapter_cls.assert_called_once()
+        return adapter_cls.call_args.kwargs
+
+    @pytest.mark.parametrize("is_flux2", [False, True])
+    def test_compiled_blocks_disable_pattern_checks(self, is_flux2):
+        kwargs = self._block_adapter_kwargs(compiled=True, is_flux2=is_flux2)
+        assert kwargs["check_forward_pattern"] is False
+        assert kwargs["check_num_outputs"] is False
+
+    @pytest.mark.parametrize("is_flux2", [False, True])
+    def test_eager_blocks_keep_pattern_checks(self, is_flux2):
+        kwargs = self._block_adapter_kwargs(compiled=False, is_flux2=is_flux2)
+        assert kwargs["check_forward_pattern"] is True
+        assert kwargs["check_num_outputs"] is True
+
+
+# ---------------------------------------------------------------------------
+# 4) Enabler registry declaration consistency (CPU)
+# ---------------------------------------------------------------------------
+
+
+@requires_cache_dit
+class TestCacheDiTEnablerRegistry:
+    def test_enabler_keys_name_registered_pipeline_classes(self):
+        """Every CUSTOM_CACHE_DIT_ENABLERS key must name a real registered
+        pipeline class; a typo or rename would otherwise turn every enable
+        attempt for that pipeline into the 'no enabler registered' error."""
+        import tensorrt_llm._torch.visual_gen.models  # noqa: F401  (populates the registry)
+        from tensorrt_llm._torch.visual_gen.cache.cache_dit_enablers import (
+            CUSTOM_CACHE_DIT_ENABLERS,
+        )
+        from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY
+
+        unknown = sorted(set(CUSTOM_CACHE_DIT_ENABLERS) - set(PIPELINE_REGISTRY))
+        assert not unknown, (
+            f"CUSTOM_CACHE_DIT_ENABLERS keys not in the pipeline registry: {unknown}"
+        )