[BugFix] Fix Wan2.2-VACE-Fun cache and serving lifecycle

[david6666666]

Purpose

This PR builds on #4249 after rebasing it onto current upstream/main and fixes the Wan2.2-VACE issues found while reproducing the supplied performance script with:

• MODEL_PATH=Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers
• INPUT_VIDEO=mkt_part_000.mp4 (61 frames)

I also validated Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers against the original alibaba-pai/Wan2.2-VACE-Fun-A14B checkpoint. Under the same vLLM-Omni pipeline, prompt, input video, seed, size, frame count, and 30 inference steps, both produced a byte-identical MP4. This lets the PR use the standard diffusers model layout and avoid adding original-format key remapping / .pth component loading code.

What changed

1. Register Wan22VACEPipeline in the Cache-DiT custom enabler registry so it uses the Wan2.2 block-level path instead of the generic transformer-level path.
2. Pass TaylorSeer calibrator_config through enable_cache_for_wan22() for both single-transformer and dual-transformer Wan2.2 paths.
3. Keep the API server's AsyncOmni engine context alive until the HTTP shutdown task completes. Previously omni_run_server_worker() exited build_async_omni immediately after serve_http() returned its shutdown task, which could close the engine while the server was still accepting requests.
4. Support dual-transformer Wan2.2 VACE denoising by switching between transformer and transformer_2 at the configured boundary timestep.
5. Include VACE conditioning blocks in layerwise offload discovery by declaring WanVACETransformer3DModel._layerwise_offload_blocks_attrs = ["vace_blocks", "blocks"].
6. Add a recipe and supported-model docs for Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.

Issue reproduction notes

Issue	Local result	Fix
#4632 Cache-DiT not working for Wan22VACEPipeline	Reproduced in code: #4249 had no Wan22VACEPipeline entry in CUSTOM_DIT_ENABLERS, so CacheDiTBackend.enable() would fall back to the generic path.	Added Wan22VACEPipeline -> enable_cache_for_wan22 and a registry regression test.
#4633 TaylorSeer silently disabled	Reproduced in code: enable_cache_for_wan22() did not resolve or pass calibrator_config, unlike the generic path and S2V custom path.	Added _resolve_calibrator_config() usage and passed the same calibrator into both ParamsModifier and cache_dit.enable_cache().
#4634 subsequent V2V requests return 504	Reproduced in code: the server worker awaited shutdown_task after exiting build_async_omni, so the engine context could shut down while HTTP serving continued.	Moved await shutdown_task inside the engine context and added a lifecycle regression test.
#4635 HSDP OOM on 64GB NPU	Could not reproduce the NPU allocator behavior on local NVIDIA B300. Code inspection found VACE-specific memory pressure: vace_blocks were not included in layerwise offload discovery.	Added VACE vace_blocks to layerwise offload discovery. The NPU allocator limit itself is outside this repo.

Validation

Unit / CPU tests

pytest tests/diffusion/cache/test_cache_dit.py tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py tests/diffusion/offloader/test_layerwise_backend.py tests/entrypoints/openai_api/test_video_server.py::test_server_worker_keeps_engine_alive_until_http_shutdown -q
# 27 passed

pre-commit run --files docs/models/supported_models.md recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md tests/diffusion/cache/test_cache_dit.py tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py tests/entrypoints/openai_api/test_video_server.py vllm_omni/diffusion/cache/cache_dit_backend.py vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py vllm_omni/entrypoints/openai/api_server.py
# passed

Local GPU validation

Platform: single NVIDIA B300 for offline runs and vllm serve.

30-step V2V equivalence test, 736x1280, 61 frames, seed=1, boundary_ratio=0.875:

original alibaba-pai/Wan2.2-VACE-Fun-A14B: 248.7s total, 8.3s/step
Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers: 248.6s total, 8.3s/step
SHA256(original mp4) == SHA256(Pyros13 mp4)
MAE=0, RMSE=0, max pixel diff=0, PSNR=inf, mean SSIM=1.0

Pyros13 smoke test after removing original-format loader code:

BASELINE: 28.2s total, 14.1s/step  # 2 steps

Cache-DiT evidence after the fix:

Using custom cache-dit enabler for model: Wan22VACEPipeline
Calibrator Config: TaylorSeer_O(1)

Online serving validation for #4634:

POST /v1/videos/sync request 1: HTTP 200
POST /v1/videos/sync request 2: HTTP 200
GET /health after both requests: HTTP 200

Validation videos are attached in the PR discussion.

[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.

[david6666666]

Validation update for Wan2.2-VACE using the supplied mkt_part_000.mp4 and alibaba-pai/Wan2.2-VACE-Fun-A14B.

Environment note: local validation was on NVIDIA B300 GPUs, so the Ascend/NPU allocator behavior from #4635 cannot be reproduced exactly here. I validated the vLLM-Omni code paths that are relevant to the fix.

30-step videos for each issue scenario:

• [Bug]: Cache-DiT not working for Wan22VACEPipeline (fallback to generic enabler, 0% speedup) #4632 Cache-DiT / custom Wan22VACEPipeline enabler: https://github.com/david6666666/vllm-omni/releases/download/codex-wan22-vace-validation-20260624-0159/issue4632_cache_dit_30step.mp4
• [Bug]: enable_cache_for_wan22 does not pass calibrator_config to cache_dit.enable_cache() (TaylorSeer silently disabled) #4633 Cache-DiT + TaylorSeer calibrator: https://github.com/david6666666/vllm-omni/releases/download/codex-wan22-vace-validation-20260624-0159/issue4633_taylorseer_30step.mp4
• [Bug]: V2V request causes DiffusionExecutor to close, subsequent requests return 504 #4634 online server lifecycle, request 1: https://github.com/david6666666/vllm-omni/releases/download/codex-wan22-vace-validation-20260624-0159/issue4634_server_req1_30step.mp4
• [Bug]: V2V request causes DiffusionExecutor to close, subsequent requests return 504 #4634 online server lifecycle, request 2: https://github.com/david6666666/vllm-omni/releases/download/codex-wan22-vace-validation-20260624-0159/issue4634_server_req2_30step.mp4
• [Bug]: HSDP load OOM on 64GB HBM devices (requires full model load before sharding) #4635 4-card HSDP + Ulysses-SP without explicit --hsdp-shard-size 4 --hsdp-replicate-size 2: https://github.com/david6666666/vllm-omni/releases/download/codex-wan22-vace-validation-20260624-0159/issue4635_hsdp_usp4_30step.mp4

Key local validation from the 30-step rerun:

• [Bug]: Cache-DiT not working for Wan22VACEPipeline (fallback to generic enabler, 0% speedup) #4632 Cache-DiT completed: CACHE_DIT: 169.4s total, 5.6s/step.
• [Bug]: Cache-DiT not working for Wan22VACEPipeline (fallback to generic enabler, 0% speedup) #4632 logs confirm Using custom cache-dit enabler for model: Wan22VACEPipeline and block-level cache contexts for num_inference_steps=30.
• [Bug]: enable_cache_for_wan22 does not pass calibrator_config to cache_dit.enable_cache() (TaylorSeer silently disabled) #4633 TaylorSeer completed: TAYLORSEER: 161.5s total, 5.4s/step.
• [Bug]: enable_cache_for_wan22 does not pass calibrator_config to cache_dit.enable_cache() (TaylorSeer silently disabled) #4633 logs confirm TaylorSeer enabled with order=1 and Calibrator Config: TaylorSeer_O(1).
• [Bug]: V2V request causes DiffusionExecutor to close, subsequent requests return 504 #4634 online server handled two sequential 30-step /v1/videos/sync requests: request 1 HTTP 200 in 258.9s, request 2 HTTP 200 in 257.2s, then /health returned HTTP 200.
• [Bug]: V2V request causes DiffusionExecutor to close, subsequent requests return 504 #4634 logs show both POST /v1/videos/sync HTTP/1.1 requests returned 200 OK; [AsyncOmni] Shutting down appeared only after the validation script terminated the server.
• [Bug]: HSDP load OOM on 64GB HBM devices (requires full model load before sharding) #4635 4-card HSDP command shape was --use-hsdp --ulysses-degree 4 on CUDA_VISIBLE_DEVICES=0,1,2,7, with no explicit HSDP shard/replicate args.
• [Bug]: HSDP load OOM on 64GB HBM devices (requires full model load before sharding) #4635 HSDP auto-resolved to replicate_size=1, shard_size=4, world_size=4.
• [Bug]: HSDP load OOM on 64GB HBM devices (requires full model load before sharding) #4635 logs show VACE HSDP sharding: Sharded 48 modules + root for FSDPWanVACETransformer3DModel.
• [Bug]: HSDP load OOM on 64GB HBM devices (requires full model load before sharding) #4635 30-step HSDP completed: BASELINE: 82.3s total, 2.7s/step.

Unit/regression checks already run:

• pytest tests/diffusion/cache/test_cache_dit.py tests/diffusion/offloader/test_layerwise_backend.py tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py tests/entrypoints/openai_api/test_video_server.py::test_server_worker_keeps_engine_alive_until_http_shutdown -q
• pytest tests/entrypoints/openai_api/test_video_server.py -q
• pre-commit run --files ...

[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.

[david6666666]

Follow-up from review/testing addressed through f8f8410:

• Switched the PR to use Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers instead of adding support for the original alibaba-pai/Wan2.2-VACE-Fun-A14B layout.
• Removed the original-format detection/key-remapping/.pth component loading code and the corresponding original-format unit tests.
• Updated docs and recipe to point at the diffusers checkpoint.
• Kept the actual issue fixes: VACE Cache-DiT registration, TaylorSeer calibrator passthrough, server lifecycle fix, dual-transformer VACE boundary switching, and VACE vace_blocks layerwise offload discovery.
• Dropped the optional None block-group fallback in layerwise_backend.py; supported VACE configs include vace_layers, so vace_blocks is initialized as a real ModuleList.

Why this replacement is safe from the local validation:

Input: mkt_part_000.mp4
Resolution: 736x1280
Frames: 61
Steps: 30
Seed: 1
boundary_ratio: 0.875

alibaba-pai/Wan2.2-VACE-Fun-A14B:          248.7s total, 8.3s/step
Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers:   248.6s total, 8.3s/step

SHA256(original mp4) == SHA256(Pyros13 mp4)
MAE=0, RMSE=0, max pixel diff=0, PSNR=inf, mean SSIM=1.0

Validation after cleanup:

pytest tests/diffusion/cache/test_cache_dit.py tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py tests/diffusion/offloader/test_layerwise_backend.py tests/entrypoints/openai_api/test_video_server.py::test_server_worker_keeps_engine_alive_until_http_shutdown -q
# 27 passed

pre-commit run --files docs/models/supported_models.md examples/offline_inference/vace/vace_video_generation.md recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md tests/diffusion/cache/test_cache_dit.py tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py tests/entrypoints/openai_api/test_video_server.py vllm_omni/diffusion/cache/cache_dit_backend.py vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py vllm_omni/entrypoints/openai/api_server.py
# passed

Pyros13 2-step smoke after removing original loader:
# BASELINE: 28.2s total, 14.1s/step

[hsliuustc0106]

Wan2.2-VACE-Fun bugfix looks correct. All gate checks pass, regression tests cover Cache-DiT, TaylorSeer calibrator, dual-transformer switching, layerwise offload, and server lifecycle issues documented in PR description.