[None][feat] Qwen-Image TeaCache/Cache-DiT Support - #16339
Conversation
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59030 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughQwen-Image now supports TeaCache and Cache-DiT during visual generation. The pipeline manages cache state across true CFG passes, Cache-DiT receives a dedicated enabler, tests cover integration and coefficient selection, and documentation records support and tuning guidance. ChangesQwen-Image cache acceleration
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant QwenImagePipeline
participant TeaCache
participant CacheDiT
participant QwenImageTransformer2DModel
Request->>QwenImagePipeline: start generation
QwenImagePipeline->>TeaCache: reset cache state
QwenImagePipeline->>CacheDiT: refresh cache context
QwenImagePipeline->>QwenImageTransformer2DModel: run conditional forward
QwenImagePipeline->>QwenImageTransformer2DModel: run unconditional forward
QwenImagePipeline->>TeaCache: collect cache statistics
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py`:
- Around line 347-376: Verify Cache-DiT’s DBCacheConfig behavior for
Qwen-Image’s two separate CFG passes, then update
enable_cache_dit_for_qwen_image() to pass enable_separate_cfg=True to
_build_db_cache_config() if automatic isolation is not guaranteed; otherwise
document or preserve the confirmed default behavior. In pipeline_qwen_image.py,
retain the existing _cache_branch handling unless verification shows it must be
coordinated with Cache-DiT.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a7e0347c-bdc1-421a-a17f-7b3fcdb45b39
📒 Files selected for processing (5)
docs/source/models/visual-generation.mdtensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.pytensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytests/unittest/_torch/visual_gen/test_cache_dit.pytests/unittest/_torch/visual_gen/test_teacache.py
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59034 [ run ] triggered by Bot. Commit: |
|
PR_Github #59030 [ run ] completed with state |
| timestep = t.expand(latents.shape[0]).to(latents.dtype) | ||
|
|
||
| # Conditional (positive) branch. | ||
| if cache_acc is not None: |
There was a problem hiding this comment.
Should we use the same conditions as above?
There was a problem hiding this comment.
Fixed, along with the other cache guard inconsistencies, by introducing cache_on = self.cache_enabled before the loop. All branch guards now use the same condition.
| latents = latents.to(latents_dtype) | ||
|
|
||
| if getattr(self, "rank", 0) == 0: | ||
| if cache_acc is not None and cache_acc.is_enabled(): |
There was a problem hiding this comment.
Maybe we can have something like self.cache_enabled()
There was a problem hiding this comment.
Done, added cache_enabled as a property on QwenImagePipeline and used cache_on = self.cache_enabled before the loop so all sites share a single consistent check.
|
PR_Github #59034 [ run ] completed with state |
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
cef291c to
cd065a7
Compare
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59525 [ run ] triggered by Bot. Commit: |
|
PR_Github #59525 [ run ] completed with state |
zhenhuaw-me
left a comment
zhenhuaw-me
left a comment
There was a problem hiding this comment.
LGTM with some naming issues.
Please always feel free to resolve the threads if the changes reflect the review comments.
| # Batched CFG: default enable_separate_cfg False (single cond+uncond batch). | ||
| # Qwen-Image uses non-batched CFG (two separate forward passes), | ||
| # so it has no enable_separate_cfg flag here. | ||
| _WAN_CFG_DEFAULT = False | ||
| _FLUX_CFG_DEFAULT = False | ||
| _LTX2_CFG_DEFAULT = False |
There was a problem hiding this comment.
| # Batched CFG: default enable_separate_cfg False (single cond+uncond batch). | |
| # Qwen-Image uses non-batched CFG (two separate forward passes), | |
| # so it has no enable_separate_cfg flag here. | |
| _WAN_CFG_DEFAULT = False | |
| _FLUX_CFG_DEFAULT = False | |
| _LTX2_CFG_DEFAULT = False | |
| # Batched CFG: default enable_separate_cfg False (single cond+uncond batch). | |
| # Qwen-Image uses non-batched CFG (two separate forward passes), | |
| # so it has no enable_separate_cfg flag here. | |
| _WAN_BATCHED_CFG_DEFAULT = False | |
| _FLUX_BATCHED_CFG_DEFAULT = False | |
| _LTX2_BATCHED_CFG_DEFAULT = False |
(I should raise this in the previous round review, sorry) The semantic is "batch_cfg", in this case, can we rename to like this?
| f"W={db_cfg.max_warmup_steps}", | ||
| ) | ||
|
|
||
| disable_target = cache_dit.enable_cache( |
There was a problem hiding this comment.
(still naming confusions...) In first read, it's confusing to assign the return values of enable_xyz() to objects named like disable_target. It seems the return values is the original pipeline according to https://github.com/vipshop/cache-dit/blob/ad9335fdcc7d648b50a7d4ff46b1f25e2abdaf45/src/cache_dit/caching/cache_interface.py#L82 . Did I misread anything?
5 participants
Summary by CodeRabbit
New Features
Documentation
residual_diff_threshold.Description
Enables both TeaCache and Cache-DiT support for Qwen-Image.
Example quality/perf comparison on B200 (umb-b200-237):
"a cat sitting on a bench in a garden"
Default:
Configs: 1328×1328, 30 steps, seed 42, all other params set at default value
E2E pipeline time: 6.167s
TeaCache:
Configs: 1328×1328, 30 steps, seed 42, teacache_thresh=0.2, all other params set at default value
E2E pipeline time: 4.493s (1.37x default)
LPIPS score against default: 0.0210
Cache-DiT:
Configs: 1328×1328, 30 steps, seed 42, residual_diff_threshold=0.16 (lower than default for better quality), all other params set at default value
E2E pipeline time: 4.390s (1.40x default)
LPIPS score against default: 0.0372
Our findings show that the LPIPS scores are low enough that we could use more aggressive TeaCache/Cache-DiT configs to enable more caching than the example shown above, leading to faster E2E perf.
Test Coverage
test_teacache.py:TestQwenImageTeacacheTable: 3 new unit tests covering Qwen-Image coefficient resolution from checkpoint path, rejection of unrelated paths, and user-supplied coefficients taking precedence over the built-in table.test_cache_dit.py:TestCacheDiTRealPipelineForward: 1 new GPU integration test verifying that Cache-DiT skips at least one transformer block during a real 16-step Qwen-Image forward pass.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.