[xuechendi]

Purpose

Fix Cache-DiT on Wan2.2 when only one expert is loaded, and make the backend discover its cache targets instead of assuming pipeline.transformer.

The flow-shift / effective-boundary part of the expert split that this PR originally carried is now on main via #7335 (refresh from forward() with the resolved schedule), so after merging main this PR keeps only what #7335 does not cover.

Changes

• cachedit/model_specific.py enable_cache_for_wan22: in single-expert mode, cache and refresh whichever expert is actually loaded. With a Wan2.2 MoE checkpoint and boundary_ratio=1.0, Wan22Pipeline loads only transformer_2 and pipeline.transformer is None, so the old code crashed on pipeline.transformer.blocks and its refresh read pipeline.transformer. Both the runner refresh and [Bugfix][Diffusion] Make Cache-DiT effective when serving Wan2.2 T2V/I2V/S2V without num_inference_steps #7335's per-schedule refresh now use a getter for the active expert. Raise a clear ValueError when neither expert is loaded.
• cachedit/backend.py: an enabler that only returns a refresh func now gets its cache targets discovered from the pipeline's DiT modules (filtered by BlockAdapter.is_cached). Before, the target was pipeline.transformer, which is None for the low-noise-only run and also left out transformer_2 in the dual-expert case, so the second expert's hooks stayed installed after disable().
• wan2_2: resolve_wan_boundary_ratio() (engine override -> request -> 0.875) replaces the three copies in T2V / I2V / VACE forward() and the batch-settings check.

Test Plan

pytest tests/diffusion/cache tests/diffusion/models/wan2_2 tests/diffusion/test_diffusion_model_runner.py

New tests in tests/diffusion/cache/test_cache_dit.py: test_wan22_custom_enabler_low_noise_expert_only (enable, runner refresh and per-schedule refresh all target transformer_2), test_wan22_custom_enabler_without_any_transformer_raises, test_resolve_wan_boundary_ratio.

Test Result

On Intel Arc Pro B70 (XPU), head bd20d0cad:

• tests/diffusion/cache/test_cache_dit.py: 27 passed.
• cache + wan2_2 + model runner + offloader + Wan quantization suites: 1036 passed, 16 failed; the same 16 fail on origin/main (98c204144) on this node (offline model downloads / environment), none in the touched code.

(anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

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, CODEOWNERS; @princepride via module of the changed files, module named in the PR description, CODEOWNERS; @SamitHuang via module of the changed files, module named in the PR description, CODEOWNERS

@xuechendi, 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 triage note

Resolved as of bd20d0cad174: the high-priority or low-quality signal noted on an earlier commit no longer applies.

[LinzeShi]

The runner refreshes the cache before forward() reads the request's flow_shift, so this still uses the previous value. A request that overrides it (e.g. flow_shift=12.0 as in the 480p I2V example) on a server started with 5.0 still gets the 17/23 split, but denoises with a different one.

[xuechendi]

Good catch, thanks. The refresh API only gets (pipeline, num_inference_steps), so the split couldn't see the request. Fixed in 3fee614:

• The runner now calls an optional pipeline.prepare_cache_refresh(first_req) right before cache_backend.refresh(...).
• Wan22Pipeline (skipped for DMD) and Wan22I2VPipeline use it to apply the request's sample_solver / flow_shift to the scheduler. They go through a new apply_wan_request_scheduler() helper, which forward() now reuses instead of duplicating the logic. Calling it again in forward() for the same request is a no-op.
• VACE overrides the hook as a no-op. Its forward() ignores per-request scheduler settings, so its split already matches.

New tests:

• test_wan22_step_split_follows_request_flow_shift_override: a server started with 5.0 gets a request with flow_shift=12.0, and the split matches the 12.0 schedule instead of 17/23.
• test_refresh_cache_prepares_pipeline_before_backend_refresh: the hook runs before the backend refresh.

[LinzeShi]

Could we add a test for the low noise only case (transformer is None, transformer_2 loaded)? The existing taylorseer test only covers transformer_2=None, and the new test only covers the step split.

[xuechendi]

Added in 3fee614:

• test_wan22_custom_enabler_low_noise_expert_only: with transformer=None and transformer_2 loaded, the cache is installed on transformer_2 and the returned refresh refreshes transformer_2.
• test_wan22_custom_enabler_without_any_transformer_raises: covers the ValueError when neither expert is loaded.

[xuechendi]

@LinzeShi @hsliuustc0106 , PTAL
fixed the issue @LinzeShi raised. Verified e2e

Setup: Wan2.2-T2V-A14B-Diffusers , server started with --boundary-ratio 0.875 --flow-shift 5.0.
Each request is 480×832, 17 frames, 40 steps, seed 42, with a per-request flow_shift.

Expected split, from build_wan_scheduler("unipc", fs).set_timesteps(40, shift=fs) with boundary 875:

• flow_shift=5.0: 17 / 23
• flow_shift=12.0: 26 / 14

With this PR (3fee614):

request 1: flow_shift=5.0 (same as server default) -> HTTP 200, 78 s

(APIServer pid=689230) INFO 10-15 03:52:09 [serving_video.py:470] Boundary ratio parse: request=0.875 gen_params=0.875
[10-15 03:52:09] [Cache-DiT] ✅ Refreshed cache context: blocks_136994329969776, Collected Context Config: DBCache_F1B0_W4I1M0MC3_R0.24_N17_CFG1, Calibrator Config: None
[10-15 03:52:09] [Cache-DiT] ✅ Refreshed cache context: blocks_136994306414624, Collected Context Config: DBCache_F1B0_W2I1M20MC3_R0.24_N23_CFG1, Calibrator Config: None

request 2: flow_shift=12.0 (per-request override) -> HTTP 200, 67 s

(APIServer pid=689230) INFO 10-15 03:53:27 [serving_video.py:470] Boundary ratio parse: request=0.875 gen_params=0.875
[10-15 03:53:27] [Cache-DiT] ✅ Refreshed cache context: blocks_136994329969776, Collected Context Config: DBCache_F1B0_W4I1M0MC3_R0.24_N26_CFG1, Calibrator Config: None
[10-15 03:53:27] [Cache-DiT] ✅ Refreshed cache context: blocks_136994306414624, Collected Context Config: DBCache_F1B0_W2I1M20MC3_R0.24_N14_CFG1, Calibrator Config: None

The high-noise (W4…M0) and low-noise (W2…M20) contexts get N17/N23 at flow_shift=5.0 and N26/N14 at the 12.0 override. Both match the schedule forward() denoises with.

[LinzeShi]

LGTM

[xuechendi]

@hsliuustc0106 PTAL

[hsliuustc0106]

Reviewed head 3fee614a0058f184915e60d271192dd4f6eedc0b. The request-level flow-shift ordering and low-noise-only target selection are addressed.

One pre-existing correctness case remains in the same cache-split path: when the engine is started without --boundary-ratio, self.boundary_ratio is None. The new prepare_cache_refresh hook applies solver/shift only, while _split_wan22_inference_steps treats a missing boundary as “all high-noise”. For a 40-step dual-expert run this refreshes the contexts with 40/0 steps. forward instead uses the request's boundary or defaults to 0.875, giving 17/23 at flow_shift=5.0. A request-level boundary override has the same mismatch when no engine boundary is set.

Please share resolution of the effective boundary (engine override → request → 0.875) between cache preparation and denoising, and cover both an omitted boundary and a request override. Keep the startup value separate from request-effective state so successive requests can choose different boundaries. This is an existing gap, not a regression introduced by this PR.

Validation: source/call-chain review and static parsing; no model execution or independent hardware measurements.

[xuechendi]

@hsliuustc0106 thanks, fixed in c7a6424.

• New resolve_wan_boundary_ratio(engine_boundary_ratio, sampling_params) in pipeline_wan2_2.py resolves engine override → request → 0.875. forward() in T2V, I2V and VACE and _validate_generation_batch all use it now. Each pipeline used to repeat this inline.
• prepare_cache_refresh (T2V, I2V, VACE) stores the result in self._request_boundary_ratio, and _split_wan22_inference_steps reads that before pipeline.boundary_ratio. So without --boundary-ratio, a 40-step run at flow_shift=5.0 now splits 17/23 instead of 40/0. The startup self.boundary_ratio is never written, so each request can pick its own boundary.

Tests (tests/diffusion/cache/test_cache_dit.py):

• test_wan22_step_split_uses_effective_boundary: runs on all three pipelines for three cases: boundary omitted everywhere (→ 0.875), request override with no engine boundary, and engine override winning over the request. Each case checks the split against the schedule forward() denoises with, and checks that pipeline.boundary_ratio is unchanged.
• test_wan22_step_split_follows_successive_request_boundaries: three requests in a row (0.95 / omitted / 0.8) on an engine without a boundary get three different splits, each matching its own schedule.

B60 (XPU): I ran every non-e2e test file that imports the changed modules (37 files). Result: 1590 passed, 80 failed. The same 80 tests also fail on 3fee614 (1580 passed there), e.g. MammothModa2 model path missing on the node and test_wan_device_reduction. Nothing new fails.

[hsliuustc0106]

Reviewed current head 7ae90faaa6c4f5c9b3249d468c737c6bc77f8ebb.

The request-specific scheduler/flow-shift ordering, effective-boundary precedence and successive-request isolation address the earlier review concerns. I also inspected low-noise-only target selection and cleanup of both experts, the runner integration, and the corresponding regression tests. The named scheduler/boundary regressions and author-provided before/after serving logs support the correction; I did not independently execute hardware tests.

Current-head GitHub Actions and docs checks pass. The ready-triggered CUDA pipeline build #17778 reports success for this exact SHA; its entry/bootstrap and ready routing were checked against the trusted repository configuration. This is pipeline-level CI evidence, not a claim that I inspected individual job logs or independently validated XPU/AMD behavior. AMD and Intel CI report failures whose causes I have not verified; CUDA does not establish their correctness.

No remaining actionable correctness finding in the reviewed changes. Final squash-and-merge remains with the maintainer.

[hsliuustc0106]

fix conflicts

[xuechendi]

@hsliuustc0106 conflicts fixed in bd20d0c (merge of main).

#7335 landed on main with the same flow-shift / effective-boundary fix, done from forward() (refresh_cache_for_denoise_schedule). I resolved toward main's version and dropped this PR's prepare_cache_refresh runner hook, _request_boundary_ratio and apply_wan_request_scheduler. The PR is now limited to what main still lacks:

• single-expert Cache-DiT on whichever expert is loaded (low-noise-only transformer_2 with boundary_ratio=1.0), for both the runner refresh and [Bugfix][Diffusion] Make Cache-DiT effective when serving Wan2.2 T2V/I2V/S2V without num_inference_steps #7335's per-schedule refresh
• backend cache-target discovery, so disable() also unhooks transformer_2
• the shared resolve_wan_boundary_ratio helper

The title and description are updated. The diff is smaller than the one you approved, so please take another look.

CI: the AMD failure was test_qwen_image_fused_qk_norm_rope.py, already fixed on main by #9229 and pulled in by this merge. The Intel failure was test_qwen2_5_omni_expansion failing to start because xpu:0 had only 11.7/31.9 GiB free (another process held the memory), so it's infra, not this PR.

B60 at bd20d0c: test_cache_dit.py 27 passed; cache + wan2_2 + model runner + offloader + Wan quant suites 1036 passed, and the 16 failures there are the same as on origin/main on that node.

[hsliuustc0106]

Reviewed current head bd20d0cad1745380aa2f05134cf0fce8a952f066 after the main merge and narrower scope.

Revalidated the remaining six-file diff: low-noise-only cache enable/refresh targets transformer_2, backend target discovery also covers cleanup, and the shared boundary resolver preserves engine override → request → default precedence. The request-specific schedule refresh now comes from merged #7335; the old runner hook is no longer duplicated. The earlier flow-shift and low-noise-only review concerns are addressed in this implementation and the named regression tests.

Evidence: inspected the production paths and test_wan22_custom_enabler_low_noise_expert_only, missing-expert, boundary-precedence and shifted-schedule tests; reused the earlier Wan request-level evidence and read the author's current-head 27-test receipt. I did not execute tests or benchmarks locally.

Current-head GitHub checks and CUDA ready pipeline #17845 pass. Intel pipeline #10608 failed; its job-level cause is unavailable here and is not classified as infrastructure or treated as a passing platform test. This review makes no new Intel correctness or performance claim.

[hsliuustc0106]

Re-checked after #7335 landed: this PR is still dirty / conflicted against main.

Please rebase onto current main and resolve the overlap in model_specific.py / pipeline_wan2_2*.py carefully — #7335 just changed the same expert-budget path. After rebase + green CI (CUDA / GHA / NPU), ping me and I’ll re-approve. Prior review points look addressed; no new logic blockers beyond the conflict.