From 974481c3e7a02accee4dcef2a564406f5ffe49ce Mon Sep 17 00:00:00 2001 From: Kamila Kuatova Date: Thu, 1 Oct 2026 16:08:33 -0400 Subject: [PATCH 1/3] [Bugfix] Calibrate Z-Image TeaCache coefficients and add ZImageAdapter The TeaCache polynomial for ZImageTransformer2DModel was a placeholder copied from Qwen-Image. On Z-Image it predicts about one third of the measured step-to-step output change, so the accumulated distance stays under rel_l1_thresh for too long and TeaCache skips steps that matter. Users got visibly different (blurred) images instead of a faster version of the same one. - Replace the placeholder with coefficients calibrated on Z-Image-Turbo: 70 Parti prompts x {25, 50} steps, 1024x1024, guidance_scale=4.0, 5,110 consecutive-step pairs, 4th-order polyfit (R^2 0.69). - Add ZImageAdapter to TeaCacheCoefficientEstimator so the calibration can be reproduced, and type-annotate the adapter registry so mypy passes. - Add CPU tests for the Z-Image TeaCache defaults. A100-40GB, 36 prompt/seed pairs vs. uncached output, rel_l1_thresh=0.2: mean SSIM 0.786 -> 0.980 at 25 steps and 0.563 -> 0.902 at 50 steps; speed-up 1.08x / 1.54x (the previous 1.96x / 2.85x came from skipping steps that should have run). Fixes #8270 Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kamila Kuatova --- tests/diffusion/cache/test_teacache_config.py | 28 +++++++++++++++++++ .../cache/teacache/coefficient_estimator.py | 19 +++++++------ vllm_omni/diffusion/cache/teacache/config.py | 17 ++++++----- 3 files changed, 49 insertions(+), 15 deletions(-) create mode 100644 tests/diffusion/cache/test_teacache_config.py diff --git a/tests/diffusion/cache/test_teacache_config.py b/tests/diffusion/cache/test_teacache_config.py new file mode 100644 index 00000000000..f2bcefc3d36 --- /dev/null +++ b/tests/diffusion/cache/test_teacache_config.py @@ -0,0 +1,28 @@ +# 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, ZImageAdapter +from vllm_omni.diffusion.cache.teacache.config import _MODEL_COEFFICIENTS, TeaCacheConfig + +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 len(zimage) == 5 + assert zimage != _MODEL_COEFFICIENTS["QwenImageTransformer2DModel"] + + 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 diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index 49ba806d55b..6321065e7e3 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index b83f600d6e7..4b5b1e1a300 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/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,16 @@ 7.61309272e-01, ], # Z-Image transformer coefficients - # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future + # Calibrated with ZImageAdapter + TeaCacheCoefficientEstimator on Z-Image-Turbo: + # 70 Parti prompts x {25, 50} steps, 1024x1024, guidance_scale=4.0, 5,110 step + # pairs (R^2 0.69). The previous Qwen-Image placeholder predicted ~1/3 of the + # measured output change and over-skipped steps (see #8270). "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": [ From 50f3b2f4fa7e07fbd904d3eda1f4a19a54db9220 Mon Sep 17 00:00:00 2001 From: Kamila Kuatova Date: Thu, 1 Oct 2026 18:38:30 -0400 Subject: [PATCH 2/3] [Bugfix] Address review on Z-Image TeaCache calibration - Document the measured input rel-L1 range (0.024-0.336), the polynomial's validity window, the GPU-side collection used for calibration, and the rel_l1_thresh=0.2 operating point in the config comment. - Add a CPU test that constructs DataCollectionHook directly and checks the extractor is resolved at init and unknown types fail fast. - Correct two doc claims contradicted by the calibration: borrowed coefficients are not "within 5-10% of optimal", and the 1.5x-2.0x TeaCache speed-up does not hold for few-step distilled models. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kamila Kuatova --- docs/design/feature/teacache.md | 2 +- .../diffusion/cache_acceleration/teacache.md | 2 +- tests/diffusion/cache/test_teacache_config.py | 17 ++++++++++++++++- vllm_omni/diffusion/cache/teacache/config.py | 15 +++++++++++---- 4 files changed, 29 insertions(+), 7 deletions(-) diff --git a/docs/design/feature/teacache.md b/docs/design/feature/teacache.md index 8577cff1f05..d539d388926 100644 --- a/docs/design/feature/teacache.md +++ b/docs/design/feature/teacache.md @@ -260,7 +260,7 @@ 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 diff --git a/docs/user_guide/diffusion/cache_acceleration/teacache.md b/docs/user_guide/diffusion/cache_acceleration/teacache.md index 026b86ec7f9..496694039e0 100644 --- a/docs/user_guide/diffusion/cache_acceleration/teacache.md +++ b/docs/user_guide/diffusion/cache_acceleration/teacache.md @@ -190,5 +190,5 @@ cache_config={"rel_l1_thresh": 0.1} ## 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 diff --git a/tests/diffusion/cache/test_teacache_config.py b/tests/diffusion/cache/test_teacache_config.py index f2bcefc3d36..48a7dc576ab 100644 --- a/tests/diffusion/cache/test_teacache_config.py +++ b/tests/diffusion/cache/test_teacache_config.py @@ -5,8 +5,13 @@ import pytest -from vllm_omni.diffusion.cache.teacache.coefficient_estimator import _MODEL_ADAPTERS, ZImageAdapter +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] @@ -26,3 +31,13 @@ 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") diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 4b5b1e1a300..5acfa0679bd 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -46,10 +46,17 @@ 7.61309272e-01, ], # Z-Image transformer coefficients - # Calibrated with ZImageAdapter + TeaCacheCoefficientEstimator on Z-Image-Turbo: - # 70 Parti prompts x {25, 50} steps, 1024x1024, guidance_scale=4.0, 5,110 step - # pairs (R^2 0.69). The previous Qwen-Image placeholder predicted ~1/3 of the - # measured output change and over-skipped steps (see #8270). + # 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 drops to ~0 by + # x~0.40, so it is only valid inside that range. 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": [ -7.54613422e01, -9.23596156e01, From a4800c2364bf82a2f4f8a82711df0d05c644aad9 Mon Sep 17 00:00:00 2001 From: Kamila Kuatova Date: Thu, 1 Oct 2026 21:21:52 -0400 Subject: [PATCH 3/3] [Bugfix] Z-Image TeaCache: document fit limits, pin coefficients, fix stale doc claims - Config comment: note the zero crossing near x~0.40 (the hook applies the polynomial unclamped), scope the calibration to Z-Image-Turbo, and link the Flux2 placeholder to the tracking issue #8411. - Test: pin the calibrated coefficients with pytest.approx so a recalibration must update the test deliberately. - Docs: scope the remaining unconditional "1.5x-2.0x" claims, the suggested rel_l1_thresh range, and the troubleshooting advice to non-distilled models; make the design-doc adapter example inherit DefaultAdapter to match the registry typing. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kamila Kuatova --- docs/design/feature/teacache.md | 6 +++--- .../diffusion/cache_acceleration/teacache.md | 12 ++++++------ tests/diffusion/cache/test_teacache_config.py | 3 ++- vllm_omni/diffusion/cache/teacache/config.py | 11 ++++++++--- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/docs/design/feature/teacache.md b/docs/design/feature/teacache.md index d539d388926..3f23dc6f4fc 100644 --- a/docs/design/feature/teacache.md +++ b/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. @@ -268,7 +268,7 @@ The polynomial coefficients rescale L1 distances between consecutive modulated i 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). diff --git a/docs/user_guide/diffusion/cache_acceleration/teacache.md b/docs/user_guide/diffusion/cache_acceleration/teacache.md index 496694039e0..4468dbe29ef 100644 --- a/docs/user_guide/diffusion/cache_acceleration/teacache.md +++ b/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,11 +178,11 @@ 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) --- diff --git a/tests/diffusion/cache/test_teacache_config.py b/tests/diffusion/cache/test_teacache_config.py index 48a7dc576ab..a8320039866 100644 --- a/tests/diffusion/cache/test_teacache_config.py +++ b/tests/diffusion/cache/test_teacache_config.py @@ -19,8 +19,9 @@ def test_zimage_has_calibrated_coefficients(): """Z-Image must not fall back to the Qwen-Image placeholder polynomial (#8270).""" zimage = _MODEL_COEFFICIENTS["ZImageTransformer2DModel"] - assert len(zimage) == 5 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 diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 5acfa0679bd..8a189e3b77e 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -50,8 +50,12 @@ # 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 drops to ~0 by - # x~0.40, so it is only valid inside that range. The previous Qwen-Image + # 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), @@ -83,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,