[Bugfix] Calibrate Z-Image TeaCache coefficients and add ZImageAdapter

docs/design/feature/teacache.md

@@ -20,7 +20,7 @@ This section describes how to add TeaCache to a diffusion transformer model. We
 
 ### What is TeaCache?
 
-TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides **1.5x-2.0x speedup** with minimal quality loss.
+TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides up to **1.5x-2.0x speedup** with minimal quality loss on models with calibrated coefficients; few-step distilled models gain less.
 
 The core insight is that the modulated input (after normalization and timestep conditioning) changes gradually across timesteps. By measuring the L1 distance between consecutive modulated inputs and comparing it to a threshold, TeaCache decides whether to execute the full transformer blocks or reuse the cached residual from the previous step.
 
@@ -260,15 +260,15 @@ The polynomial coefficients rescale L1 distances between consecutive modulated i
 
 | Approach | Performance | Effort |
 |----------|-------------|--------|
-| Using defaults from similar model | Within 5-10% of optimal | Low |
+| Using defaults from similar model | Unreliable: Z-Image with Qwen-Image's coefficients under-predicted output change ~3x and over-skipped steps (#8270); treat as a placeholder and calibrate | Low |
 | Estimating custom coefficients | Best performance | Medium |
 
 #### Implement Data Collection Adapter
 
 Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
 
 ```python
-class YourModelAdapter:
+class YourModelAdapter(DefaultAdapter):
     """Adapter for coefficient estimation on your model."""
 
     @staticmethod
@@ -398,7 +398,7 @@ images = omni.generate(
 **Verify:**
 
 1. **Check logs** - Look for TeaCache initialization messages
-2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x)
+2. **Compare performance** - Measure speedup vs baseline (expect up to 1.5x-2.0x on 50-step runs; less on few-step distilled models)
 3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical)
 
 See more detailed examples in [user guide for teacache](../../user_guide/diffusion/cache_acceleration/teacache.md).

docs/user_guide/diffusion/cache_acceleration/teacache.md

@@ -15,7 +15,7 @@
 
 ## Overview
 
-TeaCache accelerates diffusion model inference by caching transformer computations when consecutive timesteps are similar, providing **1.5x-2.0x speedup** with minimal quality loss. It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters without sacrificing generation quality.
+TeaCache accelerates diffusion model inference by caching transformer computations when consecutive timesteps are similar, providing up to **1.5x-2.0x speedup** with minimal quality loss on models with calibrated coefficients and enough steps to skip (few-step distilled models gain less; see the Summary). It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters without sacrificing generation quality.
 
 See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
 
@@ -119,7 +119,7 @@ In `OmniDiffusionConfig`
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `rel_l1_thresh` | float | `0.2` | Similarity threshold for cache reuse. Lower values prioritize quality (less caching), higher values prioritize speed (more caching). Suggested range: 0.1-0.8 |
+| `rel_l1_thresh` | float | `0.2` | Similarity threshold for cache reuse. Lower values prioritize quality (less caching), higher values prioritize speed (more caching). Suggested range: 0.1-0.3; higher values mostly cost quality on calibrated models |
 | `coefficients` | list[float] \| None | `None` | Polynomial coefficients for rescaling L1 distance. Must contain exactly 5 elements if provided. If `None`, uses model-specific defaults based on transformer type. |
 
 Users can find the default model coefficients in [`vllm_omni/diffusion/cache/teacache/config.py`](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/cache/teacache/config.py), for example:
@@ -149,7 +149,7 @@ _MODEL_COEFFICIENTS = {
 **Good for:**
 
 - Production deployments requiring faster inference, tolerant of minimal quality loss
-- Scenarios where 1.5-2x speedup is valuable
+- Scenarios where a model-dependent 1.2-2x speedup is valuable
 - Useful for single-card acceleration
 
 **Not for:**
@@ -178,17 +178,17 @@ cache_config={"rel_l1_thresh": 0.1}
 **Symptoms**: Actual speedup is less than expected (< 1.3x)
 
 **Solutions**:
-1. Increase the threshold to enable more aggressive caching:
+1. Increase the threshold to enable more aggressive caching (non-distilled models only; on few-step distilled models such as Z-Image-Turbo a higher threshold mostly costs quality, see the Summary):
    ```python
-   cache_config={"rel_l1_thresh": 0.8}
+   cache_config={"rel_l1_thresh": 0.3}
    ```
-2. Ensure you're using sufficient inference steps (35+ recommended)
+2. Ensure you're using sufficient inference steps (35+ recommended for non-distilled models; few-step distilled models cannot add steps)
 3. Check that your model architecture is supported (see Supported Models section)
 
 ---
 
 
 ## Summary
 
-1. ✅ **Enable TeaCache** - Set `cache_backend="tea_cache"` to get 1.5x-2.0x speedup with optimized defaults
+1. ✅ **Enable TeaCache** - Set `cache_backend="tea_cache"` to get up to 1.5x-2.0x speedup with optimized defaults on 50-step runs; few-step distilled models (e.g. Z-Image-Turbo at 25 steps) change too much per step to skip many, so expect ~1.1x there
 2. ✅ **(Optional) Customize** - Adjust thresholds and polynomial coefficients for specific speed/quality trade-offs

tests/diffusion/cache/test_teacache_config.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for TeaCache model-specific defaults (coefficients and estimator adapters)."""
+
+import pytest
+
+from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
+    _MODEL_ADAPTERS,
+    DataCollectionHook,
+    ZImageAdapter,
+)
+from vllm_omni.diffusion.cache.teacache.config import _MODEL_COEFFICIENTS, TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import extract_zimage_context
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_zimage_has_calibrated_coefficients():
+    """Z-Image must not fall back to the Qwen-Image placeholder polynomial (#8270)."""
+    zimage = _MODEL_COEFFICIENTS["ZImageTransformer2DModel"]
+    assert zimage != _MODEL_COEFFICIENTS["QwenImageTransformer2DModel"]
+    # Pin the calibrated fit so a recalibration has to update this test deliberately.
+    assert zimage == pytest.approx([-7.54613422e01, -9.23596156e01, 5.75318402e01, -3.76790311e00, 2.27176809e-01])
+
+    config = TeaCacheConfig(transformer_type="ZImageTransformer2DModel")
+    assert config.coefficients == zimage
+    assert config.rel_l1_thresh == 0.2
+
+
+def test_zimage_estimator_adapter_registered():
+    assert _MODEL_ADAPTERS["ZImage"] is ZImageAdapter
+    assert ZImageAdapter.model_class_name == "ZImagePipeline"
+    assert ZImageAdapter.uses_tf_config is True
+
+
+def test_data_collection_hook_resolves_extractor_at_init():
+    """The estimator hook binds its extractor in __init__, so an unknown type fails before any forward."""
+    hook = DataCollectionHook("ZImageTransformer2DModel")
+    assert hook.extractor_fn is extract_zimage_context
+    assert hook.current_trajectory == []
+
+    with pytest.raises(ValueError, match="Unknown model type"):
+        DataCollectionHook("NotARegisteredTransformer")

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import os
 from typing import Any
@@ -25,13 +25,9 @@ class DataCollectionHook(ModelHook):
     def __init__(self, transformer_type: str):
         super().__init__()
         self.transformer_type = transformer_type
-        self.extractor_fn = None
+        self.extractor_fn = get_extractor(transformer_type)
         self.current_trajectory: list[tuple[np.ndarray, np.ndarray]] = []
 
-    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
-        self.extractor_fn = get_extractor(self.transformer_type)
-        return module
-
     def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
         ctx = self.extractor_fn(module, *args, **kwargs)
         # NOTE: We upcast to float32 to also handle bfloat16.
@@ -56,7 +52,7 @@ def stop_collection(self) -> list[tuple[np.ndarray, np.ndarray]]:
 class DefaultAdapter:
     """Default adapter for standard diffusers pipelines."""
 
-    model_class_name = None
+    model_class_name: str | None = None
     uses_tf_config = True
 
     @classmethod
@@ -132,11 +128,18 @@ class StableAudioAdapter(DefaultAdapter):
     model_class_name = "StableAudioPipeline"
 
 
-_MODEL_ADAPTERS: dict[str, type] = {
+class ZImageAdapter(DefaultAdapter):
+    """Adapter for Z-Image coefficient estimation."""
+
+    model_class_name = "ZImagePipeline"
+
+
+_MODEL_ADAPTERS: dict[str, type[DefaultAdapter]] = {
     "Bagel": BagelAdapter,
     "StableAudio": StableAudioAdapter,
     "Flux2": Flux2Adapter,
     "LongCat": LongCatAdapter,
+    "ZImage": ZImageAdapter,
 }
 
 _EPSILON = 1e-6

vllm_omni/diffusion/cache/teacache/config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from dataclasses import dataclass
 
@@ -46,13 +46,27 @@
         7.61309272e-01,
     ],
     # Z-Image transformer coefficients
-    # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future
+    # Calibrated on Z-Image-Turbo with ZImageAdapter + the TeaCacheCoefficientEstimator
+    # hook path (per-step rel-L1 computed on the GPU instead of storing full CPU
+    # trajectories): 70 Parti prompts x {25, 50} steps, 1024x1024, guidance_scale=4.0,
+    # 5,110 step pairs, 4th-order fit (R^2 0.69). Measured input rel_l1 range
+    # 0.024-0.336 (mean 0.075); the quartic peaks near x~0.28 and crosses zero
+    # near x~0.40, so it is only valid inside that range: the hook applies it
+    # unclamped, and a step with input distance ~0.40 would rescale to ~0 and be
+    # served from cache despite being a large step. Validated only on
+    # Z-Image-Turbo; other checkpoints that reach TeaCache through this class
+    # name inherit these values uncalibrated. The previous Qwen-Image
+    # placeholder predicted ~1/3 of the measured output change and over-skipped
+    # steps (see #8270). With the default rel_l1_thresh=0.2 Z-Image-Turbo skips
+    # about half of its steps at 50 steps (~1.5x) and few at 25 steps (~1.1x),
+    # because the distilled model changes >20% per step early on; raising the
+    # threshold to 0.25-0.30 buys ~10% more speed for a visible quality drop.
     "ZImageTransformer2DModel": [
-        -4.50000000e02,
-        2.80000000e02,
-        -4.50000000e01,
-        3.20000000e00,
-        -2.00000000e-02,
+        -7.54613422e01,
+        -9.23596156e01,
+        5.75318402e01,
+        -3.76790311e00,
+        2.27176809e-01,
     ],
     # Estimated TeaCache polynomial coefficients for StableAudioDiTModel.
     "StableAudioDiTModel": [
@@ -73,7 +87,8 @@
         6.78098549e-01,
     ],
     # Flux2 transformer coefficients
-    # Copied from Qwen-Image, need to be tuned specifically for Flux2 in future
+    # Copied from Qwen-Image, still uncalibrated; tracked in #8411 together with
+    # the Flux2Klein borrow and the Ming subclasses of ZImageTransformer2DModel.
     "Flux2Transformer2DModel": [
         -4.50000000e02,
         2.80000000e02,