Repository navigation
[Bugfix] Calibrate Z-Image TeaCache coefficients and add ZImageAdapter - #8408
linyueqian merged 8 commits into
Conversation
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 vllm-project#8270
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Kamila Kuatova <kuatova_kamila@list.ru>
|
This PR appears to belong to: docs/design/module/cache_management.md, docs/design/module/diffusion/index.md. Module owners: @Isotr0py @princepride @SamitHuang Routing: @Isotr0py via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @princepride via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @SamitHuang via module of the changed files, module named in the PR description, semantic router, CODEOWNERS @kuatovakamila, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot: three questions on the performance claim@kuatovakamila this PR reads as a performance or value claim:
Before the full evidence checklist, three short questions:
When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution. |
Omni ReviewBot routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
AI assistance: Claude Code helped draft the calibration/benchmark scripts, the unit test, and the PR text, and ran the |
|
This is a correctness fix, not a performance claim: the PR reduces TeaCache's speed-up on Z-Image (2.85× → 1.54× at 50 steps) in exchange for output that matches the uncached image (SSIM 0.56 → 0.90). The "speed-up drops" sentence is explaining that regression, not claiming a gain. Answering the three questions anyway:
|
|
Raising the threshold past 0.20 buys little speed (1.54× → 1.67× at 50 steps) for a steady quality loss, so the PR keeps the global default of 0.2 and does not add a Z-Image entry to |
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Scan:
| Category | Result |
|---|---|
| Tests / verification | 6 finding(s) below |
| Security | no finding reported |
| Docs / comments | 3 finding(s) below |
| Behavior / compatibility | no finding reported |
| Correctness | no finding reported |
Validated:
- [resolved] Concern: removing DataCollectionHook.initialize_hook breaks hook lifecycle — base ModelHook.initialize_hook is identity (vllm_omni/diffusion/hooks/base.py:68-77 'return module'); test_teacache_extractors.py:783/944/1002 call TeaCacheHook.initialize_hook which keeps its own override (teacache/hook.py:91-108). Residual checked: nothing reads hook.extractor_fn before construction; failure for an unknown transformer_type now surfaces at construction (fail-fast, better).
- [claim-verified] PR numbers plug back in: new quartic p(0.09)≈0.282 vs claimed mean predicted output diff 0.278; old Qwen placeholder p(0.09)≈0.078 vs claimed 0.089 — consistent with the reported 0.32 predicted/actual ratio.
- [claim-verified] PR claim 'Both tests fail on main': the diff shows main shipped Z-Image == Qwen-Image placeholder coefficients and no 'ZImage' entry in _MODEL_ADAPTERS, so test 1 fails the inequality assert and test 2 KeyErrors — confirmed from the diff itself.
- [claim-verified] 'Both tests fail on main': plausible by construction (base placeholder equals Qwen's values per the diff; no ZImage adapter key at base), but the base tree is not inspectable from this PR-time checkout - accepted as plausible, not directly verified.
- [validated] 'ZImagePipeline' is a real class: vllm_omni/diffusion/models/z_image/pipeline_z_image.py:164
class ZImagePipeline(nn.Module, ...)and registered in vllm_omni/diffusion/registry.py:63; 'ZImageTransformer2DModel' is the real transformer class (z_image_transformer.py:598), so ZImageAdapter.model_class_name → get_transformer() → _MODEL_COEFFICIENTS keys agree end-to-end. - [validated] get_extractor('ZImageTransformer2DModel') resolves: teacache/extractors.py:1554 '"ZImageTransformer2DModel": extract_zimage_context'.
Of 14 candidates, 6 survive as publishable comments after merging 4 duplicates and dropping 4 non-comment verification/exclusion records (3, 4, 5, 13 — safety analyses and a self-excluded pre-existing Ming observation, none requesting change). The centerpiece is candidate 6 (major): the calibrated quartic's operating point at the retained rel_l1_thresh=0.2 leaves TeaCache nearly inert for Z-Image — verified by hand-evaluating the shipped polynomial (p(0)=0.227, p(0.1)~0.33 vs threshold 0.2) and confirmed by the PR's own 1.08x/SSIM 0.980 benchmark cell — asking for either a tuned per-model default threshold or an explicit quality-for-speed statement, with the new test's assertion updated if the default moves; candidate 11's matching observation is folded in. Supporting comments (all minor): the leftover byte-identical Flux2 placeholder with reachability evidence (8, absorbing 2 and 9), the design-doc row and user-guide speedup claim this PR's own measurements falsify (10, absorbing 1), the missing measured input rel_l1 range in the new calibration comment (0, math checked), and the provenance comment vs the GPU-side collection path deviation admitted in the Test Plan (7). One nit survives: the Bagel entry's factually wrong 'Qwen's coefficients' comment (12). One candidate sub-claim was corrected during verification: 6's '|p(d)|>=0.2 for all d>0.07' is false at x=0.4 (p0.081) and was removed from the rewrite; the threshold argument stands on the verified values and the PR's own mean-per-step figure. The merged review engages the PR's primary change (the calibrated coefficients and their shipped operating point) first, with doc-truth and leftover-placeholder findings as direct consequences.
Verdict: REQUEST CHANGES
CI at
974481c3e7a0(2026-10-01T22:24:16.539275+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
Findings
- [P1] The diff ships the calibrated quartic (constant term 2.27176809e-01, config.py:… —
vllm_omni/diffusion/cache/teacache/config.py
The diff ships the calibrated quartic (constant term 2.27176809e-01, config.py:53-58) but leaves the default threshold at 0.2 — there is no 'ZImageTransformer2DModel' entry in _MODEL_DEFAULT_REL_L1_THRESH (in-tree precedent: '"MiniMaxH3DiTModel": 0.17', config.py:101). The polynomial evaluates to p(0)=0.227 and p(0.1)≈0.33, so a single typical step already meets the threshold: with the runtime accumulating abs(rescaled) each step and resetting on crossing (hook.py:263, 271), the accumulator trips almost every step and the cache nearly never skips — consistent with the PR's own 25-step cell (1.08x speedup, SSIM 0.980). Every Z-Image tea_cache user gets near-zero acceleration by default, with no warning. Fix the shipped operating point in this PR: add a tuned 'ZImageTransformer2DModel' entry to _MODEL_DEFAULT_REL_L1_THRESH selected from the already-run 6x3x2 SSIM/speedup sweep, or state in the PR and the config comment that at 0.2 the cache intentionally trades all speedup for quality; if the default changes, updateassert config.rel_l1_thresh == 0.2(tests/diffusion/cache/test_teacache_config.py:22).
Evidence: In diff — vllm_omni/diffusion/cache/teacache/config.py:53-58: "ZImageTransformer2DModel": [ ... 2.27176809e-01,; tests/diffusion/cache/test_teacache_config.py:22: assert config.rel_l1_thresh == 0.2. Unchanged by this diff, present in the PR-time tree — config.py:99-102: _DEFAULT_REL_L1_THRESH = 0.2 / _MODEL_DEFAULT_REL_L1_THRESH = { / "MiniMaxH3DiTModel": 0.17, / } (no Z-Image entry); config.py:133: threshold = _MODEL_DEFAULT_REL_L1_THRESH.get(self.transformer_type, _DEFAULT_REL_L1_THRESH) resolves Z-Image to 0.2; hook.py:86: self.rescale_func = np.poly1d(config.coefficients) applies the shipped constants directly to the raw rel-distance; hook.py:262-263: rescaled_distance = float(self.rescale_func(rel_distance)) / state.accumulated_rel_l1_distance += abs(rescaled_distance); hook.py:271: state.accumulated_rel_l1_distance = 0.0 # Reset accumulator. Hand evaluation of np.poly1d([-75.4613422, -92.3596156, 57.5318402, -3.76790311, 0.227176809]): p(0)=0.227, p(0.1)=0.326, p(0.12)=0.428, p(0.45)=-1.33, p(0.5)=-3.54 — p(0)>0.2 means even a near-zero input diff rescales above the default threshold, so a single computed step trips the accumulator and the cache skips almost nothing. coefficient_estimator.py:176: return np.polyfit(x, y, poly_order).tolist() confirms coefficients map input rel-L1 diffs to predicted output rel-L1 diffs.
- [P2] config.py:49-52 records the calibration setup (70 Parti prompts x {25,50} steps… —
vllm_omni/diffusion/cache/teacache/config.py
config.py:49-52 records the calibration setup (70 Parti prompts x {25,50} steps, 1024x1024, guidance_scale=4.0, R^2 0.69) but not the measured input rel_l1 range, unlike the sibling HunyuanImage3 entry (config.py:70: 'Maps time_embed rel_l1 (range ~0.12-0.54) ...'). Fitted from the new coefficients via np.poly1d (hook.py:86), the quartic peaks at x≈0.28 (p≈1.19) and collapses to ≈0.08 at x=0.40, crossing zero near x≈0.41 — beyond the calibration range it under-predicts step-to-step change, so the accumulated distance grows too slowly and TeaCache over-skips, the #8270 failure mode (hook.py:263's abs() masks the sign flip but not the collapse). Please add the observed min-max input rel_l1 range and the settings it was measured under to the comment so readers can judge extrapolation to other resolutions/guidance scales.
Evidence: config.py:49-52 (this diff) # 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). — no input rel_l1 range recorded. Unchanged by this diff, present in the PR-time tree — config.py:70 # Maps time_embed rel_l1 (range ~0.12-0.54) to output rel_l1 (range ~0.01-0.53). (sibling convention); hook.py:86 self.rescale_func = np.poly1d(config.coefficients) (highest-degree-first ordering — makes p(0.28)≈1.191 the peak and p(0.40)≈0.082 the collapse, zero-crossing ≈0.41); hook.py:262-263 rescaled_distance = float(self.rescale_func(rel_distance)) / state.accumulated_rel_l1_distance += abs(rescaled_distance) (abs() hides the sign flip, not the collapse).
Suggestion: # 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). Measured input rel_l1 range: -.
# Valid only within that range: the quartic turns over near x~0.3 and
# under-predicts beyond it (see #8270).
# The previous Qwen-Image placeholder predicted ~1/3 of the
# measured output change and over-skipped steps (see #8270).
- [P2] This diff removes exactly this borrowed-placeholder pattern for Z-Image, and it… —
vllm_omni/diffusion/cache/teacache/config.py
This diff removes exactly this borrowed-placeholder pattern for Z-Image, and its own table entry says the borrowed Qwen polynomial predicted ~1/3 of the measured output change (#8270) — yet the identical placeholder remains one block below for Flux2Transformer2DModel (config.py:79-86: the same [-4.5e02, 2.8e02, -4.5e01, 3.2e00, -2e-02] under '# Copied from Qwen-Image, need to be tuned specifically for Flux2 in future'), and it is reachable: Flux2Pipeline has no entry in CUSTOM_TEACACHE_ENABLERS (backend.py:123-129), so TeaCacheBackend.enable() takes the generic branch and resolves transformer_type fromtransformer.__class__.__name__(backend.py:173); Flux2Transformer2DModel is a real served class (vllm_omni/diffusion/models/flux2/flux2_transformer.py:756, listed for Flux2Pipeline / black-forest-labs/FLUX.2-dev in docs/user_guide/diffusion/cpu_offload.md:43). The Flux2Klein entry reusing FLUX.1's coefficients (config.py:18-26) is a second live borrow, though its enabler passes the type explicitly (backend.py:109). State in the PR description why Flux2 is left as-is, or file a tracking issue for calibrating the Flux2 entries — mirroring how the new Z-Image comment cites #8270.
Evidence: All decisive lines are OUTSIDE the diff (pre-existing in the PR-time tree; the diff only rewrites the Z-Image block at config.py:48-59). config.py:79-86 verbatim: # Copied from Qwen-Image, need to be tuned specifically for Flux2 in future / "Flux2Transformer2DModel": [ / -4.50000000e02, / 2.80000000e02, / -4.50000000e01, / 3.20000000e00, / -2.00000000e-02, — byte-identical to the Z-Image placeholder this diff deletes. backend.py:173 (generic branch of TeaCacheBackend.enable(), reached because CUSTOM_TEACACHE_ENABLERS at backend.py:123-129 lists Flux2KleinPipeline but not Flux2Pipeline): transformer_type = transformer.__class__.__name__, then backend.py:178-179 teacache_config = TeaCacheConfig( / transformer_type=transformer_type,. Class exists: vllm_omni/diffusion/models/flux2/flux2_transformer.py:756 class Flux2Transformer2DModel(nn.Module):. Docs: docs/user_guide/diffusion/cpu_offload.md:43 | Flux2Pipeline | \black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | .... Second borrow: config.py:18-20 # Flux2 Klein transformer coefficients/ # Same as FLUX.1 (similar dual-stream architecture)/ "Flux2Klein": [with values identical to FluxTransformer2DModel at config.py:11-17; its enabler backend.py:109 transformer_type="Flux2Klein",`.
- [P2] The config.py:49 comment says the coefficients were 'Calibrated with ZImageAdap… —
vllm_omni/diffusion/cache/teacache/config.py
The config.py:49 comment says the coefficients were 'Calibrated with ZImageAdapter + TeaCacheCoefficientEstimator', but the PR's own Test Plan states collection computed the rel-L1 distances on the GPU per step instead of storing activations (~13 GB of host RAM per 50-step 1024x1024 CFG trajectory), while the committed DataCollectionHook stores every step as float32 host numpy (coefficient_estimator.py:41-42:ctx.hidden_states.detach().float().cpu().numpy()thenself.current_trajectory.append(...); stop_collection() feeds the full stored trajectory to np.polyfit). The in-repo path as documented is not the path that produced these coefficients — please note the GPU-resident deviation in the comment (or land the GPU-resident collector) so the next person recalibrating Z-Image can reproduce it.
Evidence: Tree-verified anchors (PR-time head 974481c): vllm_omni/diffusion/cache/teacache/config.py:49 # Calibrated with ZImageAdapter + TeaCacheCoefficientEstimator on Z-Image-Turbo: — added by this diff, attributes calibration to the in-repo estimator. vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:41-42 (lines 34/41-42 inside this diff's DataCollectionHook.new_forward) model_output_cpu = ctx.hidden_states.detach().float().cpu().numpy() / self.current_trajectory.append((modulated_input_cpu, model_output_cpu)), with unchanged-by-this-diff coefficient_estimator.py:48-49 def stop_collection(self) -> list[tuple[np.ndarray, np.ndarray]]: / return list(self.current_trajectory) and :176 return np.polyfit(x, y, poly_order).tolist() — proving the committed path stores the full float32 host trajectory and fits from it, i.e. the store-everything 'stock hook' the PR Test Plan says was NOT used (quote from provided PR-body evidence, not independently fetchable: 'Collection computed the rel-L1 distances on the GPU per step instead of storing activations (the stock hook would need ~13 GB of host RAM per 50-step 1024² CFG trajectory)').
Suggestion: # Calibrated with ZImageAdapter + TeaCacheCoefficientEstimator on Z-Image-Turbo
# (collection computed per-step rel-L1 on the GPU; the stock DataCollectionHook
# stores full CPU trajectories, ~13 GB per 50-step 1024x1024 CFG run):
- [P2] This PR is the change that falsifies the design doc's borrowing row: its own ad… —
docs/design/feature/teacache.md
This PR is the change that falsifies the design doc's borrowing row: its own added comment in config.py:51-52 records that the borrowed Qwen-Image polynomial predicted ~1/3 of the measured output change on Z-Image and over-skipped steps (#8270) — not 'Within 5-10% of optimal' (docs/design/feature/teacache.md:263: '| Using defaults from similar model | Within 5-10% of optimal | Low |'). Per the PR's measurement table, the placeholder hit 0.32x of actual output change (R^2 -0.18), and the calibrated default yields 1.08x at 25 steps versus the user guide's '1.5x-2.0x speedup with optimized defaults' (docs/user_guide/diffusion/cache_acceleration/teacache.md:193). Update both claims in this PR (or scope them to the models where they hold) so the docs match shipped behavior.
Evidence: Decisive lines, quoted verbatim. (1) Added BY this diff: vllm_omni/diffusion/cache/teacache/config.py:51-52 'pairs (R^2 0.69). The previous Qwen-Image placeholder predicted ~1/3 of the' / 'measured output change and over-skipped steps (see #8270).' — the PR's own record that borrowing under-predicted ~3x. (2) Unchanged by this diff, present in the PR-time tree: docs/design/feature/teacache.md:263 '| Using defaults from similar model | Within 5-10% of optimal | Low |' — contradicted by (1) for Z-Image. (3) Unchanged by this diff, present in the PR-time tree: docs/user_guide/diffusion/cache_acceleration/teacache.md:193 '1. ✅ Enable TeaCache - Set cache_backend="tea_cache" to get 1.5x-2.0x speedup with optimized defaults'. Neither doc file appears in the PR diff.
Suggestion: | Using defaults from similar model | Unpredictable; Z-Image under-predicted ~3x (#8270) - calibrate | Low |
- [P2] This PR highlights a gap in the design doc's calibration checklist (docs/design… —
docs/design/feature/teacache.md:371
This PR highlights a gap in the design doc's calibration checklist (docs/design/feature/teacache.md, unchanged by this diff): the Data Statistics Guide (lines 366-371) gates only coefficient magnitude (line 371 accepts -1e6..1e6, so the shipped max |c|=92.36 passes trivially) and only warns on correlation (line 370); it defines no R² or fit-quality acceptance gate, so the R²=0.69 Z-Image fit passes every documented check. Since this PR does not touch the doc, consider either extending it here or filing a follow-up: add an explicit fit-quality threshold (R² or equivalent) to the Data Statistics Guide so the next weak calibration is mechanically rejected instead of passing on magnitude alone.
Evidence: Unchanged by this diff, present in the PR-time tree — docs/design/feature/teacache.md:371 | **Coefficient magnitude** | -1e6 to 1e6 | > 1e8: numerical instability | and :370 | **Output Diffs (y)** | Should correlate with x | No correlation: check extractor |; whole-doc grep for R^2/R²/correlat/calibrat matches only these checklist lines, so no R² gate exists. Shipped coefficients (in-diff, config.py) -7.54613422e01, -9.23596156e01, 5.75318402e01, -3.76790311e00, 2.27176809e-01 give max |c|=92.36 < 1e6. PR diff contains only tests/diffusion/cache/test_teacache_config.py, coefficient_estimator.py, config.py — the doc is not edited.
- [P2] This PR makes Z-Image tea_cache actually work while its in-tree Ming twins stay… —
vllm_omni/diffusion/cache/teacache/config.py:140
This PR makes Z-Image tea_cache actually work while its in-tree Ming twins stay crash-only: MingZImageTransformer2DModel (vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py:21) and MingImageTransformer2DModel (vllm_omni/diffusion/models/ming_image/transformer.py:13) subclass ZImageTransformer2DModel but keep distinct class names, and the backend derives transformer_type from transformer.class.name (teacache/backend.py:173), so a Ming pipeline with cache_backend="tea_cache" hits the loud KeyError at config.py:140-144 rather than silently inheriting the new coefficients — pre-existing and not worsened here. Since the premise of this PR is that Z-Image users were being served borrowed coefficients, please fold the Ming gap into the same tracking effort as the Flux2 entries (or note the unsupported combination next to the new calibration comment at config.py:49) so the full sibling picture — Flux2/Flux2Klein silent borrows, Ming loud failure — lands in one place.
Evidence: Unchanged by this diff (the diff only rewrites config.py:48-59 coefficient values/comment and the SPDX headers), present in the PR-time tree: vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py:21 class MingZImageTransformer2DModel(ZImageTransformer2DModel):; vllm_omni/diffusion/models/ming_image/transformer.py:13 class MingImageTransformer2DModel(ZImageTransformer2DModel):; vllm_omni/diffusion/cache/teacache/backend.py:173 transformer_type = transformer.__class__.__name__ feeding backend.py:178-179 teacache_config = TeaCacheConfig( / transformer_type=transformer_type,; vllm_omni/diffusion/cache/teacache/config.py:140-141 if self.transformer_type not in _MODEL_COEFFICIENTS: / raise KeyError(. _MODEL_COEFFICIENTS (config.py:9-97) contains no Ming key, so the distinct Ming class names miss the lookup loudly instead of inheriting the Z-Image entry.
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
| @@ -25,13 +25,9 @@ class DataCollectionHook(ModelHook): | |||
| def __init__(self, transformer_type: str): | |||
There was a problem hiding this comment.
[P2] The only runtime-machinery change in this diff ships with no test selecting it:…
The only runtime-machinery change in this diff ships with no test selecting it: DataCollectionHook now resolves its extractor eagerly in init (coefficient_estimator.py:28) and its initialize_hook override is removed, so it falls back to base ModelHook.initialize_hook (hooks/base.py:68). No test in the repo constructs DataCollectionHook — the sole instantiation is TeaCacheCoefficientEstimator.init (coefficient_estimator.py:200) — and the only cache-hook tests that call initialize_hook exercise TeaCacheHook, whose override remains (test_teacache_extractors.py:777-783; hook.py:91); the SP and offloader hook tests cover unrelated hooks. The PR's new tests cover only TeaCacheConfig.post_init and the _MODEL_ADAPTERS registry. Add a core_model/cpu test that constructs DataCollectionHook(transformer_type=...) directly — going through TeaCacheCoefficientEstimator would load a real pipeline (coefficient_estimator.py:198) — and asserts extractor_fn is resolved at init with no None window before new_forward, so the path this PR changed is exercised by CI.
Evidence: vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:28 self.extractor_fn = get_extractor(transformer_type) — inside init (lines 25-29), changed by this diff from None-until-initialize_hook. vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:20-49 — class DataCollectionHook defines only init/new_forward/start_collection/stop_collection; no initialize_hook override. Unchanged by this diff, present in the PR-time tree: vllm_omni/diffusion/hooks/base.py:68 def initialize_hook(self, module: nn.Module) -> nn.Module: (the base fallback now used). vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:200 self.hook = DataCollectionHook(self.transformer_type) — sole instantiation in the repo; its surrounding init line 198 self.pipeline = adapter.load_pipeline(model_path, device, dtype) shows the estimator path cannot serve as a CPU unit test. Unchanged by this diff, present in the PR-time tree: tests/diffusion/cache/test_teacache_extractors.py:777-783 hook = TeaCacheHook( ... hook.initialize_hook(minimax_h3_module) (likewise lines 944 and 1002), and vllm_omni/diffusion/cache/teacache/hook.py:91 def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: (TeaCacheHook's override remains). Repo-wide grep confirms DataCollectionHook and TeaCacheCoefficientEstimator appear in no file under tests/ — the no-coverage claim holds; note the draft's 'only hook-lifecycle test in the repo' is overbroad since tests/diffusion/distributed/test_sp_plan_hooks.py:734 and tests/diffusion/offloader/test_layerwise_backend.py:97 also call initialize_hook on unrelated hooks.
Omni ReviewBot: finding feedback[p1] The diff ships the calibrated quartic (constant term 2.27176809e-01, config.py:… — The diff ships the calibrated quartic (constant term 2.27176809e-01, config.py:53-58) but leaves the default threshold at 0.2 — there is no 'ZImageTransformer2DModel' entry in _MODEL_DEFAULT_REL_L1_THRESH (in-tree precedent: '"MiniMaxH3DiTModel": 0.17', config.py:101). The polynomial evaluates to p(0)=0.227 and p(0.1)≈0.33, so a single typical step already meets the threshold: with the runtime accumulating abs(rescaled) each step and resetting on crossing (hook.py:263, 271), the accumulator trips almost every step and the cache nearly never skips — consistent with the PR's own 25-step cell (1.08x speedup, SSIM 0.980). Every Z-Image tea_cache user gets near-zero acceleration by default, with no warning. Fix the shipped operating point in this PR: add a tuned 'ZImageTransformer2DModel' entry to _MODEL_DEFAULT_REL_L1_THRESH selected from the already-run 6x3x2 SSIM/speedup sweep, or state in the PR and the config comment that at 0.2 the cache intentionally trades all speedup for quality; if the default changes, update Evidence: In diff — vllm_omni/diffusion/cache/teacache/config.py:53-58: This finding appears in the bot's COMMENT review, but GitHub could not place it as an inline diff comment. If you are the PR author and disagree, react 👎 to this comment. The disagreement will be shown to the maintainer; it does not approve or merge the PR. |
- 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 <noreply@anthropic.com> Signed-off-by: Kamila Kuatova <kuatova_kamila@list.ru>
|
Thanks for the detailed review. Pushed a follow-up commit; responses by finding: [P1] Threshold / operating point. Checked the claim against the calibration data rather than the polynomial's constant term. The constant (p(0) = 0.227) is outside the measured range (min input rel-L1 is 0.024, where p = 0.17), and on the 5,110 measured steps 53% rescale below 0.2 (the actual output change is below 0.2 on 56% of them), so the cache does skip: that's the 1.54× at 50 steps. At 25 steps only 37% of steps have an actual output change under 0.2 because the distilled model moves >20% per step early on, which is why 25-step runs get ~1.1×. That is the model, not the fit; the previous 1.96× there was wrong images. I went with your second option: [P2] Input range in the comment. Added: measured input rel-L1 range 0.024–0.336 (mean 0.075), with a note that the quartic peaks near x≈0.28 and drops to ~0 by x≈0.40, so it's only valid inside the measured range. Your extrapolation concern is real; 95% of measured steps are below 0.17. [P2] Provenance of the collector. The comment now says the per-step rel-L1 was computed on the GPU rather than via the stock CPU-trajectory [P2] Doc claims. Updated both lines in this PR: the design-doc "within 5-10% of optimal" row now says borrowed coefficients are unreliable (with the Z-Image numbers), and the user guide's "1.5x-2.0x" is scoped to 50-step runs with the few-step caveat. Note [P2] Flux2 / Flux2Klein borrows and the Ming subclasses. Added to the description as out of scope with a follow-up issue to track all three together (Flux2 silent borrow, Flux2Klein borrow, Ming [P2] R² gate in the design doc. Left for the follow-up rather than inventing a threshold from one calibration. [P2, inline] Local checks on the new head: pre-commit (ruff, ruff-format, mypy-3.10, SPDX, typos, CI marks) pass except the pre-existing markdownlint errors above; |
Omni ReviewBot attempt recordReview attempt ended as stale. |
Omni ReviewBot attempt recordReview attempt ended as stale. |
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Scan:
| Category | Result |
|---|---|
| Tests / verification | 2 finding(s) below |
| Security | no finding reported |
| Docs / comments | 3 finding(s) below |
| Behavior / compatibility | no finding reported |
| Correctness | 1 finding(s) below |
Validated:
- [resolved] initialize_hook-removal lifecycle concern (raised by review bot, resolved): base ModelHook.initialize_hook is identity — vllm_omni/diffusion/hooks/base.py:68-77
return module; TeaCacheHook keeps its own lazy override (teacache/hook.py:91, :103self.extractor_fn = get_extractor(self.config.transformer_type)). Residual checked: test_teacache_extractors.py:783/944/1002 never touch DataCollectionHook (repo-wide grep), nothing reads hook.extractor_fn before init, and the runtime TeaCacheHook lazy path is untouched by this diff. - [claim-verified] PR's 'Ming subclasses raise KeyError' statement: MingZImageTransformer2DModel (ming_zimage_transformer.py:21) and MingImageTransformer2DModel (ming_image/transformer.py:13) have their own class names -> KeyError at config.py:147-151, wrapped as a loud ValueError at backend.py:183-189; pre-existing, author tracks as follow-up — not re-raised
- [claim-verified] Doc corrections truthful: design doc '~3x' matches predicted/actual 0.32; user guide '~1.1x' matches the 25-step 1.08x sweep; 'on 50-step runs' matches 1.54x
- [validated] coefficient_estimator.py:128-143 registry typing accurate: BagelAdapter (:92), Flux2Adapter (:110), LongCatAdapter (:116), StableAudioAdapter (:125), ZImageAdapter (:131) all subclass DefaultAdapter, so
dict[str, type[DefaultAdapter]]holds. - [validated] "ZImagePipeline" is loadable: class at vllm_omni/diffusion/models/z_image/pipeline_z_image.py:164, registered in vllm_omni/diffusion/registry.py:63, documented at docs/models/supported_models.md:36 — DefaultAdapter.load_pipeline's model_class_name kwarg resolves for the new adapter.
- [validated] initialize_hook removal safe: TeaCacheHook keeps its own override (hook.py:91-108); all other initialize_hook references are unrelated hook classes; nothing reads DataCollectionHook.extractor_fn before init (grep: only coefficient_estimator.py:28,32 + new test)
Of 11 candidates, 4 survive as published comments (all minor), 3 merge as duplicates, and 4 drop because their own producing lens had already cancelled them (excluded/resolved/no-issue self-verdicts, not refuted findings). The kept set stays on the PR's primary change: (1) the shipped quartic crosses zero at x≈0.404 just above its calibrated input range while hook.py applies it unclamped, so out-of-envelope steps can silently rescale to ≈0 — document or clamp (config.py:53); (2) the entry is keyed by exact class name, so non-Turbo Z-Image checkpoints silently inherit Turbo-calibrated values — add a scoping sentence (config.py:60); (3) the doc correction is incomplete — the same page's intro, Best Practices and Troubleshooting (raise rel_l1_thresh to 0.8, 35+ steps) plus the design-doc mirror still assert what the PR's own sweep and config warning contradict (anchored at teacache.md:193, absorbing candidates 4 and 9); and (4) the declared follow-up issue for the Flux2 borrowed coefficients and the Ming KeyError is linked nowhere (config.py:86, absorbing candidate 5). Nothing blocks the merge; the Troubleshooting-advice and zero-crossing comments are the two most consequential.
Checked, no defect found:
tests/diffusion/cache/test_teacache_config.py:16— Check found the code correct: the GPU 'Diffusion · Cache Test' lane selects-m 'core_model and cuda'(test-ready.yml:63) so it never runs this cpu-marked file, but the CPU diffusion sweepstests/diffusion -m 'core_model and cpu'collect it in the CUDA merge lane (test-merge.yml:23), the AMD merge lane (test-amd-merge.yml:46 — among cache tests it ignores only test_teacache_extractors.py), and the ready pipeline's 'Simple · Diffusion Test' (test-ready.yml:31) — so a regression in these defaults fails CI at both ready and merge time; no change needed.
Verdict: COMMENT
CI at
4f253f6386ad(2026-10-02T00:48:51.904646+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
Findings
- [P2] The diff's own calibration comment restricts the shipped quartic (config.py:61-… —
vllm_omni/diffusion/cache/teacache/config.py
The diff's own calibration comment restricts the shipped quartic (config.py:61-65) to the measured input range 0.024-0.336 because it 'peaks near x~0.28 and drops to0 by x0.40' (config.py:53-54) — but the runtime applies it unclamped. hook.py, unchanged by this diff, buildsnp.poly1d(config.coefficients)(hook.py:86), evaluates it on the raw per-step rel_l1 (hook.py:259, hook.py:262), accumulatesabs(rescaled_distance)(hook.py:263), and compares to rel_l1_thresh (hook.py:268). Evaluating the shipped coefficients: p(0.40)≈+0.08, p(0.41)≈-0.14, so the fit crosses zero at x≈0.404. A step whose input rel_l1 lands near the crossing — reachable just outside the calibrated envelope (other guidance scales, resolutions, or non-Turbo Z-Image checkpoints sharing this entry) — rescales to ≈0, contributes ≈0 to the accumulator, and is silently served from cache despite being one of the largest steps. State this out-of-range failure mode in the config comment, or add a clamp/floor on the rescale input in the hook (noting a clamp there affects all TeaCache models).
Evidence: In-diff: config.py:53-54 # 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.; config.py:61-65 coefficients -7.54613422e01, -9.23596156e01, 5.75318402e01, -3.76790311e00, 2.27176809e-01. Unchanged by this diff, present in the PR-time tree: hook.py:86 self.rescale_func = np.poly1d(config.coefficients) (no clamp); hook.py:259 rel_distance = (mean_diff / (mean_prev + 1e-8)).item(); hook.py:262 rescaled_distance = float(self.rescale_func(rel_distance)); hook.py:263 state.accumulated_rel_l1_distance += abs(rescaled_distance); hook.py:268-269 if state.accumulated_rel_l1_distance < rel_l1_thresh: / return False # Use cache. Evaluating np.poly1d of the shipped coefficients: p(0.40)≈+0.082, p(0.41)≈-0.144 → zero crossing at x≈0.404, so a step with rel_distance≈0.404 adds ≈0 to the accumulator and is reused from cache.
Suggestion: # Outside the calibrated range the fit must not be trusted: it crosses zero
near x~0.404, so a step with input distance ~0.40 rescales to ~0 and is
silently skipped despite being the largest step. Base (non-Turbo) Z-Image
checkpoints share this entry but were not calibrated.
- [P2] These coefficients are keyed by the shared transformer class name (config.py:60… —
vllm_omni/diffusion/cache/teacache/config.py
These coefficients are keyed by the shared transformer class name (config.py:60) and documented as 'Calibrated on Z-Image-Turbo' (config.py:49), but the runtime resolves the entry solely viatransformer.__class__.__name__(backend.py:173 → config.py:152_MODEL_COEFFICIENTS[self.transformer_type]), so any checkpoint whose transformer is this class — e.g. a future base Z-Image, not only the Turbo checkpoint the fit was computed on — silently inherits these values. In-tree that currently covers every user (docs/models/supported_models.md:36 lists only Tongyi-MAI/Z-Image-Turbo), so this is just future-proofing: add one scoping sentence to the provenance comment, e.g.# Validated only on Z-Image-Turbo; other checkpoints that reach TeaCache via this shared class name (e.g. base Z-Image) inherit these values uncalibrated.
Evidence: config.py:60 (this diff): "ZImageTransformer2DModel": [ and config.py:49 (this diff): # Calibrated on Z-Image-Turbo with ZImageAdapter + the TeaCacheCoefficientEstimator — the fit is Turbo-specific but keyed by class name. Unchanged by this diff, present in the PR-time tree: backend.py:173 transformer_type = transformer.__class__.__name__ feeding config.py:152 self.coefficients = _MODEL_COEFFICIENTS[self.transformer_type] — runtime lookup is by class name only, so any checkpoint of this class inherits the values. Unchanged by this diff, present in the PR-time tree: docs/models/supported_models.md:36 | "ZImagePipeline" | Z-Image | "Tongyi-MAI/Z-Image-Turbo" | ✅︎ ... — only the Turbo checkpoint is listed, and it is the sole Z-Image mention in that file.
Suggestion: # Validated only on Z-Image-Turbo; other checkpoints that reach TeaCache via
this shared class name (e.g. base Z-Image) inherit these values uncalibrated.
- [P2] The diff corrects only this page's Summary (line 193), but the rest of the page… —
docs/user_guide/diffusion/cache_acceleration/teacache.md
The diff corrects only this page's Summary (line 193), but the rest of the page still states what the PR's own measurements contradict: the intro (line 18) still promises an unconditional '1.5x-2.0x speedup with minimal quality loss'; therel_l1_threshtable row (line 122) still says 'Suggested range: 0.1-0.8'; Best Practices (line 152) still cites 'Scenarios where 1.5-2x speedup is valuable'; and Troubleshooting 'Limited Speedup' (lines 181-185) tells users seeing <1.3x to raiserel_l1_threshto 0.8 and use 35+ steps — while the shipped config comment warns 0.25-0.30 already 'buys ~10% more speed for a visible quality drop' (config.py:58-59), the fit is 'only valid inside' the measured rel_l1 range 0.024-0.336 (config.py:52-54) which 0.8 lies far outside, and few-step distilled models cannot add steps. The same stale claim is mirrored in docs/design/feature/teacache.md:23 and :401. Align intro, config table, Best Practices and Troubleshooting with the calibrated envelope: scope the threshold-raise/35+ steps advice to non-distilled models within the validated range and cross-reference the new few-step Summary note.
Evidence: docs/user_guide/diffusion/cache_acceleration/teacache.md:183 (unchanged by diff, present in PR-time tree) cache_config={"rel_l1_thresh": 0.8} with :181 1. Increase the threshold to enable more aggressive caching: and :185 2. Ensure you're using sufficient inference steps (35+ recommended); :18 (unchanged) providing **1.5x-2.0x speedup** with minimal quality loss; :122 (unchanged) Suggested range: 0.1-0.8; :152 (unchanged) - Scenarios where 1.5-2x speedup is valuable; :193 is the Summary line this diff rewrites. Contradicting calibration added by this diff at vllm_omni/diffusion/cache/teacache/config.py:52-59: 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. ... raising the / threshold to 0.25-0.30 buys ~10% more speed for a visible quality drop. Mirrored stale claim outside the diff: docs/design/feature/teacache.md:23 It provides **1.5x-2.0x speedup** with minimal quality loss. and :401 2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x).
Suggestion: 1. Increase the threshold to enable more aggressive caching (non-distilled
models only — on few-step distilled models like Z-Image-Turbo a higher
threshold mostly costs quality, see the Summary note):
cache_config={"rel_l1_thresh": 0.3}- Ensure you're using sufficient inference steps (35+ recommended for
non-distilled models; few-step distilled models cannot add steps)
- [P2] This diff calibrates Z-Image but leaves the other self-acknowledged borrowed-co… —
vllm_omni/diffusion/cache/teacache/config.py
This diff calibrates Z-Image but leaves the other self-acknowledged borrowed-coefficient entries in place: Flux2Transformer2DModel still ships Qwen-Image's polynomial verbatim (config.py:86-92, '# Copied from Qwen-Image, need to be tuned specifically for Flux2 in future' — values identical to config.py:31-35), Flux2Klein still ships FLUX.1's (config.py:20-26), and MingZImageTransformer2DModel (ming_zimage_transformer.py:21) has no _MODEL_COEFFICIENTS entry, so cache_backend="tea_cache" on a Ming pipeline raises ValueError (KeyError from config.py:148, wrapped at backend.py:185-189; no Ming entry in CUSTOM_TEACACHE_ENABLERS at backend.py:123-129). If the PR description defers these to a follow-up issue, create and link that issue and reference its number at config.py:86 so the declared deferrals survive the merge.
Evidence: Unchanged by this diff, present in the PR-time tree — vllm_omni/diffusion/cache/teacache/config.py:86-87: ' # Copied from Qwen-Image, need to be tuned specifically for Flux2 in future' / ' "Flux2Transformer2DModel": [' with values at :88-92 identical to QwenImageTransformer2DModel at :31-35; config.py:19 ' # Same as FLUX.1 (similar dual-stream architecture)' over the Flux2Klein entry at :20-26 identical to FluxTransformer2DModel at :11-16; ming_zimage_transformer.py:21 'class MingZImageTransformer2DModel(ZImageTransformer2DModel):' with no Ming key in _MODEL_COEFFICIENTS (config.py:9-104) and no Ming entry in backend.py:123-129 CUSTOM_TEACACHE_ENABLERS; backend.py:185-189: 'raise ValueError(
f"Invalid TeaCache configuration: {e}. "' wrapping the KeyError raised at config.py:148 ('raise KeyError(
f"Cannot find coefficients for {self.transformer_type}. "')
Suggestion: # Copied from Qwen-Image, still uncalibrated — tracked in #
- [P2] The PR tightens
_MODEL_ADAPTERStodict[str, type[DefaultAdapter]](coeffic… —docs/design/feature/teacache.md:271
The PR tightens_MODEL_ADAPTERStodict[str, type[DefaultAdapter]](coefficient_estimator.py:137), but the design doc's adapter recipe (docs/design/feature/teacache.md:271-301) still teaches a standaloneclass YourModelAdapter:with only@staticmethodmethods and noDefaultAdapterbase — following it produces a class that fails the new registry annotation. All five real adapters, including the newZImageAdapter(coefficient_estimator.py:131), subclassDefaultAdapter. Update the example to inherit fromDefaultAdapterso the documented recipe matches the typing this diff introduces.
Evidence: Unchanged by this diff (doc diff edits only line 263), present in the PR-time tree: docs/design/feature/teacache.md:271 class YourModelAdapter: with :274 @staticmethod (no DefaultAdapter base anywhere in the block) and :301 _MODEL_ADAPTERS["YourModel"] = YourModelAdapter. Introduced by this diff: vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:137 _MODEL_ADAPTERS: dict[str, type[DefaultAdapter]] = {, while :131 class ZImageAdapter(DefaultAdapter): and :92 class BagelAdapter(DefaultAdapter):, :110 class Flux2Adapter(DefaultAdapter):, :116 class LongCatAdapter(DefaultAdapter):, :125 class StableAudioAdapter(DefaultAdapter): show every real adapter inherits DefaultAdapter — the doc's non-inheriting example violates the new value type.
- [P3] This regression test pins the calibrated Z-Image polynomial only as '5 elements… —
tests/diffusion/cache/test_teacache_config.py
This regression test pins the calibrated Z-Image polynomial only as '5 elements' and 'different from Qwen-Image's placeholder' — any other wrong 5-vector (a bad re-fit, a botched conflict resolution) would still pass, while the shipped values are the entire fix for #8270. Only an exact revert to the old placeholder trips the != check, since the placeholder was verbatim Qwen-Image's entry. Pin the calibrated constants themselves so an intentional recalibration is forced to touch this test: assert zimage == pytest.approx([-7.54613422e01, -9.23596156e01, 5.75318402e01, -3.76790311e00, 2.27176809e-01]) (values from vllm_omni/diffusion/cache/teacache/config.py:60-65).
Evidence: In-diff: tests/diffusion/cache/test_teacache_config.py:21-23 zimage = _MODEL_COEFFICIENTS["ZImageTransformer2DModel"] / assert len(zimage) == 5 / assert zimage != _MODEL_COEFFICIENTS["QwenImageTransformer2DModel"] — the only value assertions; line 26 assert config.coefficients == zimage compares the config back to the same dict entry (circular), so no calibrated constant is pinned. The shipped values the comment cites are at vllm_omni/diffusion/cache/teacache/config.py:60-65 (line range verified): "ZImageTransformer2DModel": [ / -7.54613422e01, / -9.23596156e01, / 5.75318402e01, / -3.76790311e00, / 2.27176809e-01,. Unchanged-context check in the same file: config.py:31-35 Qwen-Image's entry is -4.50000000e02, 2.80000000e02, -4.50000000e01, 3.20000000e00, -2.00000000e-02, identical to the pre-PR Z-Image placeholder shown in the diff's removed lines, so an exact revert would fail the != assertion — the draft's 'merge revert' example is slightly overstated but its core claim holds.
Suggestion: assert zimage == pytest.approx([-7.54613422e01, -9.23596156e01, 5.75318402e01, -3.76790311e00, 2.27176809e-01])
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
… 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 vllm-project#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 <noreply@anthropic.com> Signed-off-by: Kamila Kuatova <kuatova_kamila@list.ru>
|
Addressed the second-pass comments in one more commit:
Gates (ruff, mypy-3.10, SPDX, typos, CI marks) pass; |
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Scan:
| Category | Result |
|---|---|
| Tests / verification | 2 finding(s) below |
| Security | no finding reported |
| Docs / comments | 3 finding(s) below |
| Behavior / compatibility | no finding reported |
| Correctness | no finding reported |
Validated:
- [resolved] Concern: removing DataCollectionHook.initialize_hook could break hook lifecycle or leave extractor_fn unset — base ModelHook.initialize_hook is identity and TeaCacheHook keeps its own override (hook.py:91-108); residual checked: the only constructor caller is TeaCacheCoefficientEstimator (coefficient_estimator.py:200), whose adapter transformer names (Bagel, StableAudioDiTModel, Flux2Transformer2DModel, LongCatImageTransformer2DModel, ZImageTransformer2DModel) are all in EXTRACTOR_REGISTRY (extractors.py:1549-1564), so eager get_extractor in init changes only failure timing, not behavior; unknown types now fail with the ValueError at extractors.py:1630-1635 before any forward, pinned by the new test.
- [claim-verified] 'no per-model entry is added to _MODEL_DEFAULT_REL_L1_THRESH' — true at this head (config.py:112-114); test_teacache_config.py:28 pins rel_l1_thresh==0.2 so a future retune must update the test deliberately.
- [claim-verified] Doc claims 'under-predicted ~3x' (design doc:263) and '~1.1x at 25 steps / ~1.5x at 50 steps' (user guide:193) match the PR's measured table (predicted/actual 0.32; 1.08x/1.54x in the discussion sweep); r=0.80 linear vs R^2=0.69 quartic are distinct quantities (0.64 <= 0.69), not self-contradictory.
- [validated] Ming subclasses (ming_flash_omni/ming_zimage_transformer.py:21, ming_image/transformer.py:13) have distinct class names, so they do NOT inherit the new coefficients; with tea_cache they hit the loud descriptive KeyError at config.py:152-156 (re-raised as ValueError at backend.py:183-189) — pre-existing, disclosed in-diff (config.py:89-91) and in the PR body.
- [validated] Runtime wiring closes end-to-end: backend.py:179 builds TeaCacheConfig(transformer_type=transformer.class.name); ZImagePipeline.transformer = ZImageTransformer2DModel (pipeline_z_image.py:228); extractor map key ZImageTransformer2DModel -> extract_zimage_context (extractors.py:1554) — the calibrated entry is reached via the real --cache-backend tea_cache path.
- [validated] Design-doc 'nearly identical' claim vs the PR's own sweep: at the shipped default (threshold 0.2, 50 steps) the PR reports mean SSIM 0.902 / min 0.762 — the basis of the second published comment.
Of 16 candidates, 3 verified comments survive, all at minor and all engaging the PR's core: (0) the newly calibrated quartic crosses zero near x~0.40 and is applied unclamped by hook.py, a residual the cited tracker #8411 demonstrably does not cover — kept as a tracker-scoping or warn-once ask on the primary change's blast radius, independently engaged by all three lenses; (1) the design doc's surviving 'should be nearly identical' verification step, which the PR's own default-threshold sweep (mean SSIM 0.902, min 0.762 at 50 steps) contradicts, in a PR whose stated purpose is fixing exactly such doc claims; and (11) the Test Plan citing a cuda-select CI lane that never runs the new cpu-marked test file, rewritten to fold in its dup's confirmation that the CPU merge lanes do cover it. Three dups fold into survivors 0 and 11 (the adversary and behavior lenses' parallel takes on the quartic hazard and the CI-lane check). Ten drops are verified: nine are sound self-pruned non-issues (eager extractor resolution has no victim, the doc adapter example is self-consistent, no stale doc consumers, tracker grouping defensible, threshold-guidance nits unfounded), and one correctly declines to re-raise the review bot's P1 threshold question already answered with the 5,110-step measurement in the discussion.
Checked, no defect found:
vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:28— No issue: eager extractor resolution in init (replacing resolution in the removed initialize_hook override) changes only failure timing. The sole production caller (TeaCacheCoefficientEstimator, coefficient_estimator.py:200) passes adapter-derived transformer class names all present in EXTRACTOR_REGISTRY ('Bagel':1551, 'ZImageTransformer2DModel':1554, 'StableAudioDiTModel':1556, 'Flux2Transformer2DModel':1557, 'LongCatImageTransformer2DModel':1558); TeaCacheHook keeps its own lazy override (hook.py:103) and is untouched; removing the override is safe because base ModelHook provides a pass-through initialize_hook (hooks/base.py:68) that the framework still calls (base.py:235); unknown types now fail loudly at construction, which the new test pins (test_teacache_config.py:43-44).
Verdict: COMMENT
CI at
a4800c2364bf(2026-10-02T02:08:08.473745+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
Findings
- [P2] The shipped Z-Image quartic (config.py:64-70) is only valid on its measured rel… —
vllm_omni/diffusion/cache/teacache/config.py
The shipped Z-Image quartic (config.py:64-70) is only valid on its measured rel_l1 range 0.024-0.336 and crosses zero at x≈0.404 (p(0.403)≈+0.017), but hook.py:262-263 applies it unclamped:rescaled_distance = float(self.rescale_func(rel_distance))/state.accumulated_rel_l1_distance += abs(rescaled_distance). So a step whose input distance lands near the root — plausible for the non-Turbo checkpoints the new comment itself says 'inherit these values uncalibrated' (config.py:56-58), with nothing bounding their rel_distance to the fitted range — rescales to ≈0, adds ≈0 to the accumulator (abs() also hides the negative lobe), and is served from cache despite being a large change: the same under-prediction/over-skip shape as #8270. The cited tracker doesn't cover it: config.py:90-91 scope #8411 to 'the Flux2Klein borrow and the Ming subclasses of ZImageTransformer2DModel'. Either add this extrapolation hazard to #8411's scope, or make it loud — a warn-once when rel_distance leaves the fitted range, or a per-model calibrated-range field on TeaCacheConfig checked in the hook — so future checkpoints inheriting this entry degrade visibly instead of silently skipping large steps.
Evidence: hook.py:262-263 — unchanged by this diff, present in the PR-time tree: rescaled_distance = float(self.rescale_func(rel_distance)) / state.accumulated_rel_l1_distance += abs(rescaled_distance); hook.py:86 self.rescale_func = np.poly1d(config.coefficients) — raw polynomial, no clamp (grep for clamp/clip/min/max/warn in hook.py: 0 matches). config.py:54 # near x~0.40, so it is only valid inside that range: the hook applies it; config.py:56-58 # 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.; config.py:90-91 # Copied from Qwen-Image, still uncalibrated; tracked in #8411 together with / # the Flux2Klein borrow and the Ming subclasses of ZImageTransformer2DModel. Root re-derived: p(0.336)≈+0.99, p(0.403)≈+0.017, p(0.405)≈-0.027 → first root x≈0.404.
- [P2] This diff qualifies the speedup half of the Overview sentence ('up to 1.5x-2.0x —
docs/user_guide/diffusion/cache_acceleration/teacache.md
This diff qualifies the speedup half of the Overview sentence ('up to 1.5x-2.0x ... on models with calibrated coefficients and enough steps to skip') but leaves the same sentence's absolute tail 'without sacrificing generation quality' unqualified — in tension even with its own 'minimal quality loss' qualifier. The PR's own calibration note (vllm_omni/diffusion/cache/teacache/config.py:63) records a 'visible quality drop' at threshold 0.25-0.30, and the author's sweep at the shipped default (rel_l1_thresh=0.2, 50 steps) reports mean SSIM 0.902 (min 0.762) / 25.9 dB vs uncached. Suggest: '...making it ideal for production deployments where inference speed matters and modest, threshold-dependent quality drift is acceptable.'
Evidence: docs/user_guide/diffusion/cache_acceleration/teacache.md:18 (line modified by this diff) 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. — first half qualified, tail absolute, both in the post-diff sentence. vllm_omni/diffusion/cache/teacache/config.py:62-63 (comment added by this diff) 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. — the PR's own in-tree acknowledgment of quality loss, decisive without the out-of-tree SSIM figures.
Suggestion: It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters and modest, threshold-dependent quality drift is acceptable.
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
| 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] |
There was a problem hiding this comment.
[P2] The new tests/diffusion/cache/test_teacache_config.py carries `pytestmark = [py…
The new tests/diffusion/cache/test_teacache_config.py carries pytestmark = [pytest.mark.core_model, pytest.mark.cpu] (line 16), but the Test Plan cites the L2 job 'Diffusion · Cache Test' as covering evidence, and that job runs timeout 15m pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model" (.buildkite/cuda/test-ready.yml:63) — a selector that collects none of this module's items, since they are cpu-marked and no hook auto-adds cuda (the only collection hook adds xdist groups only; cuda marks come solely from the @hardware_test decorator, which this module doesn't use). The tests do run per-PR via the L1 CPU lane pytest -sv tests/diffusion -m 'core_model and cpu' (.buildkite/cuda/test-ready.yml:31, mirrored on AMD at .buildkite/amd/test-amd-merge.yml:46), and the author's direct run passed 3, so coverage exists — only the plan's attribution is wrong. Correct the Test Plan bullet to name that CPU lane (or keep the direct pytest invocation) so the plan matches what CI actually executes for this file.
Evidence: tests/diffusion/cache/test_teacache_config.py:16 pytestmark = [pytest.mark.core_model, pytest.mark.cpu] — the module requests CPU, not CUDA. Decisive CI rule, unchanged by this diff, present in the PR-time tree: .buildkite/cuda/test-ready.yml:61-63 — label "Diffusion · Cache Test", command timeout 15m pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model", which cannot match cpu-only items. No auto-marking undermines this: tests/helpers/fixtures/pytest_collection.py:16-30 (only collection hook in tests/) does item.add_marker(pytest.mark.xdist_group(name="serial")) for unmarked tests and nothing else; tests/helpers/mark.py:160 marks.extend([pytest.mark.cuda, _cards_mark(num_cards)]) attaches cuda only inside the opt-in hardware_marks/hardware_test decorators, unused by the new module. Covering lane, also unchanged by this diff: .buildkite/cuda/test-ready.yml:31 pytest -sv tests/diffusion -m 'core_model and cpu' (label 'Simple · Diffusion Test' at :28) selects the module, and .buildkite/amd/test-amd-merge.yml:46 runs the AMD equivalent over tests/diffusion with the cpu selector.
| 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) |
There was a problem hiding this comment.
[P2] This diff qualifies the checklist's speedup bullet (line 401 now reads 'expect…
This diff qualifies the checklist's speedup bullet (line 401 now reads 'expect up to 1.5x-2.0x on 50-step runs; less on few-step distilled models') but leaves line 402's verification step asserting the cached output 'should be nearly identical' unconditionally. The PR's own sweep (discussion table) contradicts that at the shipped default operating point (rel_l1_thresh 0.2, Z-Image-Turbo, 50 steps): mean SSIM 0.902, min 0.762 vs uncached — visibly different; 'nearly identical' holds only at 0.15, where nothing is skipped (1.00x). The diff's own calibration note concedes the same direction: ~half the steps are already skipped at 0.2/50 steps and 'raising the threshold to 0.25-0.30 buys ~10% more speed for a visible quality drop' (config.py). Qualify the step, e.g. '(nearly identical at conservative thresholds; expect visible drift at the default threshold on few-step distilled models)'.
Evidence: docs/design/feature/teacache.md:401 2. **Compare performance** - Measure speedup vs baseline (expect up to 1.5x-2.0x on 50-step runs; less on few-step distilled models) — this is the line this diff reworded (hunk @@ -398,7 +398,7 @@, old text 'expect 1.5x-2.0x'); the next line is unchanged by the diff, read in the PR-time tree: docs/design/feature/teacache.md:402 3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical). Corroboration added by this same diff (not pre-existing): vllm_omni/diffusion/cache/teacache/config.py:60-63 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. The SSIM sweep figures cited in the comment (0.902 mean / 0.762 min at 0.20; 1.00x at 0.15) are from the PR discussion, not the repo: a tree-wide grep for 0.902|0.762|SSIM returns no hit in either teacache doc or config.py, so those exact numbers rest on the discussion thread, which a reviewer must check on GitHub.
Suggestion: 3. Verify output quality - Visually compare cached vs uncached outputs (nearly identical at conservative thresholds; expect visible drift at the default threshold on few-step distilled models)
|
|
||
| - 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 |
There was a problem hiding this comment.
[P3] This diff rewrites both this 'Good for' bullet and the Summary, and the two hun…
This diff rewrites both this 'Good for' bullet and the Summary, and the two hunks now state different ranges: the Summary says to expect ~1.1x on few-step distilled models (Z-Image-Turbo at 25 steps; matching config.py's calibration note), which is below this bullet's 1.2x floor. Lower the floor to 1.1x so both rewritten hunks state one consistent range.
Evidence: docs/user_guide/diffusion/cache_acceleration/teacache.md:152 - Scenarios where a model-dependent 1.2-2x speedup is valuable (this diff rewrites it from '- Scenarios where 1.5-2x speedup is valuable'); same file, same diff, teacache.md:193 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; corroborated in the same diff by vllm_omni/diffusion/cache/teacache/config.py:61 about half of its steps at 50 steps (~1.5x) and few at 25 steps (~1.1x),. Both doc lines are newly rewritten by this PR, so the 1.2x floor vs ~1.1x mismatch is introduced here, not pre-existing.
Suggestion: - Scenarios where a model-dependent 1.1-2x speedup is valuable
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Scan:
| Category | Result |
|---|---|
| Tests / verification | 2 finding(s) below |
| Security | no finding reported |
| Docs / comments | 2 finding(s) below |
| Behavior / compatibility | no finding reported |
| Correctness | no finding reported |
Validated:
- [resolved] initialize_hook-removal lifecycle concern: base ModelHook.initialize_hook (vllm_omni/diffusion/hooks/base.py:68) is the inherited fallback, and the only DataCollectionHook( constructor callers are coefficient_estimator.py:200 (adapter-derived transformer_type) and the new test, so extractor_fn is always bound before any forward; residual checked — TeaCacheHook keeps its own lazy binding (hook.py:103) unchanged, and test_teacache_extractors.py:783/944/1002 exercise TeaCacheHook, not DataCollectionHook
- [claim-verified] '5,110 step pairs' arithmetic: 70 prompts × ((25−1)+(50−1)) = 70×73 = 5,110 — exact.
- [claim-verified] The PR's numbers hang together: 0.089/0.278 = 0.32 matches 'predicted/actual 0.32'; '~3x under-predicted' = 1/0.32; sweep speedups 1.08x@25 / 1.54x@50 match the doc's '~1.1x' / '~1.5x' and '0.25-0.30 buys ~10% more speed' (1.54→1.67 = +8%).
- [claim-refuted] Test Plan bullet 'CI-aligned L2 job "Diffusion · Cache Test"' cannot cover the new tests: that lane runs '-m 'core_model and cuda'' (.buildkite/cuda/test-ready.yml:63) and tools/run_ready_jobs.sh runs test-ready.yml commands verbatim (tools/run_ready_jobs.sh:5,44), so the cpu-only module is deselected there; real coverage is the merge-gate lane plus the direct pytest run.
- [validated] vllm_omni/diffusion/cache/teacache/extractors.py:1554 — "ZImageTransformer2DModel": extract_zimage_context is registered; get_extractor raises ValueError("Unknown model type: ...") at extractors.py:1631-1632, exactly matching the new test's pytest.raises(match="Unknown model type")
- [validated] ZImageAdapter.model_class_name="ZImagePipeline" matches a real pipeline: vllm_omni/diffusion/registry.py:63 and vllm_omni/diffusion/models/z_image/pipeline_z_image.py:164; adapter follows the sibling pattern (BagelAdapter/Flux2Adapter/LongCatAdapter/StableAudioAdapter)
Of 16 candidates, 13 are self-cancelled verification records that their producing samples already marked '[record, not raised]', excluded, resolved, or no_issue (lifecycle of the init extractor move, Ming KeyError, removed placeholder consumers, polynomial shape, SPDX header, comment length, threshold range, and the unclamped-quartic hazard), and they are dropped rather than published as review comments. The unclamped-quartic clamp ask (candidates 0, 4, 10) is dropped on the merits as well: the diff's own comment states the zero-crossing hazard and cache-serving consequence verbatim, and the PR discussion shows the author explicitly considered and declined clamping in hook.py as a cross-model behavior change deferred to #8411, so re-raising it is re-litigating an answered question. The merged review publishes two comments that engage the primary change: (1) the Turbo-only calibration scope is disclosed only in a code comment while base Z-Image users — who reach the same ZImageTransformer2DModel entry per pipeline_z_image.py:228 and backend.py:173 — get no warning in the user guide (keep 6, minor, a protected blast-radius-scoping finding); and (2) the Test Plan cites the 'Diffusion · Cache Test' ready lane, whose 'core_model and cuda' selector can never collect the new cpu-marked module, so the PR body should cite the merge-gate 'core_model and cpu' lane instead (keep 7, minor; candidate 1 is a duplicate of 7).
Checked, no defect found:
docs/user_guide/diffusion/cache_acceleration/teacache.md:122— No issue: the narrowed suggested range (0.1-0.8 → 0.1-0.3) leaves no dependents behind. A repo-wide sweep of every rel_l1_thresh / --tea-cache-rel-l1-thresh literal found values only in {0.1, 0.17, 0.2, 0.25, 0.3}: examples default to 0.2 (image_edit.py:107 and :436, text_to_audio.py:138) with 0.25 only in usage docstrings; the MiniMax-H3 recipe uses 0.17; tests use 0.17/0.2/0.3 (the parametrized backend test covers only (None, 0.17) and (0.2, 0.2), test_cache_backends.py:487). The design doc's own troubleshooting still says 'Try 0.1-0.2' (teacache.md:462) and the guide's new 0.3 suggestion (teacache.md:183) sits at the range edge. No example, recipe, test, or doc recommends a value above 0.3.
Verdict: COMMENT
CI at
7fade114ea45(2026-10-02T07:18:45.335589+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
See inline comments below.
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
| 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] |
There was a problem hiding this comment.
[P2] This diff adds tests/diffusion/cache/test_teacache_config.py marked `core_model…
This diff adds tests/diffusion/cache/test_teacache_config.py marked core_model, cpu (line 16), but the Test Plan cites the 'Diffusion · Cache Test' ready-lane job as CI validation — that lane runs pytest -sv tests/diffusion/cache -m 'core_model and cuda' (.buildkite/cuda/test-ready.yml:63) and tools/run_ready_jobs.sh replays test-ready.yml's pytest commands verbatim (tools/run_ready_jobs.sh:5,44), so it deselects every test in this cpu-marked module and cannot evidence it. The selector that actually runs this file is the merge-gate lane pytest -sv tests/diffusion -m 'core_model and cpu' (.buildkite/cuda/test-merge.yml:23; AMD equivalent at .buildkite/amd/test-amd-merge.yml:46, which adds card-marker exclusions and ignores but still collects this module). Cite that lane (or a direct run) in the Test Plan so the cited evidence names a command that can collect this module, and note the ready lane deselects it.
Evidence: tests/diffusion/cache/test_teacache_config.py:16 (added by this diff): pytestmark = [pytest.mark.core_model, pytest.mark.cpu] — no test in the module carries a cuda marker. Decisive CI selectors, unchanged by this diff but present in the PR-time tree: .buildkite/cuda/test-ready.yml:61-63 — - label: "Diffusion · Cache Test" / timeout 15m pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model" (the and cuda term deselects this cpu-only module, so a green run of the cited lane proves nothing about it); .buildkite/cuda/test-merge.yml:20,23 — - label: "Simple · Diffusion Test" / "pytest -sv tests/diffusion -m 'core_model and cpu'" (matches both markers and the directory: the lane that actually runs the file); .buildkite/amd/test-amd-merge.yml:46 — "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" (does not ignore the new file); tools/run_ready_jobs.sh:5,44 — # Extract steps from .buildkite/cuda/test-ready.yml that contain pytest, synthesize / BUILDKITE_REL=".buildkite/cuda/test-ready.yml" (replay consumes that same YAML).
| ## 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 |
There was a problem hiding this comment.
[P2] This diff replaces the ZImageTransformer2DModel placeholder with coefficients c…
This diff replaces the ZImageTransformer2DModel placeholder with coefficients calibrated only on Z-Image-Turbo and states that scope only in a code comment (config.py:56-58: 'Validated only on Z-Image-Turbo; other checkpoints that reach TeaCache through this class name inherit these values uncalibrated'). Both Z-Image checkpoints reach this single entry: pipeline_z_image.py:228 unconditionally constructs ZImageTransformer2DModel (registry.py:63-67 has one ZImagePipeline, no base/Turbo split), and teacache/backend.py:173 auto-detects transformer_type from the class name — so the base Tongyi-MAI/Z-Image checkpoint silently inherits Turbo-fitted values. Neither the user guide Summary nor the supported-models row (docs/user_guide/diffusion_features.md:171) signals this. State the calibration scope where users read — append to the Summary bullet at docs/user_guide/diffusion/cache_acceleration/teacache.md:193: 'Coefficients are calibrated on Z-Image-Turbo only; the base Z-Image checkpoint inherits them uncalibrated.'
Evidence: Decisive lines: in-diff — vllm_omni/diffusion/cache/teacache/config.py:56-58 # 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. (scope caveat lives only in a code comment). Unchanged by this diff, present in the PR-time tree — vllm_omni/diffusion/models/z_image/pipeline_z_image.py:228 self.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_config) (unconditional, both checkpoints); vllm_omni/diffusion/registry.py:63-67 "ZImagePipeline": ("z_image", "pipeline_z_image", "ZImagePipeline",) (single pipeline, no Turbo variant); vllm_omni/diffusion/cache/teacache/backend.py:173 transformer_type = transformer.__class__.__name__ (auto-detect keys _MODEL_COEFFICIENTS lookup); docs/user_guide/diffusion_features.md:171 | **Z-Image** | ✅ | ✅ | ✅ | ... (supported-models row with no calibration-scope note). Anchor: docs/user_guide/diffusion/cache_acceleration/teacache.md:193 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.
Suggestion: 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. Coefficients are calibrated on Z-Image-Turbo only; the base Z-Image checkpoint inherits them uncalibrated.
|
|
||
| ```python | ||
| class YourModelAdapter: | ||
| class YourModelAdapter(DefaultAdapter): |
There was a problem hiding this comment.
[P2] teacache.md:271 now subclasses DefaultAdapter but keeps the example's raw-diffu…
teacache.md:271 now subclasses DefaultAdapter but keeps the example's raw-diffusers load_pipeline override (returning YourModelPipeline.from_pretrained(...) at :279) — exactly what the base class exists to replace: DefaultAdapter.load_pipeline loads through vLLM's DiffusersPipelineLoader (coefficient_estimator.py:78-80). TeaCacheCoefficientEstimator drives whatever load_pipeline returns via self.pipeline.forward(DiffusionRequestBatch(requests=[req])) (coefficient_estimator.py:198, 217), which a raw diffusers pipeline does not implement, so a reader following the example verbatim crashes at the very next 'Collect Data and Estimate' step. All five in-tree adapters rely on the inherited loader — including ZImageAdapter, the pattern this PR just added (coefficient_estimator.py:131-134, sets only model_class_name). Rewrite the example in that minimal pattern (set model_class_name = "YourModelPipeline", keep get_transformer) and drop the load_pipeline and install_hook overrides (the latter duplicates DefaultAdapter verbatim, coefficient_estimator.py:87-89).
Evidence: docs/design/feature/teacache.md:271 class YourModelAdapter(DefaultAdapter): (line changed by this diff) with unchanged-by-this-diff override at :275 def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any: returning :279 pipeline = YourModelPipeline.from_pretrained( — present in the PR-time tree. Decisive consumer, unchanged by this diff, present in PR-time tree: vllm_omni/diffusion/cache/teacache/coefficient_estimator.py:198 self.pipeline = adapter.load_pipeline(model_path, device, dtype) and :217 self.pipeline.forward(DiffusionRequestBatch(requests=[req])). Base class, unchanged by this diff: coefficient_estimator.py:78-80 loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config) / return loader.load_model(load_device=device); and :131-134 class ZImageAdapter(DefaultAdapter): with only model_class_name = "ZImagePipeline" (no load_pipeline override).
Suggestion: class YourModelAdapter(DefaultAdapter):
"""Adapter for coefficient estimation on your model."""
model_class_name = "YourModelPipeline"
@staticmethod
def get_transformer(pipeline: Any) -> tuple[Any, str]:
"""Extract transformer from pipeline."""
return pipeline.transformer, "YourTransformer2DModel"
Register your adapter
_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
|
|
||
| - 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 |
There was a problem hiding this comment.
[P3] This diff rewrites the good-for bullet to 'a model-dependent 1.2-2x speedup' (l…
This diff rewrites the good-for bullet to 'a model-dependent 1.2-2x speedup' (line 152), but the Summary the same diff adds (line 193) tells Z-Image-Turbo users at 25 steps to 'expect ~1.1x there' — the bullet's lower bound excludes a case this PR itself documents later in the same file. Lower it to 1.1 so the two sections agree.
Evidence: Both decisive lines are added by this diff and read verbatim in the PR-time tree: docs/user_guide/diffusion/cache_acceleration/teacache.md:152 - Scenarios where a model-dependent 1.2-2x speedup is valuable vs same file :193 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`.
Suggestion: - Scenarios where a model-dependent 1.1-2x speedup is valuable
hsliuustc0106
left a comment
There was a problem hiding this comment.
Approving — this is the fix for #8270 with the calibration evidence to back it.
What I verified:
- The new Z-Image coefficients replace the Qwen-Image placeholder exactly as measured upthread (SSIM 0.786/0.563 → 0.980/0.902 at thr 0.2, speed-up 1.08×/1.54×), and the test pins them with
pytest.approxso a future recalibration must consciously update it — plus asserts Z-Image ≠ Qwen-Image so the placeholder can't silently return. ZImageAdapteris a cleanDefaultAdaptersubclass registered in_MODEL_ADAPTERS, andextract_zimage_contextalready exists on main, so the estimator path needs no new extractor.- Moving extractor resolution from
initialize_hookto__init__is strictly safer (unknown transformer types now fail at construction instead of at first forward);DataCollectionHookhas no consumers outside the estimator module and the base class provides the defaultinitialize_hook. - Docs now tell the truth about borrowed coefficients ("Unreliable … treat as a placeholder and calibrate", with the #8270 numbers) and about few-step distilled models expecting ~1.1×, and the Flux2 comment now points at #8411 for the remaining borrows. Good.
Two non-blocking notes, both already disclosed in the config comment but worth tracking: (1) the quartic crosses zero near x~0.40 and is applied unclamped, so an extreme input distance beyond the measured 0.024–0.336 range would rescale toward ~0 and wrongly serve a large step from cache — a follow-up clamp (x or rescaled value) would make the fit safe outside its calibration range; (2) values are validated on Z-Image-Turbo only; other checkpoints resolving through this class name inherit them uncalibrated — #8411 is the right place to extend coverage.
vllm-project#8408) Signed-off-by: Kamila Kuatova <kuatova_kamila@list.ru> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Purpose
Fixes #8270.
The TeaCache polynomial for
ZImageTransformer2DModelwas a placeholder copied from Qwen-Image (the config comment said "need to be tuned specifically for Z-Image in future"). On Z-Image it predicts about one third of the measured step-to-step output change, so the accumulated rescaled distance stays underrel_l1_threshfor too long and TeaCache skips steps that matter. With--cache-backend tea_cache, users got visibly different, blurred images instead of a faster version of the same image (mean SSIM vs. uncached 0.56 at 50 steps).This PR:
Tongyi-MAI/Z-Image-Turbo: 70 Parti prompts sampled across categories × {25, 50} steps, 1024×1024,guidance_scale=4.0, 5,110 consecutive-step pairs, 4th-ordernp.polyfit(Pearson r = 0.80, R² = 0.69), followingdocs/design/feature/teacache.md.ZImageAdaptertoTeaCacheCoefficientEstimatorso the calibration is reproducible, and type-annotates the adapter registry somypypasses on the file.tests/diffusion/cache/test_teacache_config.py).Why the borrowed polynomials fail, evaluated on the measured Z-Image input distances:
No default-threshold change:
rel_l1_threshstays at 0.2. With correct coefficients that is a quality-first operating point: Z-Image-Turbo changes more than 20% per step early in the trajectory, so at 25 steps few steps are skippable (~1.1×) and at 50 steps about half are (~1.5×). The sweep in the comments shows 0.25–0.30 buys ~10% more speed for a visible quality drop (min SSIM 0.95 → 0.86 at 25 steps), so no per-model entry is added to_MODEL_DEFAULT_REL_L1_THRESH; users can still passcache_config={"rel_l1_thresh": 0.25}.Out of scope, to be tracked in a follow-up issue:
Flux2Transformer2DModelandFlux2Kleinstill use borrowed coefficients, and the Ming subclasses ofZImageTransformer2DModel(MingZImageTransformer2DModel,MingImageTransformer2DModel) have no_MODEL_COEFFICIENTSentry at all, sotea_cacheraises aKeyErrorfor them (pre-existing). The calibration collector used here computed per-step rel-L1 on the GPU instead of storing CPU trajectories; landing that as an option onDataCollectionHookis a candidate follow-up too.Reproduce the bug on
main(same prompt and seed, the example's default 50 steps;tea_cache.pngcomes out blurred):Test Plan
vLLM Version: 0.30.0 (torch 2.13.0+cu132, CUDA 13.2, Python 3.12)
vLLM-Omni Commit: rebased on
7b15d22f4; issue measured at0173cb374pre-commit run --from-ref <base> --to-ref HEAD(ruff, ruff-format, mypy 3.10, SPDX, typos, torch.cuda, CI marks): all pass.pytest tests/diffusion/cache/test_teacache_config.py -p no:cacheprovider(CPU): 3 passed. The module fails onmain(ImportError: cannot import name 'ZImageAdapter'; the coefficient assertion would also fail since the placeholder equals Qwen-Image's).tools/run_ready_jobs.sh --label-substr "Cache Test"on 1× A100-40GB:8 passed, 209 deselected in 157.80s(includestest_teacache[riverclouds/qwen_image_random]and the GPU extractor tests). L1 "Simple · Diffusion Test" result is in the comments.TeaCacheCoefficientEstimator-style collection withZImageAdapteron 70 Parti prompts (nateraw/parti-prompts, seeded shuffle) × {25, 50} steps, 1024×1024,guidance_scale=4.0, thennp.polyfit(x, y, 4). Collection computed the rel-L1 distances on the GPU per step instead of storing activations (the stock hook would need ~13 GB of host RAM per 50-step 1024² CFG trajectory).guidance_scale=4.0, viaOmni(model, cache_backend="tea_cache", cache_config={"rel_l1_thresh": ..., "coefficients": ...}). Each TeaCache image is compared with its uncached twin (same prompt, seed, steps) withtorchmetricsSSIM/PSNR (data_range=1.0). First request per process (torch.compile warm-up) excluded; timings are means over 18 images per cell. VRAM was not measured (unchanged code path).Test Result
The speed-up drops because the previous 2–3× came from skipping steps that should have run. At 25 steps there is little for TeaCache to skip on this distilled model; the gain is at 50 steps. Results for
rel_l1_thresh0.25 and 0.30 will be added in a comment to document the quality/speed trade-off.Side-by-side (uncached | placeholder | SGLang | calibrated), same prompt and seed: