user-configured TeaCache coefficients

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>

Author: o-stoner

tensorrt_llm/_torch/visual_gen/config.py

@@ -160,7 +160,9 @@ class TeaCacheConfig(StrictBaseModel):
         teacache_thresh: Distance threshold for cache decisions (lower = more caching)
         use_ret_steps: Use aggressive warmup mode (5 steps) vs minimal (1 step)
         coefficients: Polynomial coefficients for rescaling embedding distances
-                     Applied as: rescaled_distance = poly(raw_distance)
+                     Applied as: rescaled_distance = poly(raw_distance).
+                     None means use the pipeline built-in coefficient table (checkpoint path
+                     matching). A non-None list overrides that table.
         ret_steps: Number of warmup steps (always compute, initialized at runtime)
         cutoff_steps: Step to stop caching (always compute after, initialized at runtime)
         num_steps: Total inference steps (set at runtime)
@@ -171,7 +173,7 @@ class TeaCacheConfig(StrictBaseModel):
     teacache_thresh: float = PydanticField(0.2, gt=0.0)
     use_ret_steps: bool = False
 
-    coefficients: List[float] = PydanticField(default_factory=lambda: [1.0, 0.0])
+    coefficients: Optional[List[float]] = None
 
     # Runtime state fields (initialized by TeaCacheBackend.refresh)
     ret_steps: Optional[int] = None
@@ -186,8 +188,8 @@ class TeaCacheConfig(StrictBaseModel):
     @model_validator(mode="after")
     def validate_teacache(self) -> "TeaCacheConfig":
         """Validate TeaCache configuration."""
-        # Validate coefficients
-        if len(self.coefficients) == 0:
+        # Validate coefficients (when provided)
+        if self.coefficients is not None and len(self.coefficients) == 0:
             raise ValueError("TeaCache coefficients list cannot be empty")
 
         # Validate ret_steps if set

tensorrt_llm/_torch/visual_gen/pipeline.py

@@ -20,6 +20,16 @@
     from .config import DiffusionModelConfig
 
 
+def _teacache_coefficients_are_explicit_user_override(teacache_cfg: Any) -> bool:
+    """Return True if teacache.coefficients should skip built-in variant table matching.
+
+    Only None means auto (use checkpoint path table). Any non-empty list, including
+    the identity polynomial [1.0, 0.0], is treated as an explicit user override.
+    """
+
+    return teacache_cfg.coefficients is not None
+
+
 class BasePipeline(nn.Module):
     """
     Base class for diffusion pipelines.
@@ -250,10 +260,85 @@ def post_load_weights(self) -> None:
         if self.transformer is not None and hasattr(self.transformer, "post_load_weights"):
             self.transformer.post_load_weights()
 
+    def _apply_teacache_coefficients(self, coefficients: Optional[Dict] = None) -> None:
+        """Resolve TeaCache polynomial coefficients into model_config.teacache.
+
+        Precedence:
+
+        1. User-specified TeaCacheConfig.coefficients — any non-None list skips built-in
+           variant matching.
+
+        2. Pipeline table — if step 1 does not apply and the coefficients argument is a
+           non-empty dict (model-specific tables from the pipeline subclass), match
+           pretrained_config._name_or_path against keys and set coefficients (and optional
+           default_thresh).
+
+        3. If coefficients is still None after step 2, _setup_teacache raises: TeaCache
+           must not run without resolved coefficients.
+
+        Args:
+            coefficients: Optional mapping from variant key to coefficient list or nested
+                dict (ret_steps / standard), from the pipeline subclass.
+        """
+        teacache_cfg = self.model_config.teacache
+        if _teacache_coefficients_are_explicit_user_override(teacache_cfg):
+            logger.info(
+                "TeaCache: Using user-configured coefficients "
+                "(skipping built-in checkpoint variant matching)"
+            )
+            return
+
+        teacache_explicit = teacache_cfg.model_dump(exclude_unset=True)
+
+        if not coefficients:
+            return
+
+        checkpoint_path = getattr(self.model_config.pretrained_config, "_name_or_path", "") or ""
+
+        for model_size, coeff_data in coefficients.items():
+            # Match model size in path (case-insensitive, e.g., "1.3B", "14B", "dev")
+            path_l = checkpoint_path.lower()
+            key_l = model_size.lower()
+            if key_l not in path_l:
+                continue
+
+            if isinstance(coeff_data, dict):
+                # Select coefficient set based on warmup mode
+                mode = "ret_steps" if teacache_cfg.use_ret_steps else "standard"
+                if mode not in coeff_data:
+                    logger.warning(
+                        "TeaCache: matched variant %r but table has no %r entry "
+                        "(available keys: %s). Trying other variants.",
+                        model_size,
+                        mode,
+                        list(coeff_data.keys()),
+                    )
+                    continue
+                teacache_cfg.coefficients = coeff_data[mode]
+                logger.info(f"TeaCache: Using {model_size} coefficients ({mode} mode)")
+                # Apply model-specific default threshold if user didn't explicitly set one
+                default_thresh = coeff_data.get("default_thresh")
+                if default_thresh is not None and "teacache_thresh" not in teacache_explicit:
+                    teacache_cfg.teacache_thresh = default_thresh
+                    logger.info(f"TeaCache: Using {model_size} default threshold {default_thresh}")
+                break
+            else:
+                # Single coefficient list (no mode distinction)
+                teacache_cfg.coefficients = coeff_data
+                logger.info(f"TeaCache: Using {model_size} coefficients")
+                break
+        else:
+            raise ValueError(
+                f"TeaCache: No coefficients found for checkpoint '{checkpoint_path}'. "
+                f"Available variants: {list(coefficients.keys())}. "
+                f"Set teacache.coefficients explicitly in VisualGenArgs to use TeaCache anyway, "
+                f"or use a checkpoint path that contains one of the variant keys."
+            )
+
     def _setup_teacache(self, model, coefficients: Optional[Dict] = None):
         """Setup TeaCache optimization for the transformer model.
 
-        TeaCache caches transformer block outputs when timestep embeddings change slowly,
+        TeaCache caches transformer block outputs,
         reducing computation during the denoising loop.
 
         Args:
@@ -268,42 +353,14 @@ def _setup_teacache(self, model, coefficients: Optional[Dict] = None):
         if not teacache_cfg.enable_teacache:
             return
 
-        # Apply model-specific polynomial coefficients
-        # Coefficients are used to rescale embedding distances for cache decisions
-        if coefficients:
-            checkpoint_path = (
-                getattr(self.model_config.pretrained_config, "_name_or_path", "") or ""
+        self._apply_teacache_coefficients(coefficients)
+
+        if teacache_cfg.coefficients is None:
+            raise ValueError(
+                "TeaCache is enabled but no polynomial coefficients were resolved. "
+                "Set teacache.coefficients in VisualGenArgs, or use a pipeline and "
+                "checkpoint whose path matches a built-in coefficient table."
             )
-            for model_size, coeff_data in coefficients.items():
-                # Match model size in path (case-insensitive, e.g., "1.3B", "14B", "dev")
-                if model_size.lower() in checkpoint_path.lower():
-                    if isinstance(coeff_data, dict):
-                        # Select coefficient set based on warmup mode
-                        mode = "ret_steps" if teacache_cfg.use_ret_steps else "standard"
-                        if mode in coeff_data:
-                            teacache_cfg.coefficients = coeff_data[mode]
-                            logger.info(f"TeaCache: Using {model_size} coefficients ({mode} mode)")
-                        # Apply model-specific default threshold if user didn't explicitly set one
-                        default_thresh = coeff_data.get("default_thresh")
-                        if (
-                            default_thresh is not None
-                            and "teacache_thresh" not in teacache_cfg.model_fields_set
-                        ):
-                            teacache_cfg.teacache_thresh = default_thresh
-                            logger.info(
-                                f"TeaCache: Using {model_size} default threshold {default_thresh}"
-                            )
-                    else:
-                        # Single coefficient list (no mode distinction)
-                        teacache_cfg.coefficients = coeff_data
-                        logger.info(f"TeaCache: Using {model_size} coefficients")
-                    break
-            else:
-                raise ValueError(
-                    f"TeaCache: No coefficients found for checkpoint '{checkpoint_path}'. "
-                    f"Available variants: {list(coefficients.keys())}. "
-                    f"TeaCache is not supported for this model variant."
-                )
 
         # Initialize and enable TeaCache backend
         logger.info("TeaCache: Initializing...")

tests/unittest/_torch/visual_gen/test_teacache.py

@@ -15,6 +15,7 @@
 """Unit tests for TeaCache (CPU-only, no model weights needed)."""
 
 from types import SimpleNamespace
+from typing import Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -24,21 +25,116 @@
 from tensorrt_llm._torch.visual_gen.teacache import TeaCacheBackend
 
 
+class _PipelineTeacacheTestDouble:
+    """Minimal object for BasePipeline._setup_teacache / _apply_teacache_coefficients tests."""
+
+    def __init__(self, model_config: DiffusionModelConfig):
+        self.model_config = model_config
+        self.cache_backend = None
+
+    def _apply_teacache_coefficients(self, coefficients: Optional[dict] = None):
+        return BasePipeline._apply_teacache_coefficients(self, coefficients)
+
+
 class TestSetupTeacache:
     """Tests for _setup_teacache coefficient matching and fail-early behavior."""
 
     def _make_pipeline_mock(self, checkpoint_name, use_ret_steps=False):
-        pipeline = MagicMock()
-        pipeline.model_config = DiffusionModelConfig(
-            pretrained_config=SimpleNamespace(_name_or_path=f"/path/to/{checkpoint_name}/snapshot"),
-            teacache=TeaCacheConfig(
-                enable_teacache=True,
-                teacache_thresh=0.3,
-                use_ret_steps=use_ret_steps,
-            ),
-            skip_create_weights_in_init=True,
+        return _PipelineTeacacheTestDouble(
+            DiffusionModelConfig(
+                pretrained_config=SimpleNamespace(
+                    _name_or_path=f"/path/to/{checkpoint_name}/snapshot"
+                ),
+                teacache=TeaCacheConfig(
+                    enable_teacache=True,
+                    teacache_thresh=0.3,
+                    use_ret_steps=use_ret_steps,
+                    coefficients=None,  # auto: resolve from pipeline variant table
+                ),
+                skip_create_weights_in_init=True,
+            )
         )
-        return pipeline
+
+    def _make_pipeline_teacache_enable_only(self, checkpoint_name, use_ret_steps=False):
+        """enable_teacache only so teacache_thresh stays unset (exclude_unset omits it)."""
+        return _PipelineTeacacheTestDouble(
+            DiffusionModelConfig(
+                pretrained_config=SimpleNamespace(
+                    _name_or_path=f"/path/to/{checkpoint_name}/snapshot"
+                ),
+                teacache=TeaCacheConfig(
+                    enable_teacache=True,
+                    use_ret_steps=use_ret_steps,
+                    coefficients=None,
+                ),
+                skip_create_weights_in_init=True,
+            )
+        )
+
+    def test_setup_teacache_raises_when_no_table_and_no_user_coefficients(self):
+        """Fails if TeaCache is on but the pipeline passes no coefficient table."""
+        pipeline = self._make_pipeline_mock("FLUX.1-dev")
+        with pytest.raises(ValueError, match="no polynomial coefficients were resolved"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), None)
+
+    def test_setup_teacache_raises_when_empty_table(self):
+        """Same as no table: nothing to resolve from."""
+        pipeline = self._make_pipeline_mock("FLUX.1-dev")
+        with pytest.raises(ValueError, match="no polynomial coefficients were resolved"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), {})
+
+    def test_matching_variant_selects_ret_steps_mode(self):
+        """Nested table: use_ret_steps=True selects ret_steps coefficients."""
+        pipeline = self._make_pipeline_mock("FLUX.1-dev", use_ret_steps=True)
+        coefficients = {
+            "dev": {"standard": [1.0, 2.0, 3.0], "ret_steps": [4.0, 5.0]},
+        }
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), coefficients)
+
+        assert pipeline.model_config.teacache.coefficients == [4.0, 5.0]
+
+    def test_flat_list_table_entry(self):
+        """Table value may be a plain list (no standard/ret_steps nesting)."""
+        pipeline = self._make_pipeline_mock("FLUX.1-dev")
+        coefficients = {"dev": [11.0, 22.0, 33.0]}
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), coefficients)
+
+        assert pipeline.model_config.teacache.coefficients == [11.0, 22.0, 33.0]
+
+    def test_default_thresh_from_table_when_user_did_not_set_teacache_thresh(self):
+        """default_thresh applies when teacache_thresh was not explicitly set (exclude_unset)."""
+        pipeline = self._make_pipeline_teacache_enable_only("FLUX.1-dev")
+        builtin = {
+            "dev": {
+                "standard": [1.0, 2.0],
+                "ret_steps": [3.0, 4.0],
+                "default_thresh": 0.42,
+            },
+        }
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), builtin)
+
+        assert pipeline.model_config.teacache.coefficients == [1.0, 2.0]
+        assert pipeline.model_config.teacache.teacache_thresh == 0.42
+
+    def test_explicit_identity_coefficients_still_skip_table(self):
+        """[1.0, 0.0] is a user override: no variant lookup, no ValueError on unknown path."""
+        pipeline = self._make_pipeline_mock("FLUX.1-unknown-variant")
+        pipeline.model_config.teacache = TeaCacheConfig(
+            enable_teacache=True,
+            teacache_thresh=0.3,
+            coefficients=[1.0, 0.0],
+        )
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(
+                pipeline,
+                MagicMock(),
+                {"dev": {"standard": [99.0, 99.0]}},
+            )
+
+        assert pipeline.model_config.teacache.coefficients == [1.0, 0.0]
 
     def test_matching_variant_selects_coefficients(self):
         """Picks coefficients whose key appears in checkpoint path."""
@@ -69,3 +165,54 @@ def test_disabled_teacache_is_noop(self):
 
         BasePipeline._setup_teacache(pipeline, MagicMock(), {"dev": [1.0]})
         assert pipeline.cache_backend is None
+
+    def test_user_configured_coefficients_skip_variant_matching(self):
+        """Explicit TeaCacheConfig.coefficients skips dict lookup (no ValueError)."""
+        pipeline = self._make_pipeline_mock("FLUX.1-unknown-variant")
+        pipeline.model_config.teacache = TeaCacheConfig(
+            enable_teacache=True,
+            teacache_thresh=0.3,
+            coefficients=[0.25, 0.5, 0.75],
+        )
+        builtin = {
+            "dev": {"standard": [1.0, 2.0, 3.0], "ret_steps": [4.0, 5.0]},
+        }
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), builtin)
+
+        assert pipeline.model_config.teacache.coefficients == [0.25, 0.5, 0.75]
+
+    def test_user_configured_coefficients_take_precedence_over_builtin_table(self):
+        """User coefficients are not overwritten when a built-in variant would also match."""
+        pipeline = self._make_pipeline_mock("FLUX.1-dev")
+        pipeline.model_config.teacache = TeaCacheConfig(
+            enable_teacache=True,
+            teacache_thresh=0.3,
+            coefficients=[9.0, 8.0, 7.0],
+        )
+        builtin = {
+            "dev": {"standard": [1.0, 2.0, 3.0], "ret_steps": [4.0, 5.0]},
+        }
+        with patch.object(TeaCacheBackend, "enable"):
+            BasePipeline._setup_teacache(pipeline, MagicMock(), builtin)
+
+        assert pipeline.model_config.teacache.coefficients == [9.0, 8.0, 7.0]
+
+    def test_apply_teacache_coefficients_only(self):
+        """_apply_teacache_coefficients updates config without enabling backend."""
+        pipeline = self._make_pipeline_mock("FLUX.1-unknown-variant")
+        pipeline.model_config.teacache = TeaCacheConfig(
+            enable_teacache=True,
+            coefficients=[0.1, 0.2],
+        )
+        # Would raise if variant matching ran
+        BasePipeline._apply_teacache_coefficients(pipeline, {"dev": {"standard": [99.0]}})
+        assert pipeline.model_config.teacache.coefficients == [0.1, 0.2]
+
+
+class TestTeaCacheConfigValidation:
+    """TeaCacheConfig validation (no pipeline)."""
+
+    def test_empty_coefficients_rejected(self):
+        with pytest.raises(ValueError, match="cannot be empty"):
+            TeaCacheConfig(enable_teacache=True, coefficients=[])