[fusheng-ji]

Purpose

Part of #8482 (the explicit-branch half). Stacked on #8475: the first commit here is #8475's change rebased onto current main (same patch); please review only the second commit. I'll rebase onto main as soon as #8475 merges.

TeaCacheHook picks the per-branch cache state for sequential CFG from _forward_cnt % 2. In request mode the runner refreshes the hook before every request, so the parity restarts at zero and stays aligned. Under step execution the hook is never refreshed, so a CFG wave followed by a non-CFG wave with an odd number of forwards flips the parity. The next CFG step's positive call then lands in the negative state and continues from the other branch's residual and distance. This was raised in the #8475 review and is pre-existing for Flux / Qwen-Image; #8475 widens it to every CFGParallelMixin user.

Change:

• CFGParallelMixin stamps transformer.cfg_branch ("positive" / "negative") around each predict_noise call on all three paths: sequential CFG, each CFG-parallel rank, and no CFG. A new _predict_noise_for_cfg_branch helper restores the previous value afterwards, also when predict_noise raises, so the stamp never leaks into calls made outside the helper.
• TeaCacheHook prefers explicit sources, most specific first: the extractor's per-call hint ctx.extra_states["teacache_branch"] added for MammothModa2 in [Perf] Support TeaCache for MammothModa2 #5357, then the mixin stamp. Both go through the same validation, which keeps [Perf] Support TeaCache for MammothModa2 #5357's error message for invalid hints. The old inference (CFG rank, else forward-count parity) moves unchanged into _infer_cfg_branch and is used only when neither is present. MammothModa2 does not use the mixin, so in practice the two explicit sources never meet.
• CachedTransformer declares cfg_branch = None next to do_true_cfg.
• With do_true_cfg=True and negative_kwargs=None, the helper now raises a ValueError that names the branch, instead of a bare TypeError from **None.

In request mode the stamp equals the old parity, so TeaCache decisions there are unchanged. Every TeaCache model that runs CFG as two calls (Flux, Flux Kontext, Qwen-Image, Flux2, Flux2Klein, LongCat-Image) goes through the mixin. Z-Image batches both branches into one call and is unaffected.

Not in this PR: request boundaries under step execution (the second half of #8482, which needs a maintainer decision).

Test Plan

CUDA_VISIBLE_DEVICES='' VLLM_TARGET_DEVICE=cpu python -m pytest -q \
  tests/diffusion/cache/test_teacache_cfg_branch.py \
  tests/diffusion/distributed/test_cfg_parallel.py \
  -m 'core_model and cpu' --run-level=core_model

• test_interleaved_non_cfg_wave_keeps_branch_states_aligned: a real TeaCacheHook on a probe module, driven through CFGParallelMixin as CFG wave -> non-CFG wave -> CFG wave. It asserts the hook selects positive, negative, positive, positive, negative. Before this change it selected negative, positive for the last wave.
• test_unstamped_callers_keep_the_parity_fallback: direct calls without a stamp keep the old parity behaviour.
• test_extractor_hint_takes_precedence_over_stamp, test_invalid_stamp_is_rejected: the [Perf] Support TeaCache for MammothModa2 #5357 hint wins over the stamp, and an invalid stamp raises.
• test_cfg_helper_stamps_branch_per_call_and_restores_it, test_cfg_parallel_rank_stamps_its_branch[0/1], test_cfg_branch_stamp_is_restored_when_predict_noise_raises: the mixin stamps the right branch on each path and restores it.

The CFG-parallel test also patches cfg_parallel.is_cfg_group_initialized, because _get_cfg_world_size_or_one() reports 1 without an initialized group.

vLLM Version: 0.30.0

vLLM-Omni Commit: main c1e84ce9 (includes #5357) + #8475 (rebased; the CPU results below were re-run there)

Test Result

• The 5 new behaviour tests fail on [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 alone and pass with this change; the fallback test passes on both.
• tests/diffusion/cache/, test_cfg_parallel.py, test_diffusion_step_pipeline.py, the Flux, Flux2, Qwen-Image and LongCat-Image model suites, and MammothModa2's TeaCache tests: 350 passed with this PR vs 340 on main. Both trees show the same environment-dependent failures/errors (ProcessGroupNCCL is only supported with GPUs, and ReplicatedLinear has no cpu_linear when building the real MammothModa2 module on CPU); none are new. MammothModa2's TeaCache tests pass on both (9/9).
• pre-commit passes (the SPDX hook normalized the copyright line of hook.py and base.py). mypy-3.10 on the touched files: 3 errors vs 5 on main, none new.
• No GPU run: request-mode TeaCache decisions are unchanged by construction, and the step-execution path that changes is covered by the CPU test above.

AI assistance: Codex and Claude Code were used to trace the cache path, draft the change, tests and this description. All changes were reviewed and the tests run by the author.

[fusheng-ji]

Self-review

• Stacking. The first commit is [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475's change rebased onto current main (same patch); this PR's change is only 461418b1. I'll rebase onto main once [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 merges.
• Composition with [Perf] Support TeaCache for MammothModa2 #5357. main now also has MammothModa2's per-call extractor hint (ctx.extra_states["teacache_branch"]). The hook checks it first, then the mixin stamp, then the old inference; both explicit sources share the same validation and [Perf] Support TeaCache for MammothModa2 #5357's error message.
• Root cause. For sequential CFG, TeaCacheHook chose the branch state from _forward_cnt % 2. Request mode refreshes the hook before every request, so the parity restarts at 0. Step execution never refreshes it, so an odd number of non-CFG forwards between CFG waves flips the parity and swaps the branch histories. This is the case raised in the [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 review.
• Fix. CFGParallelMixin stamps transformer.cfg_branch around each predict_noise call on all three paths: sequential CFG, each CFG-parallel rank, and no CFG. The previous value is restored in a finally, so a stale stamp can't leak into calls made outside the helper. The hook checks explicit sources first (the [Perf] Support TeaCache for MammothModa2 #5357 extractor hint, then the stamp); the old rank/parity inference is moved unchanged into _infer_cfg_branch and only runs when neither is present.
• Behaviour change. None in request mode, where the stamp always equals the old parity. Every TeaCache model with two-call CFG (Flux, Flux Kontext, Qwen-Image, Flux2, Flux2Klein, LongCat-Image) goes through the mixin. Z-Image batches both branches into one call and is unaffected. The only other difference: do_true_cfg=True with negative_kwargs=None now raises a ValueError naming the branch instead of a TypeError from **None.
• Tests. A real TeaCacheHook driven through the mixin as CFG wave, non-CFG wave, CFG wave selects positive, negative, positive, positive, negative. On [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 the last wave selects negative, positive. Separate tests cover the per-path stamping, restore-on-error, and the unchanged parity fallback for unstamped callers. Further tests cover hint-over-stamp precedence and an invalid stamp. The 5 behaviour tests fail on [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 and pass here. The related suites (now including MammothModa2's TeaCache tests) pass 350 vs 340 on main; the remaining failures are the same on both and need a GPU or a full CPU build.
• Lint. pre-commit passes. mypy on the touched files: 3 errors vs 5 on main, none new.

Known limits: request boundaries under step execution (the second half of #8482) are not addressed here. There was no GPU run, because request-mode decisions don't change and the step-execution path is covered by the CPU test.

[hsliuustc0106]

This PR touches vllm_omni/diffusion/, tests/diffusion/ (5 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be:

@princepride @Isotr0py @SamitHuang

Could one of you take a look when you get a chance? Thanks!

[vllm-omni-review-bot]

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

Module owners: @princepride @Isotr0py @SamitHuang

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

@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 zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 54a7abc3-b66e-4b7f-93ec-67500d6f0d71) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 120s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 7bf61e2b-fd49-441c-8cff-7c5af1d3998f) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 600s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 55b9456f-8250-4867-906f-1c304896e39e) — check zcode login and the CLI version; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

0 actionable finding(s).

CI at 0faff74962e7 (2026-10-10T07:26:15.033528+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/zcode/GLM-5.3-Flash could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

CFGParallelMixin.predict_noise_maybe_with_cfg sets transformer.do_true_cfg for the call, then _predict_noise_for_cfg_branch writes transformer.cfg_branch to "positive" or "negative" around every predict_noise on the sequential, CFG-parallel, and no-CFG paths and restores the previous value in a finally. TeaCacheHook chooses the cache context from the extractor hint teacache_branch when that hint is present, otherwise from this stamp, and only then from the previous rank or forward-count parity. Under step execution a CFG wave, an odd non-CFG wave, and another CFG wave therefore stay on positive, negative, positive, positive, negative instead of swapping the last wave. Request-mode two-call CFG still matches the old parity because the hook counter restarts at zero each request.

Change flow

flowchart LR
  waves["[EXISTING] CFG wave, non-CFG wave, CFG wave"]:::existing
  mixin["[CHANGED] CFGParallelMixin.predict_noise_maybe_with_cfg"]:::changed
  stamp["[NEW] _predict_noise_for_cfg_branch"]:::new
  hook["[CHANGED] TeaCacheHook hint, then stamp, then parity"]:::changed
  states["[EXISTING] teacache_positive and teacache_negative"]:::existing
  waves --> mixin --> stamp --> hook --> states
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

No actionable findings.

---

_{🤖 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!}

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 12 days

@fusheng-ji this pull request has had no human commit, comment or review since 2026-10-04. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[fusheng-ji]

This PR is still active. I've merged the latest main (b164fea) and rerun the CFG/TeaCache CPU checks: 43 passed, with no new mypy diagnostics. Seven MammothModa2 cases encounter the same no-GPU setup errors on main. This remains stacked on #8475, so please focus on the explicit branch-stamping change. @princepride @Isotr0py, could you review the CFG → non-CFG → CFG behavior and extractor-hint/stamp/fallback precedence when available? Request-boundary resets under step execution remain outside this PR.

[Bounty-hunter]

LGTM