From 40a7b4c329182925faf79adb9e5b576ae42fceab Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 16 Jul 2026 09:24:46 -0700 Subject: [PATCH 1/7] reorder cache-dit and torch.compile application order Signed-off-by: BrianLi23 --- .../visual_gen/cache/cache_dit_enablers.py | 42 ++++++++++++++++++- .../visual_gen/models/flux/pipeline_flux.py | 4 +- .../visual_gen/models/flux/pipeline_flux2.py | 3 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 9 ++-- .../visual_gen/models/wan/pipeline_wan.py | 9 ++-- .../visual_gen/models/wan/pipeline_wan_i2v.py | 7 +--- .../_torch/visual_gen/pipeline_loader.py | 9 ++++ .../_torch/visual_gen/test_teacache.py | 8 ++++ 8 files changed, 70 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py b/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py index a859e50d70a8..ed326bb987b7 100644 --- a/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py +++ b/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,38 @@ 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 _clear_stale_fake_pipeline_cached_flag() -> None: + """Clear cache_dit's leaked class-level FakeDiffusionPipeline._is_cached flag. + + CachedAdapter.create_context sets ``pipe.__class__._is_cached = True``. For + transformer-only / pipe-less adapters the pipe class is FakeDiffusionPipeline, + and cache_dit.disable_cache(transformer) never clears the class-level flag. + A later re-enable (e.g. re-wrap after torch.compile) then sees its brand-new + FakeDiffusionPipeline as "already cached", skips cache-context creation, and + asserts in mock_blocks. + """ + try: + from cache_dit.caching.block_adapters.block_adapters import FakeDiffusionPipeline + except ImportError: + return + if "_is_cached" in FakeDiffusionPipeline.__dict__: + del FakeDiffusionPipeline._is_cached + + 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 +291,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( @@ -355,6 +394,7 @@ def enable_cache_dit_for_pipeline( pipeline: Any, cache_dit_cfg: CacheDiTConfig ) -> CacheDiTEnableResult: """Dispatch to a registered enabler or raise.""" + _clear_stale_fake_pipeline_cached_flag() name = pipeline.__class__.__name__ if name not in CUSTOM_CACHE_DIT_ENABLERS: raise ValueError( diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index 0a35ce659122..e08d07047878 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/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): diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 4f0fa7391895..1b2116f7f27f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/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): diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 1336bb7249bc..45a46ea3b015 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/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 diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 41a284e040c3..46e69baf24d8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -331,16 +331,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 @@ -353,7 +351,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(): diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index f472c4c1fe45..47138e59ed51 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -342,15 +342,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 @@ -363,7 +361,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)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index 3551371fa173..884f90f97d45 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/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( diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index c53f77aaa135..8335687ec186 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -362,6 +362,8 @@ def test_ltx2_succeeds_with_explicit_coefficients(self): ): with patch.object(TeaCacheBackend, "enable"): LTX2Pipeline.post_load_weights(pipe) + # The loader enables cache acceleration after torch.compile. + BasePipeline._setup_cache_acceleration(pipe) assert pipe.cache_accelerator is not None def test_wan22_raises_when_teacache_enabled_without_both_coefficient_lists(self): @@ -449,6 +451,8 @@ def test_wan22_t2v_installs_two_teacache_backends_when_coefficients_provided(sel ) as TB: TB.side_effect = [backend_a, backend_b] WanPipeline.post_load_weights(pipe) + # The loader enables cache acceleration after torch.compile. + BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 assert mock_enable.call_count == 2 assert pipe.cache_accelerator is not None @@ -480,6 +484,8 @@ def test_wan22_t2v_transformer_gets_coefficients_and_transformer_2_gets_coeffici ) as TB: TB.return_value = MagicMock() WanPipeline.post_load_weights(pipe) + # The loader enables cache acceleration after torch.compile. + BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 cfg_high = TB.call_args_list[0][0][0] @@ -515,6 +521,8 @@ def test_wan22_i2v_transformer_gets_coefficients_and_transformer_2_gets_coeffici ) as TB: TB.return_value = MagicMock() WanImageToVideoPipeline.post_load_weights(pipe) + # The loader enables cache acceleration after torch.compile. + BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 cfg_high = TB.call_args_list[0][0][0] From f4dd43276d60f1a9957ef7a02af0af4622e29adf Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 16 Jul 2026 11:34:18 -0700 Subject: [PATCH 2/7] remove comments Signed-off-by: BrianLi23 --- tests/unittest/_torch/visual_gen/test_teacache.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index 8335687ec186..6deeb6b6bfab 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -362,7 +362,6 @@ def test_ltx2_succeeds_with_explicit_coefficients(self): ): with patch.object(TeaCacheBackend, "enable"): LTX2Pipeline.post_load_weights(pipe) - # The loader enables cache acceleration after torch.compile. BasePipeline._setup_cache_acceleration(pipe) assert pipe.cache_accelerator is not None @@ -451,7 +450,6 @@ def test_wan22_t2v_installs_two_teacache_backends_when_coefficients_provided(sel ) as TB: TB.side_effect = [backend_a, backend_b] WanPipeline.post_load_weights(pipe) - # The loader enables cache acceleration after torch.compile. BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 assert mock_enable.call_count == 2 @@ -484,7 +482,6 @@ def test_wan22_t2v_transformer_gets_coefficients_and_transformer_2_gets_coeffici ) as TB: TB.return_value = MagicMock() WanPipeline.post_load_weights(pipe) - # The loader enables cache acceleration after torch.compile. BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 @@ -521,7 +518,6 @@ def test_wan22_i2v_transformer_gets_coefficients_and_transformer_2_gets_coeffici ) as TB: TB.return_value = MagicMock() WanImageToVideoPipeline.post_load_weights(pipe) - # The loader enables cache acceleration after torch.compile. BasePipeline._setup_cache_acceleration(pipe) assert TB.call_count == 2 From aa9ce3e0e093550417abe478dfa774542d294eb0 Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 16 Jul 2026 12:40:17 -0700 Subject: [PATCH 3/7] remove stale test Signed-off-by: BrianLi23 --- .../visual_gen/cache/cache_dit_enablers.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py b/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py index ed326bb987b7..4f4c4c267711 100644 --- a/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py +++ b/tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py @@ -56,24 +56,6 @@ def _has_compiled_blocks(*block_lists: Any) -> bool: ) -def _clear_stale_fake_pipeline_cached_flag() -> None: - """Clear cache_dit's leaked class-level FakeDiffusionPipeline._is_cached flag. - - CachedAdapter.create_context sets ``pipe.__class__._is_cached = True``. For - transformer-only / pipe-less adapters the pipe class is FakeDiffusionPipeline, - and cache_dit.disable_cache(transformer) never clears the class-level flag. - A later re-enable (e.g. re-wrap after torch.compile) then sees its brand-new - FakeDiffusionPipeline as "already cached", skips cache-context creation, and - asserts in mock_blocks. - """ - try: - from cache_dit.caching.block_adapters.block_adapters import FakeDiffusionPipeline - except ImportError: - return - if "_is_cached" in FakeDiffusionPipeline.__dict__: - del FakeDiffusionPipeline._is_cached - - def _resolved_enable_separate_cfg(cfg: CacheDiTConfig, default: bool) -> bool: if cfg.enable_separate_cfg is not None: return cfg.enable_separate_cfg @@ -394,7 +376,6 @@ def enable_cache_dit_for_pipeline( pipeline: Any, cache_dit_cfg: CacheDiTConfig ) -> CacheDiTEnableResult: """Dispatch to a registered enabler or raise.""" - _clear_stale_fake_pipeline_cached_flag() name = pipeline.__class__.__name__ if name not in CUSTOM_CACHE_DIT_ENABLERS: raise ValueError( From 19c8e0225b89e6af5bf000b16e7178f7703f6f2d Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 16 Jul 2026 15:43:49 -0700 Subject: [PATCH 4/7] handle cache support per pipeline Signed-off-by: BrianLi23 --- .../visual_gen/models/flux/pipeline_flux.py | 2 ++ .../visual_gen/models/flux/pipeline_flux2.py | 2 ++ .../visual_gen/models/ltx2/pipeline_ltx2.py | 2 ++ .../models/ltx2/pipeline_ltx2_two_stages.py | 2 ++ .../visual_gen/models/wan/pipeline_wan.py | 2 ++ .../visual_gen/models/wan/pipeline_wan_i2v.py | 2 ++ tensorrt_llm/_torch/visual_gen/pipeline.py | 24 ++++++++++++++++--- .../_torch/visual_gen/test_teacache.py | 2 ++ 8 files changed, 35 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index e08d07047878..8e9a4b4ed121 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -50,6 +50,8 @@ class FluxPipeline(BasePipeline): Supports FLUX.1-dev (50 steps, guidance) and FLUX.1-schnell (4 steps, no guidance). """ + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + def __init__(self, pipeline_config): if ( pipeline_config.visual_gen_mapping is not None diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 1b2116f7f27f..0af916483af8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -105,6 +105,8 @@ class Flux2Pipeline(BasePipeline): Follows WAN pipeline pattern for DiffusionModelLoader integration. """ + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + # Hidden state layers per text encoder type (auto-detected at load time) _TEXT_ENCODER_CONFIG = { "Mistral3ForConditionalGeneration": { diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 45a46ea3b015..1e6446b57e1f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -620,6 +620,8 @@ class LTX2Pipeline(BasePipeline): ``transformers`` library. """ + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + def __init__(self, model_config): _sa_cfg = model_config.attention.sparse_attention_config if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa": diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index b7ef703be22d..c7ed27ab189d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -934,6 +934,8 @@ class LTX2TwoStagesPipeline(LTX2Pipeline): then decode. """ + SUPPORTED_CACHE_BACKENDS = ("teacache",) + @property def common_warmup_shapes(self) -> list: return [(512, 768, 121)] diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 8230f10823e5..5dea9408f09c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -112,6 +112,8 @@ doc="Wan 2.1 & 2.2 text-to-video family.", ) class WanPipeline(BasePipeline): + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + def __init__(self, pipeline_config): # Wan2.2 A14B two-stage denoising parameters self.transformer_2 = None diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index f83d5be29f1b..cd170a836b0e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -104,6 +104,8 @@ doc="Wan 2.1 & 2.2 image-to-video family.", ) class WanImageToVideoPipeline(BasePipeline): + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + def __init__(self, pipeline_config): # Wan2.2 14B two-stage denoising parameters self.transformer_2 = None diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 13cbe77bdf12..9eee8e03129c 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -106,6 +106,13 @@ class BasePipeline(nn.Module): Base class for diffusion pipelines. """ + # Cache-acceleration backends this pipeline supports. PipelineLoader calls + # _setup_cache_acceleration() on every pipeline after torch.compile; + # pipelines that don't opt in here skip enablement with a warning instead + # of failing at enable time (e.g. no Cache-DiT enabler / TeaCache extractor + # registered for their transformer). + SUPPORTED_CACHE_BACKENDS: Tuple[str, ...] = () + @classmethod def resolve_variant(cls, config: "DiffusionPipelineConfig") -> Type["BasePipeline"]: """Return *cls* or a more specialized subclass based on *config*. @@ -496,14 +503,25 @@ def _setup_cache_acceleration(self) -> None: cfg = self.pipeline_config - if cfg.cache_backend == "cache_dit": + backend = cfg.cache_backend + if backend is None: + return + + if backend not in self.SUPPORTED_CACHE_BACKENDS: + logger.warning( + f"cache_backend={backend!r} is configured but {self.__class__.__name__} " + f"does not support it (supported: {list(self.SUPPORTED_CACHE_BACKENDS)}); " + "continuing without cache acceleration." + ) + return + + if backend == "cache_dit": acc = CacheDiTAccelerator(self, cfg.cache_dit) acc.wrap() self.cache_accelerator = acc return - use_teacache = cfg.cache_backend == "teacache" - if not use_teacache: + if backend != "teacache": return acc = TeaCacheAccelerator(self, cfg.teacache) diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index 6deeb6b6bfab..4b9ae3dfa09f 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -62,6 +62,8 @@ def model_copy(self, update=None): class _PipelineTeacacheTestDouble: """Minimal object for BasePipeline._setup_cache_acceleration / _apply_teacache_coefficients tests.""" + SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") + def __init__(self, model_config: DiffusionModelConfig): self.model_config = model_config self.cache_accelerator = None From 341edf484d38020d8f1ec93677140b0785e4c8a9 Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 23 Jul 2026 15:29:57 -0700 Subject: [PATCH 5/7] remove supported cache backend Signed-off-by: BrianLi23 --- .../visual_gen/models/flux/pipeline_flux.py | 2 -- .../visual_gen/models/flux/pipeline_flux2.py | 2 -- .../visual_gen/models/ltx2/pipeline_ltx2.py | 2 -- .../models/ltx2/pipeline_ltx2_two_stages.py | 2 -- .../visual_gen/models/wan/pipeline_wan.py | 2 -- .../visual_gen/models/wan/pipeline_wan_i2v.py | 2 -- tensorrt_llm/_torch/visual_gen/pipeline.py | 24 +++---------------- .../_torch/visual_gen/test_teacache.py | 2 -- 8 files changed, 3 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index 8e9a4b4ed121..e08d07047878 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -50,8 +50,6 @@ class FluxPipeline(BasePipeline): Supports FLUX.1-dev (50 steps, guidance) and FLUX.1-schnell (4 steps, no guidance). """ - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - def __init__(self, pipeline_config): if ( pipeline_config.visual_gen_mapping is not None diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 0af916483af8..1b2116f7f27f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -105,8 +105,6 @@ class Flux2Pipeline(BasePipeline): Follows WAN pipeline pattern for DiffusionModelLoader integration. """ - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - # Hidden state layers per text encoder type (auto-detected at load time) _TEXT_ENCODER_CONFIG = { "Mistral3ForConditionalGeneration": { diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 1e6446b57e1f..45a46ea3b015 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -620,8 +620,6 @@ class LTX2Pipeline(BasePipeline): ``transformers`` library. """ - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - def __init__(self, model_config): _sa_cfg = model_config.attention.sparse_attention_config if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa": diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index c7ed27ab189d..b7ef703be22d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -934,8 +934,6 @@ class LTX2TwoStagesPipeline(LTX2Pipeline): then decode. """ - SUPPORTED_CACHE_BACKENDS = ("teacache",) - @property def common_warmup_shapes(self) -> list: return [(512, 768, 121)] diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 5dea9408f09c..8230f10823e5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -112,8 +112,6 @@ doc="Wan 2.1 & 2.2 text-to-video family.", ) class WanPipeline(BasePipeline): - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - def __init__(self, pipeline_config): # Wan2.2 A14B two-stage denoising parameters self.transformer_2 = None diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index cd170a836b0e..f83d5be29f1b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -104,8 +104,6 @@ doc="Wan 2.1 & 2.2 image-to-video family.", ) class WanImageToVideoPipeline(BasePipeline): - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - def __init__(self, pipeline_config): # Wan2.2 14B two-stage denoising parameters self.transformer_2 = None diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 9eee8e03129c..13cbe77bdf12 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -106,13 +106,6 @@ class BasePipeline(nn.Module): Base class for diffusion pipelines. """ - # Cache-acceleration backends this pipeline supports. PipelineLoader calls - # _setup_cache_acceleration() on every pipeline after torch.compile; - # pipelines that don't opt in here skip enablement with a warning instead - # of failing at enable time (e.g. no Cache-DiT enabler / TeaCache extractor - # registered for their transformer). - SUPPORTED_CACHE_BACKENDS: Tuple[str, ...] = () - @classmethod def resolve_variant(cls, config: "DiffusionPipelineConfig") -> Type["BasePipeline"]: """Return *cls* or a more specialized subclass based on *config*. @@ -503,25 +496,14 @@ def _setup_cache_acceleration(self) -> None: cfg = self.pipeline_config - backend = cfg.cache_backend - if backend is None: - return - - if backend not in self.SUPPORTED_CACHE_BACKENDS: - logger.warning( - f"cache_backend={backend!r} is configured but {self.__class__.__name__} " - f"does not support it (supported: {list(self.SUPPORTED_CACHE_BACKENDS)}); " - "continuing without cache acceleration." - ) - return - - if backend == "cache_dit": + if cfg.cache_backend == "cache_dit": acc = CacheDiTAccelerator(self, cfg.cache_dit) acc.wrap() self.cache_accelerator = acc return - if backend != "teacache": + use_teacache = cfg.cache_backend == "teacache" + if not use_teacache: return acc = TeaCacheAccelerator(self, cfg.teacache) diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index 4b9ae3dfa09f..6deeb6b6bfab 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -62,8 +62,6 @@ def model_copy(self, update=None): class _PipelineTeacacheTestDouble: """Minimal object for BasePipeline._setup_cache_acceleration / _apply_teacache_coefficients tests.""" - SUPPORTED_CACHE_BACKENDS = ("teacache", "cache_dit") - def __init__(self, model_config: DiffusionModelConfig): self.model_config = model_config self.cache_accelerator = None From e3d43bd72c4f13e086d002d468a84213ccc57fc2 Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 23 Jul 2026 16:48:53 -0700 Subject: [PATCH 6/7] add additional tests Signed-off-by: BrianLi23 --- .../_torch/visual_gen/test_cache_dit.py | 206 +++++++++++++++++- 1 file changed, 204 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cache_dit.py b/tests/unittest/_torch/visual_gen/test_cache_dit.py index 9a345ba4a404..ab7582af22f0 100644 --- a/tests/unittest/_torch/visual_gen/test_cache_dit.py +++ b/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,105 @@ 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: explicit raise + declaration consistency (CPU) +# --------------------------------------------------------------------------- + + +@requires_cache_dit +class TestCacheDiTEnablerRegistry: + def test_unregistered_pipeline_class_raises(self): + """A pipeline without a Cache-DiT enabler must fail loudly at enable + time (e.g. LTX2TwoStagesPipeline + cache_backend=cache_dit), not warn + and continue at full denoise cost.""" + from tensorrt_llm._torch.visual_gen.cache.cache_dit_enablers import ( + enable_cache_dit_for_pipeline, + ) + from tensorrt_llm.visual_gen.args import CacheDiTConfig + + class LTX2TwoStagesPipeline: # same name as the real unregistered variant + pass + + with pytest.raises(ValueError, match="no enabler registered.*LTX2TwoStagesPipeline"): + enable_cache_dit_for_pipeline(LTX2TwoStagesPipeline(), CacheDiTConfig()) + + 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}" + ) From ed0d8b641b901adbb7e3883f205d7efba4dfb739 Mon Sep 17 00:00:00 2001 From: BrianLi23 Date: Thu, 23 Jul 2026 17:49:16 -0700 Subject: [PATCH 7/7] delete unneeded test Signed-off-by: BrianLi23 --- .../_torch/visual_gen/test_cache_dit.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cache_dit.py b/tests/unittest/_torch/visual_gen/test_cache_dit.py index ab7582af22f0..09cbaa55d35c 100644 --- a/tests/unittest/_torch/visual_gen/test_cache_dit.py +++ b/tests/unittest/_torch/visual_gen/test_cache_dit.py @@ -577,27 +577,12 @@ def test_eager_blocks_keep_pattern_checks(self, is_flux2): # --------------------------------------------------------------------------- -# 4) Enabler registry: explicit raise + declaration consistency (CPU) +# 4) Enabler registry declaration consistency (CPU) # --------------------------------------------------------------------------- @requires_cache_dit class TestCacheDiTEnablerRegistry: - def test_unregistered_pipeline_class_raises(self): - """A pipeline without a Cache-DiT enabler must fail loudly at enable - time (e.g. LTX2TwoStagesPipeline + cache_backend=cache_dit), not warn - and continue at full denoise cost.""" - from tensorrt_llm._torch.visual_gen.cache.cache_dit_enablers import ( - enable_cache_dit_for_pipeline, - ) - from tensorrt_llm.visual_gen.args import CacheDiTConfig - - class LTX2TwoStagesPipeline: # same name as the real unregistered variant - pass - - with pytest.raises(ValueError, match="no enabler registered.*LTX2TwoStagesPipeline"): - enable_cache_dit_for_pipeline(LTX2TwoStagesPipeline(), CacheDiTConfig()) - 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