Skip to content

[Bugfix][Diffusion] Fix Cosmos3 transfer caching artifacts with control-CFG consensus - #7738

Closed
yzhautouskay wants to merge 2 commits into
vllm-project:mainfrom
yzhautouskay:yzhautouskay/cosmos3_caching_transfer_fix
Closed

yzhautouskay wants to merge 2 commits into
vllm-project:mainfrom
yzhautouskay:yzhautouskay/cosmos3_caching_transfer_fix

Conversation

@yzhautouskay

@yzhautouskay yzhautouskay commented Sep 17, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Fix Cosmos3 Transfer caching artifacts by synchronizing SeaCache decisions across control-CFG branches. Original implementation in the PR: #6922

Control CFG strengthens adherence to a control hint (such as edge, depth, or segmentation) by contrasting predictions with and without that hint, while keeping the text prompt unchanged. Unlike text-CFG branches, these branches receive different vision inputs, so their SeaCache indicators can produce different skip decisions
Independent decisions can create mismatched residual refresh steps and extrapolation histories, distorting the branch differences amplified by guidance.

This change:

  • Caches only when all participating branches agree; otherwise all compute
  • Keeps branch-specific residuals separate while aligning refresh steps

Test Plan

  • Added unit tests for control-CFG consensus: branch histories, orders 0/1/2, CFG-parallel ranks, state reset, and pipeline inputs
  • Ran BF16 and FP8 E2E – generated six-arm videos: uncached / current cache / this PR fix (consensus)
  • Consensus keeps comparable caching (e.g. 26/50 cached steps per branch), so improvement is from aligned decisions – not less caching – no performance win change

Test Result

Current cache quality issue is input-dependent because control-CFG branch schedules can diverge very differently across examples, independent of the control type alone. Some show little degradation, while others develop severe temporal/block artifacts. Effect is amplified with FP8-quantized checkpoints. This PR resolves the observed caching artifacts, from mild flickering to severe degradation, in both BF16 and FP8 runs

Rows: BF16 (top), FP8 (bottom).
Columns, left to right: uncached / current cache / this PR fix (consensus):

seg_six_arms_fix.mp4
edge_six_arms_fix.mp4

vLLM Version: 0.29.0
vLLM-Omni Commit: 896f4466

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

Signed-off-by: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
@vllm-omni-review-bot

Copy link
Copy Markdown

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

Module owners: @Isotr0py @david6666666 @princepride

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

@yzhautouskay, 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.

Signed-off-by: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Sep 17, 2026
@yzhautouskay

Copy link
Copy Markdown
Contributor Author

Self-review note: FP8 runs with Cosmos3 Transfer + SeaCache implementation from this PR exposed a transfer quality issue when control_guidance>1.0 (control-CFG enabled) that also occur in some BF16 examples to a lesser extend. Unlike text-CFG branches, control-CFG branches receive different vision inputs, so their SeaCache indicators can produce different skip decisions Independent decisions can create mismatched residual refresh steps and extrapolation histories, distorting the branch differences amplified by guidance.

This PR addresses and ultimately solves this issue by synchronizing cache decisions across control-CFG branches, aligning residual refresh and extrapolation steps while retaining separate branch residuals.
I validated the fix BF16 and FP8 E2E generations across edge, depth, segmentation, blur controls. Comparison videos are included in the PR description

@yzhautouskay

Copy link
Copy Markdown
Contributor Author

Closing this PR in favor of #7939

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority high priority issue, needs to be done asap

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants