[fusheng-ji]

Purpose

Part of #8411 (item 2, Flux2Klein).

TeaCacheHook keeps one cache state per CFG branch. It decides which branch a transformer call belongs to from transformer.do_true_cfg: when the flag is set and CFG is sequential, calls alternate positive/negative on a counter; otherwise every call goes to the positive state.

Only Flux, Flux Kontext and Qwen-Image set that flag. Every other pipeline that runs sequential CFG through CFGParallelMixin.predict_noise_maybe_with_cfg leaves it unset. For the TeaCache-enabled ones among them (Flux2Klein, Flux2 with a negative prompt, LongCat-Image), the positive and negative calls of each step therefore share one TeaCache state. Both calls see the same latents and timestep, so the negative call's modulated input equals the positive call's exactly; its distance is driven only by the polynomial's constant term, and the two branches reuse each other's residuals. Guidance collapses and the output drifts to a different composition.

Change: CFGParallelMixin._set_transformer_do_true_cfg sets self.transformer.do_true_cfg from predict_noise_maybe_with_cfg (and from PipelineParallelMixin when PP world size > 1) before predicting, where the branch layout is known. The flag is cleared for calls without CFG. Pipelines without a transformer attribute are untouched, multi-branch CFG goes through the separate predict_noise_with_multi_branch_cfg, and TeaCache is the only reader of the flag, so uncached generation is unchanged. Redundant per-pipeline writes in Flux, Flux Kontext, and Qwen-Image were removed so the mixin is the single source of truth.

Follow-up: sequential CFG still alternates TeaCache branches via _forward_cnt parity when non-CFG forwards interleave; explicit per-call branch stamping is tracked in #8483 (not in this PR).

For direct pipeline calls without an initialized CFG process group, TeaCacheHook now treats the CFG world size as 1 and uses sequential positive/negative cache states. Initialized groups still select the branch by rank, and errors from an initialized group are propagated. This matches the shared CFG helper's fallback.

No coefficient change. On FLUX.2-klein-base-4B the current (FLUX.1-borrowed) coefficients over-predict the output change by about 1.9x, which keeps TeaCache conservative; once the branches are separated they give quality comparable to the calibrated Z-Image result in #8408. A 4th-order refit on 10,220 within-branch pairs (70 Parti prompts x {25, 50} steps) reached only R² = 0.15, so I am not proposing new coefficients; details in #8411.

Test Plan and Results

Branch head: 52df1b0aa72016d5cd35a3be71f7ac19b9e51d44 (fix/teacache-cfg-branch-state).

Validation environment: Python 3.12.13, Torch 2.13.0, vLLM 0.31.0 loaded from a matching wheel, and 2 x B200. Tests use VLLM_WORKER_MULTIPROC_METHOD=spawn; CPU runs hide CUDA devices. Three-card cases are excluded.

After the uninitialized-group fallback fix (and the review cleanup that centralizes the flag and adds a PP>1 CPU regression):

Suite	Passed	Failed	Errors	Skipped
New real-hook regressions + CFG CPU + hook registry	30	0	0	0
Cache core_model and cpu	205	8	7	0
CFG full_model and cuda and B200 and cards_2	10	0	0	0
Cache core_model and cuda	8	0	0	0

• Four TeaCache CFG-group tests install TeaCacheHook through apply_teacache_hook. They verify separate positive/negative residuals and cache hits on the next step with an absent or single-rank CFG group and initialized rank selection. A low-value broken-CFG-group probe was dropped in review cleanup.
• Review cleanup (CPU, PYTHONPATH=/tmp/vllm-8474-wheel + worktree, vLLM 0.31.0): test_teacache_cfg_group.py + test_cfg_parallel.py (core_model and diffusion and cpu) → 23 passed; new test_pp_predict_noise_sets_transformer_do_true_cfg → 1 passed. GPU suites were not rerun for the cleanup commit.
• Existing CFG CPU coverage includes setting/resetting the transformer flag and pipelines without a transformer attribute. GPU CFG coverage includes sequential/parallel CFG, normalization, and multi-branch CFG. GPU cache coverage includes TeaCache image generation and Flux2/Klein extractor checks.
• The cache CPU suite's 8 failures and 7 errors match the verified baseline 539c5337, including failure messages; no new failures were found.
• Before the final fallback fix, the full diffusion core_model and cpu sweep on bba2b29 produced 7,190 passed, 89 failed, 183 errors, and 79 skipped. All 272 failing/error cases were rerun on baseline 539c5337 with matching messages. The full diffusion sweep was not repeated after the fallback fix; the focused CPU, cache CPU, and GPU suites above were rerun.
• Ruff, formatting, and git diff --check pass. All applicable pre-commit checks pass except mypy-3.10: two pre-existing extractor_fn/None typing errors in hook.py were independently reproduced on the baseline. The new test file passes mypy.

Full logs, JUnit reports, and the standalone real-hook probe are retained outside the repository under /tmp/pr8475-validation-20261005/ on the validation host.

Original Image Quality Measurements

These measurements predate the final uninitialized-group fallback fix and were not rerun for that fix. Environment: 1 x B200, frozen vllm/vllm-openai:v0.30.0 container, baseline ee8fdab1, offline Omni, 1024x1024, the #8408 protocol (6 prompts x 3 seeds; each TeaCache image against its uncached twin; SSIM with a Gaussian window, sigma 1.5, population covariance; warm-up excluded; speed-up from mean request latency).

TeaCache at the default rel_l1_thresh=0.2, 50 steps, guidance_scale=4.0 (both model cards' defaults), each TeaCache image against its uncached twin:

Model	main ee8fdab1: mean SSIM (min) / PSNR / speed-up	this PR: mean SSIM (min) / PSNR / speed-up
black-forest-labs/FLUX.2-klein-base-4B	0.592 (0.331) / 14.9 dB / 1.55x	0.900 (0.720) / 24.3 dB / 1.65x
meituan-longcat/LongCat-Image	0.550 (0.292) / 13.5 dB / 1.57x	0.888 (0.646) / 25.1 dB / 1.59x

• Uncached output is unchanged (18/18 images byte-identical): LongCat-Image main vs this PR, and FLUX.2-klein-base-4B main vs the flag-only A/B tree.
• FLUX.2-klein-4B (distilled, 4 steps, no CFG) is unaffected: per-step input change is 0.33-0.53, so TeaCache never skips at 0.2 and images match the uncached ones exactly, before and after.
• On main the damaged images are different compositions, not a blur (e.g. an isometric library turns into an incoherent scene).

An A/B tree that only adds the flag inside the Klein pipeline gives the same SSIM as this PR to the last digit. Scripts, images and logs are kept outside the repository and can be shared on request.

AI assistance: Codex and Claude Code were used to trace the cache path, design measurements, and draft changes, tests, and this description. Codex executed the latest CPU/GPU validation and baseline comparisons.

[fusheng-ji]

Self-review

• Root cause. TeaCacheHook separates CFG branches only when the hooked transformer has do_true_cfg set; only Flux, Flux Kontext and Qwen-Image set it. Pipelines using CFGParallelMixin.predict_noise_maybe_with_cfg (Flux2Klein, Flux2 with a negative prompt, LongCat-Image among the TeaCache-enabled ones) shared one state across branches. Both branch calls of a step have identical modulated inputs, so they reused each other's residuals.
• Why the shared helper. It is where the branch layout is known; one place fixes every caller, resets the flag for no-CFG calls, skips pipelines without a transformer attribute, and leaves multi-branch CFG (separate method) alone. TeaCache is the only reader, so uncached output is unchanged.
• Tests. Flag test fails on main, passes here; no-transformer pipeline still works. 19 passed in test_cfg_parallel.py; the helper's pipeline suites show the same 149 env failures on main and here.
• GPU. 1 x B200, [Bugfix] Calibrate Z-Image TeaCache coefficients and add ZImageAdapter #8408 protocol: klein-base-4B SSIM 0.592 -> 0.900 at 1.55x -> 1.65x, LongCat-Image 0.550 -> 0.888 at ~1.58x; uncached byte-identical; distilled klein-4B unaffected.
• Out of scope. No coefficient change (refit R^2 0.15, current entry is conservative); estimator fixes for Flux2 collection are a separate follow-up, see [Bug]:TeaCache coefficient gaps after #8270: Flux2/Flux2Klein still use borrowed polynomials, Ming Z-Image subclasses raise KeyError #8411.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/parallelism.md, docs/design/module/diffusion/index.md.

Module owners: @princepride @xuechendi @david6666666

Routing: @princepride via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @xuechendi via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description

@fusheng-ji, 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.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[hsliuustc0106]

Verdict: Approve. Reviewed at diff level against merge-base ee8fdab1, plus a logic-level replay of the branch-selection code (no GPU runs from my side). The change is minimal, correctly scoped, and the evidence chain in the description (SSIM/PSNR A/B, byte-identical uncached outputs, flag-only A/B tree) is unusually strong.

What I verified:

• transformer.do_true_cfg has exactly one reader in production code (vllm_omni/diffusion/cache/teacache/hook.py:148). Flux / Flux-Kontext / Qwen-Image already set the same value themselves, so their behavior is unchanged; every other mixin user with sequential CFG left it unset, so pos/neg calls both landed in teacache_positive — confirming the reported guidance collapse for the TeaCache-enabled ones (Flux2Klein, Flux2, LongCat-Image).
• The CFG-parallel path is safe: the hook selects the branch from cfg_rank (not the counter) when cfg_parallel_size > 1, so setting the flag on all ranks gives each rank a consistent branch.
• The multi-branch path (predict_noise_with_multi_branch_cfg, Bagel) is untouched, and pipelines without a .transformer attribute fall through the getattr guard.
• A pure-Python replay of predict_noise_maybe_with_cfg + the hook's branch selection confirms: on main both calls of a CFG step share one state; with this PR they separate. The two new CPU regressions cover exactly this and the no-transformer path.

Two non-blocking notes:

1. The sequential-path branch assignment relies on _forward_cnt % 2 parity. An odd number of non-CFG forwards interleaved between CFG steps (e.g. separate CFG/non-CFG batches alternating on the same transformer under continuous scheduling, since _prepare_cfg_scalars forbids mixing within a batch) flips the parity: the two branches of a step still land in different states, but their histories swap, so a branch can continue from the other branch's residual/distance state. This is pre-existing for Flux/Qwen-Image; this PR widens exposure to the other mixin users. Worth a follow-up that makes the mixin stamp an explicit branch id per call instead of relying on parity.
2. The SPDX copyright-line change rides along; fine if the hook enforces it, just noting it's unrelated churn.

Nit on the notes: none blocking — happy to see this merged as-is.

[fusheng-ji]

Thanks for the thorough review and the logic-level replay, @hsliuustc0106!

1. Parity dependency. Agreed, and I traced where it can bite. In request mode, DiffusionModelRunner._execute_request_list refreshes the TeaCache hook before each request batch (_refresh_cache_for_requests -> reset_state), so _forward_cnt restarts at 0 and the branches alternate correctly. The stepwise path never refreshes, so under step execution the counter and the per-branch state persist across scheduler waves. That allows the parity flip you describe, and more generally lets a new request continue from another request's residual. I opened [Bug]: TeaCache infers the CFG branch from forward-count parity and never resets under step execution #8482 for it. It proposes having the mixin stamp an explicit branch id per call (which removes the parity dependency everywhere, Flux/Qwen-Image included) and handling request boundaries under step execution. I'll take the explicit-branch part there and keep this PR as the minimal fix.
2. SPDX line. It was rewritten by the repository's check-spdx-header pre-commit hook when the file was touched. Happy to revert it if you'd rather keep this diff to the functional change.

[fusheng-ji]

Latest head: 52df1b0aa72016d5cd35a3be71f7ac19b9e51d44 on fix/teacache-cfg-branch-state.

f041cb2a — handle TeaCache without an initialized CFG group (unchanged from the prior note).

The shared CFG helper supports direct pipeline calls without a CFG process group, but enabling the real TeaCache hook could still trigger an assertion when the hook queried the group's world size. The hook now treats an uninitialized CFG group as world size 1 and follows sequential CFG, keeping positive and negative cache states separate. Initialized groups continue to select the branch by rank, and errors from an initialized group still propagate.

The commit adds five regressions that actually install TeaCache through apply_teacache_hook. They check branch-specific residual reuse on the next step with no group or a single-rank group, rank-based branch selection, and error propagation.

52df1b0a — review cleanup after maintainer feedback: centralize transformer.do_true_cfg on CFGParallelMixin (including PP>1), remove redundant Flux/Qwen writes, tighten TeaCache CFG tests, add test_pp_predict_noise_sets_transformer_do_true_cfg, and revert unrelated SPDX churn.

Validation after 52df1b0a:

• New regressions + CFG CPU + hook registry: 30 passed.
• Two-B200 CFG suite: 10 passed (three-card cases excluded).
• GPU cache suite: 8 passed.
• CPU cache suite: 205 passed, 8 failed, 7 errors; all failures/errors match the verified baseline 539c5337.
• Review cleanup CPU rerun (test_teacache_cfg_group.py, test_cfg_parallel.py, new PP>1 test): 24 passed; GPU suites were not rerun for the cleanup commit.
• All applicable pre-commit checks pass except two pre-existing mypy errors in hook.py, independently reproduced on the baseline. The new test file passes mypy.

The PR description reflects these results. The full diffusion CPU sweep and the original image-quality measurements were not repeated after the fallback or cleanup commits; the focused and cache CPU/GPU suites above were rerun where noted.

[fusheng-ji]

Following up on the earlier update: commit 52df1b0 adds the following changes on top of f041cb2:

• Extracts the transformer.do_true_cfg assignment into CFGParallelMixin._set_transformer_do_true_cfg, shared by the normal CFG path and the pipeline-parallel path.
• Sets the flag explicitly in the PP > 1 path, which bypasses the normal CFG helper. This preserves the flag when removing the redundant assignments from Flux, Flux Kontext, and Qwen-Image.
• Adds a CPU regression verifying that the PP path sets the flag before both CFG predictions and clears it for a subsequent non-CFG prediction.
• Simplifies test descriptions and removes the synthetic broken-group error test. Four real-hook cases remain, covering cache isolation/reuse with an absent or single-rank CFG group and branch selection for both CFG ranks.

The uninitialized-CFG-group fallback introduced in f041cb2 remains unchanged. Sequential TeaCache branch selection still uses counter parity; explicit per-call branch identification is tracked separately in #8483.

Validation clarification: the earlier comment’s “Validation after 52df1b0a” heading was too broad. On 52df1b0, the focused CFG/TeaCache CPU tests passed 23 cases, and the new PP CPU regression passed 1 case. That PP regression uses a stubbed group; it is not a multi-GPU PP run.

The previously reported 30 CPU passes, 10 two-B200 CFG passes, 8 GPU cache passes, and the 205 passed / 8 failed / 7 errors CPU cache result were obtained on f041cb2, before this cleanup. At the time of that update, GPU suites, the full diffusion CPU sweep, and image-quality measurements had not been rerun on 52df1b0. Subsequent GPU validation is recorded below.

Subsequent GPU validation on 52df1b0a

We have now rerun the two-GPU PP suite (5 passed), added an external two-GPU PP flag-transition check in FP32 and BF16 (2 passed), rerun the two-GPU CFG suite (10 passed), and rerun the GPU cache suite (8 passed).

In the additional check, both real PP ranks correctly handled CFG on → off → on, observing transformer.do_true_cfg as [True, True, False, True, True] before the corresponding predictions. The final-stage outputs matched the single-GPU baseline in both FP32 and BF16.

Limits: this does not cover the four-GPU PP=2, CFG=2 topology, and we have not validated full-model generation with PP + TeaCache + CFG enabled together. The PP transition check uses a layer-partitioned test model with real two-GPU NCCL communication; the separate GPU cache tests do not establish that combined full-model coverage. The full diffusion CPU sweep and image-quality benchmarks were not rerun for this update.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 4f86166260ed produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[fusheng-ji]

Four-GPU follow-up on 7e88fc3 (the current head of this PR):

• On one Slurm node with 4 x B200, the existing PP=2, CFG=2 predict-noise and scheduler parity tests both passed (2 passed).
• The additional real four-rank flag-transition checks passed in FP32 and BF16 (2 passed). They verified all four (PP rank, CFG rank) identities, CFG on → off → on on every rank, and final-stage outputs against the single-GPU baseline.
• Successful run: Slurm job 1095857, partition overflow, job name test, Torch 2.13.0 / vLLM 0.31.0. Logs, environment lock and JUnit reports are retained in the validation artifacts. The first attempt's pytest collection failed because pytest-xdist was missing; the reported passes come from the rerun after installing it.

This closes the previously stated four-GPU topology test gap for #8475. It does not establish full pretrained-model PP + TeaCache + CFG generation or image/video quality. Wan model adaptation, calibration and full-model validation are being developed separately in draft #8731, which also incorporates the explicit-branch dependency from #8483.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 120s (try 2 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 120s (try 2 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 600s (try 3 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (Cursor exited with 1: ActionRequiredError: You've hit your usage limit You've saved $3501 on API model usage this month with Ultra. Switch to a different model or set a Spend Limit to continue with Auto. Your usage limits will reset when your monthly cycle ends on 10/20/2026.; retrying direct/cursor).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (Cursor exited with 1: ActionRequiredError: You've hit your usage limit You've saved $3501 on API model usage this month with Ultra. Switch to a different model or set a Spend Limit to continue with Auto. Your usage limits will reset when your monthly cycle ends on 10/20/2026.; retrying direct/cursor).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (Cursor exited with 1: ActionRequiredError: You've hit your usage limit You've saved $3501 on API model usage this month with Ultra. Switch to a different model or set a Spend Limit to continue with Auto. Your usage limits will reset when your monthly cycle ends on 10/20/2026.; no fallback: the fallb).

[fusheng-ji]

Hi @Isotr0py @princepride @SamitHuang, could you take a look at this TeaCache/CFG fix when you have a chance?

The PR centralizes do_true_cfg updates in the shared CFG and PP paths so sequential positive/negative predictions use separate TeaCache states, including when the CFG process group is uninitialized.

On commit cdbb1693, validation passed 10 two-B200 CFG tests, 2 four-B200 PP=2/CFG=2 parity tests, and 2 four-rank CFG on→off→on flag tests. These results are tied to that tested commit, rather than a claim that the latest head has been revalidated.

Explicit per-call branch identification is handled separately in #8483. Feedback on the shared-helper placement and TeaCache fallback would be especially helpful. Thanks!