Repository navigation
[Hardware][Intel] Wan VAE: run wan_vae_fastpath on XPU - #8968
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/diffusion/index.md. Module owners: @david6666666 @princepride @Isotr0py Routing: @david6666666 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; @Isotr0py via module of the changed files, module named in the PR description @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. |
|
Author self-review What I checked on 103ddff:
|
|
Reviewed head
Reported CI passes; diff hygiene passes. Review was static: I did not execute fork code locally or independently verify CUDA/XPU correctness and performance. |
hsliuustc0106
added
the
high priority
ChatGPT Codex Connector
Wan VAE performance on Intel XPU, built on the existing CUDA-only
wan_vae_fastpath instead of a parallel XPU implementation.
wan_vae_fastpath on XPU
- Gate the fused Triton kernels and the registry install on
KERNEL_DEVICE_TYPES ({"cuda", "xpu"}, in diffusion/layers/vae_norm_epilogue)
instead of `is_cuda`. The kernels are plain Triton, and the conv
spatial-padding rewrite is already guarded by a per-shape bit-exactness
probe, so it adapts to oneDNN as it does to cuDNN.
- Keep the conv bias inside the convolution on XPU at the lossless level
(`_may_defer_conv_bias`). Deferring it to the following norm/residual kernel
is exact on CUDA, where ATen adds the bias after cuDNN rounds, but oneDNN
adds it inside the fp32 accumulator; deferring changed ~27% of bf16/fp16
outputs by an ulp. channels_last (tolerance-based) still defers everywhere.
Spatial-shard decode without tiling
- `is_distributed_enabled()` requires use_tiling only for "tile" mode,
`registry.initialize_model` auto-enables tiling only for "tile" mode, and
DistributedAutoencoderKLWan._decode routes spatial_shard_* before the tiling
check: the sharded decode is full-frame. A VAE without distributed state
(constructed directly) keeps the parent decode. In spatial-shard modes the
encode is distributed only with explicit --vae-use-tiling (documented).
Tests: the fast-path kernel tests are device-agnostic and run on XPU (B60
marks); XPU skips the bias-deferral mechanism tests and the fp32 fused-SiLU
kernel case (never taken at the lossless level) with a reason and adds
test_xpu_lossless_keeps_conv_bias_in_the_conv.
Measured on Intel Arc Pro B70 (lossless = bit-exact vs diffusers):
- Wan2.2-T2V-A14B VAE decode, 720p/81f bf16, 1 card: 17.97 -> 10.44 s
un-tiled (1.72x), 29.31 -> 17.90 s tiled (1.64x), rel_l2 0.
- Wan2.2-T2V-A14B e2e 720p/81f/20 steps, CFG2 x TP2, tiled VAE, module
offload: vae.decode 30.54 -> 21.61 s, request 708.4 -> 702.1 s, identical
output video.
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
103ddff to
0a14ba4
Compare
xuechendi
changed the title
|
@hsliuustc0106 thanks, both addressed in 0a14ba4. [P1] Decode without distributed initialization.
[P2] Automatic tiling.
I also updated the PR title and description (the norm compile was dropped from this PR earlier). |
|
Reviewed head
Validation: source/call-chain review and static parsing at the pinned head; no CUDA/XPU model execution or numerical/performance reproduction. |
|
@hsliuustc0106 thanks, both addressed in 6b002e3. [P2] Chunked decode without tiling.
[P2] Partial VAE group fallback.
The VAE parallelism guide's spatial-shard section now says when the fallback applies and what the chunked route does. Tests on XPU (Arc Pro B70), commit 6b002e3: the 51 test files that import the Wan VAE, the spatial-shard code, the distributed VAE executor, the registry or |
…k without tiling spatial_shard_* no longer forces vae_use_tiling, so two decode routes still depended on it: - _decode_with_chunks entered the sharded decode only under use_tiling; a fresh instance decoded the full frame on every rank. Route spatial_shard_* chunked decode before the tiling check. - With a partial VAE group (vae pp size != DiT group size) the shard gate falls back, but the parent _decode then took the local decode instead of the tile-parallel executor. Keep that fallback for latents larger than a tile, independent of native tiling. Extend test_distributed_chunked_routing to use_tiling=False and add test_partial_vae_group_decode_falls_back_to_tile_parallel. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
…tial_shard_* Spatial-shard modes may not auto-enable VAE tiling (see vllm-project#8968). Without tiling the encode runs locally and VLLM_OMNI_WAN_VAE_ENCODE_SPLIT has no effect, so the serving example and the requirements list now pass --vae-use-tiling explicitly. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
The four earlier routing/initialization findings are addressed in the source at 1dbbee11c. The supplied B70 parity and stage timings are useful; I have not independently executed them.
- Please complete the repeatability evidence for the claimed XPU decode speedup. Keep the documented B70, Wan2.2-T2V-A14B checkpoint, software, 720p/81f bf16 input and topology fixed, and provide runnable component benchmark commands plus paired
off/losslessraw timings after stated warm-up, with repeat count and spread rather than only best-of-two. For the end-to-end claim, repeat the paired 20-step CFG2 × TP2 comparison and report decode-stage and request timings separately, keeping fresh-server/cold and warmed samples separate. Bind results to this head or identify the measured SHA and unchanged relevant source. Retain the identical-output check forlossless; if retainingchannels_lastclaims, report that opt-in arm separately with its explicit numerical tolerance. Acceptance is a reproducible paired comparison showing the gain and its variability without changing workload or quality criteria. This is a performance-evidence gap, not a newly demonstrated code defect.
…tial_shard_* Spatial-shard modes may not auto-enable VAE tiling (see vllm-project#8968). Without tiling the encode runs locally and VLLM_OMNI_WAN_VAE_ENCODE_SPLIT has no effect, so the serving example and the requirements list now pass --vae-use-tiling explicitly. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
…tial_shard_* Spatial-shard modes may not auto-enable VAE tiling (see vllm-project#8968). Without tiling the encode runs locally and VLLM_OMNI_WAN_VAE_ENCODE_SPLIT has no effect, so the serving example and the requirements list now pass --vae-use-tiling explicitly. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
…pu 3.8 triton-xpu 3.8 (the vLLM v0.32.0 XPU image, vllm-project#9131) lowers tl.math.div_rn to a plain fdiv, which IGC compiles to its approximate divide: ~29% of fp32 quotients differ from ATen's x / d by an ulp, so the fp32 bit-exact RMSNorm/encoder/decoder kernel tests fail in the Intel CI. triton-xpu 3.7 lowered it to OpenCL's correctly rounded divide_cr. Add vae_norm_epilogue.div_rn: on XPU it calls divide_cr directly (exact and as fast as the approximate divide on 3.7.2 and 3.8.0), elsewhere it is tl.math.div_rn, so CUDA is unchanged. libdevice div_rn was not used: it is a software routine on XPU that made the RMSNorm epilogue ~30x slower and the lossless decode slower than off. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
- Select the device with torch.accelerator, so the benchmark runs on XPU as well as CUDA (--profile stays CUDA-only). - --tiling decodes with the VAE's spatial tiling on one GPU. - Report the median of --iters decodes with min, max, standard deviation and every raw sample; the speedup compares medians. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
|
@hsliuustc0106 thanks. Below: a CI fix found while re-measuring, and the component decode evidence at the new head. I haven't re-run the end-to-end comparison yet; it's covered at the end. Intel CI fix (a7cf00a)
The AMD failures (Qwen-Image / MammothModa2 fused QK-norm-RoPE) are in tests that import nothing this PR changes. #9229 fixed the Qwen-Image ones on main, and a1757fe merges main. Component decode benchmark, head a1757fe
python benchmarks/diffusion/bench_wan_vae_decode.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--size 1280x720 --frames 81 --dtype bf16 --fast-path off,lossless,channels_last --warmup 1 --iters 5 [--tiling]Median decode per run, in seconds. Within-run stdev is 2-36 ms everywhere (≤0.2%).
Raw samples (ms), un-tiledRaw samples (ms), tiledThe un-tiled Tests at a1757fe (B70, same image)
End to end |
xuechendi
added
ready
The two fast-path kernel test files were scheduled on XPU in full: about 8 minutes on one B70 (1030 cases). Intel CI only needs to show that the fast path is bit-exact on XPU, so the file-level marks go back to CUDA only and an _XPU_CI mark selects 88 cases: the div_rn and RMSNorm bit-exactness checks, the fused SiLU verdict, the XPU conv-bias rule, the lossless decoder and encoder end to end, the Omni VAE decode and a channels_last smoke test. 53 s wall on a B70 with a cold Triton cache. Every test still runs on XPU when invoked directly. Signed-off-by: Chendi Xue <chendi.xue@intel.com>
xuechendi
added
ready
|
Holding approve for two items:
Also please check why Intel CI is red on this XPU-focused PR. Happy to re-review once (1)+(2) are addressed. |
…-gather halos on XPU
Spatial-shard encode
- Under vae_parallel_mode=spatial_shard_{height,width}, DistributedAutoencoderKLWan._encode
shards the encoder along the split dimension with the existing halo modules, before the
native-tiling check (as _decode does for decode since vllm-project#8968). It mirrors diffusers' untiled
_encode loop and all-gathers the moments to every rank. It runs on a parameter-sharing copy
of vae.encoder kept outside the module tree, so state_dict is unchanged and the original
encoder still serves the fallback.
- Only taken when the split extent divides into stride-aligned shards
(world_size * 2**downsamples); otherwise, with the encoder fast path installed, or with a
partial VAE group, encode takes the regular path (tile/rank-split parallel when use_tiling).
- Fix: _patch_decoder_modules never replaced nn.ZeroPad2d, so the encoder's stride-2
downsample padded every rank's edge. The decoder has no ZeroPad2d.
All-gather halo exchange on XPU
- A batched isend/irecv halo costs ~1.3 ms on XCCL against ~0.1-0.2 ms for an all-gather of
the same rows; XPU now exchanges halos with one all_gather_into_tensor. Other platforms keep
point-to-point.
Measured on Intel Arc Pro B70, Wan2.2-A14B VAE bf16, 720x1280, 81 frames, I2V condition,
median of 3 (rel_l2 vs a 1-card untiled encode/decode):
- encode 4 ranks: tile 4.49 s (9.5e-3), rank split 3.17 s (3.0e-3), spatial_shard_width 3.06 s (1.2e-3)
- encode 8 ranks: tile 2.32 s (9.5e-3), rank split 1.63 s (3.8e-3), spatial_shard_width 1.83 s (1.2e-3)
- decode 4 ranks: tile 8.03 s -> spatial_shard_width 5.42 s (P2P halos 7.13 s)
- decode 8 ranks: tile 4.20 s -> spatial_shard_width 3.19 s (P2P halos 4.47 s)
Tests: CPU gloo parity (fp64, Wan2.1 and residual Wan2.2 VAEs, P2P and all-gather halos,
max diff < 1e-10), gate/dispatch unit tests, and a real-weights multi-card encode parity test.
…-gather halos on XPU
Spatial-shard encode
- Under vae_parallel_mode=spatial_shard_{height,width}, DistributedAutoencoderKLWan._encode
shards the encoder along the split dimension with the existing halo modules, before the
native-tiling check (as _decode does for decode since vllm-project#8968). It mirrors diffusers' untiled
_encode loop and all-gathers the moments to every rank. It runs on a parameter-sharing copy
of vae.encoder kept outside the module tree, so state_dict is unchanged and the original
encoder still serves the fallback.
- Only taken when the split extent divides into stride-aligned shards
(world_size * 2**downsamples); otherwise, with the encoder fast path installed, or with a
partial VAE group, encode takes the regular path (tile/rank-split parallel when use_tiling).
- Fix: _patch_decoder_modules never replaced nn.ZeroPad2d, so the encoder's stride-2
downsample padded every rank's edge. The decoder has no ZeroPad2d.
All-gather halo exchange on XPU
- A batched isend/irecv halo costs ~1.3 ms on XCCL against ~0.1-0.2 ms for an all-gather of
the same rows; XPU now exchanges halos with one all_gather_into_tensor. Other platforms keep
point-to-point.
Measured on Intel Arc Pro B70, Wan2.2-A14B VAE bf16, 720x1280, 81 frames, I2V condition,
median of 3 (rel_l2 vs a 1-card untiled encode/decode):
- encode 4 ranks: tile 4.49 s (9.5e-3), rank split 3.17 s (3.0e-3), spatial_shard_width 3.06 s (1.2e-3)
- encode 8 ranks: tile 2.32 s (9.5e-3), rank split 1.63 s (3.8e-3), spatial_shard_width 1.83 s (1.2e-3)
- decode 4 ranks: tile 8.03 s -> spatial_shard_width 5.42 s (P2P halos 7.13 s)
- decode 8 ranks: tile 4.20 s -> spatial_shard_width 3.19 s (P2P halos 4.47 s)
Tests: CPU gloo parity (fp64, Wan2.1 and residual Wan2.2 VAEs, P2P and all-gather halos,
max diff < 1e-10), gate/dispatch unit tests, and a real-weights multi-card encode parity test.
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
1. What outcome do you want to achieve?
apply the same wan_vae_fastpath that CUDA already uses, on XPU.
XPU users should get the lossless decode/encode speedup by default, and XPU should share one maintained code path with CUDA.
2. What is the highest-leverage variable for that outcome?
Whether wan_vae_fastpath runs on XPU at all. Previously, the registry install and every Triton kernel are gated on CUDA, so XPU always runs the upstream diffusers VAE. That VAE spends much of its time on copies and elementwise passes around the convolutions, which is exactly the work the fast path removes.
3. How will this work move that variable?
4. Whose customer or user needs does this address, and what are those needs?
Users serving Wan2.2 T2V/I2V video generation on Intel XPU [e.g. the team, customer or deployment, if public]. They need:
Purpose
Wan VAE performance on Intel XPU. This PR builds on the existing, CUDA-only
wan_vae_fastpath(#7056 decode, #8067 encode) instead of adding a parallel XPUimplementation.
wan_vae_fastpathon XPUregistry._apply_wan_vae_fastpath_if_enabledreturned on non-CUDA, and every Triton kernel gated onx.is_cuda.KERNEL_DEVICE_TYPES = {"cuda", "xpu"}indiffusion/layers/vae_norm_epilogue.py). On XPU,losslesskeeps the conv bias inside the convolution (see below), so it stays bit-exact.spatial_shard_*decode--vae-use-tiling, andregistry.initialize_modelforced it on, although the sharded decode is full-frame and ignores tiles.is_distributed_enabled()requiresuse_tilingand the registry auto-enables it only fortilemode.DistributedAutoencoderKLWan._decoderoutesspatial_shard_*before the tiling check; a VAE without distributed state keeps the parent decode. In spatial-shard modes the encode is distributed only with explicit--vae-use-tiling(documented in the VAE parallelism guide).Spatial-shard decode still runs without the fast path.
install_wan_vae_fastpathrefusesspatial_shard_*modes:wan_spatial_shardswapsevery causal conv for a halo-exchanging one and rebinds the same block forwards.
--vae-use-tiling.** --vae-fast-path has three levels. **
Test Plan
vLLM Version: v0.31.0 (
vllm/vllm-openai-xpu:v0.31.0base ofdocker/Dockerfile.xpu)vLLM-Omni Commit:
0a14ba414onorigin/main0d4d09cf7Hardware: Intel Arc Pro B70 (31.9 GiB), torch 2.14.0+xpu, triton-xpu 3.7.2,
diffusers 0.40.0. Image built from this branch with
docker/Dockerfile.xpu.pytest tests/diffusion/distributed/test_autoencoder_kl_wan.py \ tests/diffusion/distributed/test_wan_vae_fastpath_install.py \ tests/diffusion/distributed/test_wan_vae_encoder_fastpath.py \ tests/diffusion/distributed/test_wan_vae_fastpath_verdicts.py \ tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py \ tests/diffusion/distributed/test_wan_vae_encoder_fastpath_kernels.pywan_vae_fastpath, thespatial-shard code, the distributed VAE executor or the diffusion registry. This
includes
tests/diffusion/models/wan2_2/andtest_wan_spatial_shard.py; its 2-ranktests need 2 cards.
AutoencoderKLWanvsOmniAutoencoderKLWan+install_wan_vae_fastpath(level=...), same weights and latent;server:
X-Stage-Durationsresponse header. The arms are--vae-fast-path losslessvsoff.Test Result
1. Tests (XPU, commit
0a14ba414): the 53 affected test files: 3034 passed, 24 skipped, 0 failed.This includes the fast-path suite above,
test_pruned_real_tiny_distributed_wan_vae_retains_numerics,and the new
test_registry_auto_enables_vae_tiling_only_for_tile_mode._LOSSLESS_DEFERS_CONV_BIAS); the newtest_xpu_lossless_keeps_conv_bias_in_the_convcovers the XPU behaviour instead;losslessnever takes.2. VAE decode, one B70, Wan2.2-T2V-A14B VAE, 720p/81f bf16:
lossless--level off2b. All three
--vae-fast-pathlevels, one B70, Wan2.2-T2V-A14B VAE, 720p/81f bf16(commit
103ddff8c, best of 2 after a warmup):off(upstream)losslesschannels_last3. End to end, Wan2.2-T2V-A14B, 720p/81f, 20 steps, CFG 2 × TP 2 on 4 B70s, tiled VAE:
--vae-fast-path off--vae-fast-path losslessWan22Pipeline.vae.decodeWan22Pipeline.diffuseWan22Pipeline.text_encoder.forwardX-Inference-Time-S)X-Peak-Memory-MB)d85ab93d2ba09c9c7570ef6921b3a6b8d85ab93d2ba09c9c7570ef6921b3a6b8