[xuechendi]

1. What outcome do you want to achieve?

Lower latency for the multi-card Wan VAE encode in I2V/TI2V serving, without making the encoded latents less accurate than today's tile-parallel encode. It should stay opt-in, so default behavior doesn't change.

2. What is the highest-leverage variable for that outcome?

How much redundant work each rank does in the distributed encode. Today's tile split reuses diffusers' overlapping 256/192 tiles: at 1280×720 that's 28 tiles, so ~1.65× the frame's pixels are encoded. Those tiles are spread unevenly (3–4 per rank on 8 ranks), so the slowest rank sets the latency.

3. How will this work move that variable?

VLLM_OMNI_WAN_VAE_ENCODE_SPLIT=rank gives each rank exactly one latent-aligned piece on the squarest grid for the rank count, with a small context margin (default 2 latent px) that's cropped off afterwards.

• Less work: the encoded pixel count drops to ~1.0× the frame plus the margins, and every rank gets one piece.
• Measured on Intel Arc Pro B70, Wan2.1-style VAE (as used by I2V-A14B), 720p/81f:
  • 2 ranks: 8.58 → 5.89 s (1.46x).
  • 3 ranks: 5.83 → 4.08 s (1.43x).
  • 4 ranks: 4.46 → 3.16 s (1.41x).
  • Accuracy: error against a single-card un-tiled encode falls from rel_l2 9.5e-3 to 2.2–3.0e-3.

4. Whose customer or user needs does this address, and what are those needs?

Users serving Wan2.2 I2V/TI2V on multi-card Intel XPU nodes [e.g. name the team, customer or deployment, if you can share it publicly]. They need lower per-request latency at 720p, where the VAE encode is a fixed cost on every image-to-video request. They also need it without lower output quality, and without changing behavior for anyone who doesn't opt in. [Add the concrete need you have in mind: a latency target, a deployment, or the XPU roadmap item.]

Purpose

Faster and more accurate multi-card Wan VAE encode. It applies to --vae-patch-parallel-size > 1 with pipelines that encode video (I2V, TI2V,video-to-video); T2V does not encode and is unaffected.

	Before (tile, still the default)	After, opt-in VLLM_OMNI_WAN_VAE_ENCODE_SPLIT=rank
How the frame is split	diffusers' overlapping 256/192 sample-px tiles, blended at the seams: 28 tiles at 1280x720	one latent-aligned piece per rank on the squarest grid matching the rank count (e.g. 2x4 for 8 ranks at 90x160 latent)
Pixels encoded	~1.65x the frame (every overlap encoded twice)	~1.0x, plus VLLM_OMNI_WAN_VAE_ENCODE_MARGIN (default 2) latent px of context on inner edges
Load balance	28 tiles over 8 ranks = 3-4 tiles each	exactly one piece per rank; sizes within one latent row/column
Seams	blended overlaps	margin cropped off, pieces concatenated

Code layout.

• The split: env switches, grid choice, bounds, split and merge live in a new
  module, vllm_omni/diffusion/distributed/autoencoders/wan_encode_split.py.
• The hook: DistributedAutoencoderKLWan.tiled_encode only picks the operator.
• Isolation: this keeps the change independent of the companion PR
  xuechendi:xpu/wan-vae-perf, which edits other parts of autoencoder_kl_wan.py. A
  3-way merge of the two over main is conflict-free in either order.

Accuracy. Neither split is exact against a single-card un-tiled encode: a finite
margin does not cover the encoder's full receptive field, and tile blends. The rank
split with margin ≥ 1 is closer to the un-tiled encode than today's default (see
results).

Docs. A new "Distributed Encode Split (Wan)" section in
docs/user_guide/diffusion/parallelism/vae_parallelism.md documents both variables:

• invocation, and when the split applies;
• the defaults;
• the approximation limits (margin 0 is less accurate than tile);
• the measured table below.

The module docstring states the default margin and the margin-0 tradeoff, and links the
guide.

Test Plan

vLLM Version: v0.31.0 (vllm/vllm-openai-xpu:v0.31.0 base of docker/Dockerfile.xpu)

vLLM-Omni Commit: dc9cc80d8 on origin/main ebe2c1827

Hardware: Intel Arc Pro B70 (31.9 GiB), torch 2.14.0+xpu, triton-xpu 3.7.2,
diffusers 0.40.0.

1. Unit tests (CPU, no model):

   pytest tests/diffusion/distributed/test_wan_encode_split.py

   They cover:
   • the grid choice (squarest wins, ties to fewer rows, impossible grids skipped);
   • even_bounds (contiguous, covering, sizes within one);
   • env validation;
   • an exact split → merge round trip for 2/4/8 ranks × margins 0/1/2, with a stand-in
     encoder that is local per 8x8 block.
2. Distributed encode, Wan2.2-T2V-A14B VAE (the same Wan2.1-style VAE as I2V-A14B),
   DistributedAutoencoderKLWan:
   • --vae-parallel-mode spatial_shard_height (the I2V server's setting) on 2, 3 and 4
     ranks;
   • 720p, 81 frames, bf16, random input;
   • tile vs rank with margins 0/1/2/4;
   • each against a single-card un-tiled encode of the same input;
   • slowest rank, best of 3 after a warmup, default xccl gather.

Test Result

1. Unit tests (commit dc9cc80d8): 21 passed.

2. Distributed encode, 720p/81f, against a single-card un-tiled encode:

2 ranks:

split	time (slowest rank)	vs tile	rel_l2 (all frames)	rel_l2 (frame 0)	max_abs
tile (default)	8.584 s	—	9.54e-3	1.61e-2	0.500
rank margin 0	5.429 s	1.58x	2.37e-2	2.57e-2	1.906
rank margin 1	5.757 s	1.49x	4.13e-3	1.22e-2	0.422
rank margin 2 (default)	5.887 s	1.46x	2.23e-3	6.97e-3	0.234
rank margin 4	5.915 s	1.45x	1.60e-3	2.63e-3	0.047

3 ranks:

split	time (slowest rank)	vs tile	rel_l2 (all frames)	rel_l2 (frame 0)	max_abs
tile (default)	5.827 s	—	9.54e-3	1.61e-2	0.500
rank margin 0	3.879 s	1.50x	3.36e-2	3.75e-2	1.906
rank margin 1	3.964 s	1.47x	5.68e-3	1.74e-2	0.422
rank margin 2 (default)	4.078 s	1.43x	2.80e-3	9.88e-3	0.227
rank margin 4	4.336 s	1.34x	2.00e-3	3.12e-3	0.047

4 ranks:

split	time (slowest rank)	vs tile	rel_l2 (all frames)	rel_l2 (frame 0)	max_abs
tile (default)	4.464 s	—	9.54e-3	1.61e-2	0.500
rank margin 0	2.760 s	1.62x	4.18e-2	4.43e-2	3.203
rank margin 1	3.011 s	1.48x	5.88e-3	1.98e-2	0.469
rank margin 2 (default)	3.157 s	1.41x	3.03e-3	1.01e-2	0.281
rank margin 4	3.254 s	1.37x	2.09e-3	3.44e-3	0.047

• Margin 0 is a poor choice: it is the fastest but also the least accurate, worse
  than tile.
• Margin 2 (the default) is 3–4x more accurate than tile and 1.41–1.46x faster on
  2–4 ranks.
• The error grows slightly with the rank count: more pieces means more inner edges.
• The encode numbers were measured at 572e50fe7. dc9cc80d8 adds only documentation
  and a docstring on top of it, rebased onto the latest main.

Pre-commit (local, files in this PR): all hooks pass, including markdownlint.

Not run / not measured:

• Full I2V end to end: not measured. I2V-A14B transformer weights were not on the
  test node; the encoder measured here is the same Wan2.1-style VAE I2V-A14B uses.
• 8 ranks: not measured (the A/B used up to 4 cards).
• CUDA: not run. The change is device-agnostic Python in the encode split/merge.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-15T05:22:39.464882Z	572e50f	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[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/diffusion/index.md, docs/design/module/diffusion/parallelism.md.

Module owners: @princepride @david6666666 @Isotr0py

Routing: @princepride 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; @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.

[xuechendi]

I have provided the details in PR description. Changes is identified when comparing to lightX2V on vae encode time

[xuechendi]

Author self-review

• Scope: an opt-in split for the distributed Wan VAE encode, VLLM_OMNI_WAN_VAE_ENCODE_SPLIT=rank: one latent-aligned piece per rank, with a VLLM_OMNI_WAN_VAE_ENCODE_MARGIN context margin (default 2) that's cropped off. The default stays tile, so behavior is unchanged unless the variable is set. Decode and T2V are unaffected.
• Isolation: the logic lives in a new module, wan_encode_split.py. DistributedAutoencoderKLWan.tiled_encode only picks the operator, a 7-line change. A 3-way merge with PR #A over main is conflict-free in either order.
• Correctness: test_wan_encode_split.py, 21 tests on CPU, covers:
  • the grid choice: the squarest grid wins, ties go to fewer rows, impossible grids are skipped;
  • even_bounds: contiguous, covering, sizes within one;
  • env validation;
  • an exact split → merge round trip for 2/4/8 ranks × margins 0/1/2, using a block-local stand-in encoder, which checks every crop and offset.
• Accuracy and speed: measured on Intel Arc Pro B70, Wan2.1-style VAE (as used by I2V-A14B), 720p/81f, against a single-card un-tiled encode. With the default margin of 2, rank vs tile:

ranks	speedup vs tile	rel_l2 (rank vs tile)
2	1.46x	2.2e-3 vs 9.5e-3
3	1.43x	2.8e-3 vs 9.5e-3
4	1.41x	3.0e-3 vs 9.5e-3

Margin 0 is faster but less accurate than tile, which is why the default is 2.

• Pre-commit: all hooks pass locally.
• Not measured, listed in the description:
  • full I2V end to end: the I2V-A14B transformer weights weren't available, so this covers the encode stage only;
  • 8 ranks;
  • CUDA. The change is device-agnostic Python, but it hasn't been run there.

[hsliuustc0106]

[P2] Document the new encode controls and quality tradeoff

Both environment variables introduce user-facing behavior, but this PR updates no user documentation. The module docstring also omits the default margin and the reported accuracy degradation at margin zero. Please update the VAE parallelism guide with an invocation, activation requirements, defaults, and approximation limits.

[xuechendi]

Thanks — addressed in dc9cc80.

• Added a "Distributed Encode Split (Wan)" section to docs/user_guide/diffusion/parallelism/vae_parallelism.md. It covers:
  • invocation: both env vars, plus a serving example;
  • when the split applies: a Wan VAE in a pipeline that encodes, vae_patch_parallel_size > 1 in tile or spatial_shard_* mode, and an input larger than one tile;
  • defaults: tile, margin 2;
  • approximation limits: neither split is exact, and margin 0 is less accurate than tile;
  • the measured speed/accuracy table for 2–4 ranks.
• The module docstring now states the default margin and the margin-0 tradeoff, and links the guide.

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@xuechendi this PR reads as a performance or value claim:

• claim: - Margin 2 (the default) is 3–4x more accurate than tile and 1.41–1.46x faster on

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[hsliuustc0106]

Integration note for head dc9cc80d89b0a755cba36b839deed0a7c06a57c4 together with #8968 (0a14ba414):

Please add --vae-use-tiling to the new rank-encode serving example, and update the “automatic for these modes” requirement when rebasing/merging the two changes.

#8968 deliberately stops auto-enabling tiling for spatial_shard_*. The example selects spatial_shard_height without enabling tiling, so with both PRs applied the parent VAE _encode takes the local untiled path and never calls this PR's tiled_encode operator selection. VLLM_OMNI_WAN_VAE_ENCODE_SPLIT=rank then has no effect even on a large input. The doc additions can merge without a textual conflict while leaving this behavioral mismatch.

This is a cross-PR integration correction, not a defect in the split/crop/merge math. Source review only; the reported B70 encode gains and full I2V output quality were not independently reproduced.

[xuechendi]

@hsliuustc0106 Thanks, good catch — agreed, with #8968 applied the example would take the local untiled encode and the split would be a no-op.

Fixed in 4d21ae4 (docs/user_guide/diffusion/parallelism/vae_parallelism.md, docs only):

• the rank-encode serving example now passes --vae-use-tiling;
• the "When it applies" requirement no longer says tiling is automatic for these modes; it says to pass --vae-use-tiling / vae_use_tiling=True explicitly, because spatial_shard_* doesn't guarantee tiling and without it the split has no effect.

This holds on current main too (where tiling is still auto-enabled), so it doesn't depend on the merge order of the two PRs. markdownlint / pre-commit pass locally.

[xuechendi]

@vllm-omni-review-bot Answers from the B70 run logs:

1. Bottleneck. In tile mode, the distributed Wan VAE encode does redundant work. At 1280x720 the frame splits into 28 overlapping 256/192 tiles (4 rows x 7 columns), about 1.65x the frame's pixels, because every overlap is encoded twice. The tiles also land unevenly on some rank counts: 28 over 3 ranks is 10/9/9, and over 8 ranks it's 3–4 each. The figure comes from the tile geometry, not from a profiler trace. rank gives each rank one piece, which brings the encoded pixels down to about 1.0x the frame plus the margin. The measured gain (1.41–1.46x) stays under that 1.65x ceiling, as it should.

2. Value. Lower VAE-encode latency, a fixed cost on every I2V/TI2V request with --vae-patch-parallel-size > 1. Workload: 1280x720, 81 frames, bf16. Slowest rank, tile → rank with margin 2 (the default):

   • 2 ranks: 8.58 → 5.89 s
   • 3 ranks: 5.83 → 4.08 s
   • 4 ranks: 4.46 → 3.16 s

   It is also more accurate. Against a single-card un-tiled encode, rel_l2 drops from 9.5e-3 to 2.2–3.0e-3. This covers the encode stage only; end-to-end I2V was not measured (see the PR description).

3. A/B and ablation. The comparison uses the same code, process and input; only VLLM_OMNI_WAN_VAE_ENCODE_SPLIT changes. tile is the default and is what main does, so it serves as the base. The margin sweep (0/1/2/4) shows the speed/accuracy tradeoff. Margin 0 is the fastest (1.50–1.62x) but less accurate than tile (rel_l2 2.4–4.2e-2), so the default is 2. The PR makes one change, so there is no stacked delta.

Evidence

• Code: encode numbers measured at 572e50fe7. The current head adds only docs on top (dc9cc80d8, 4d21ae4ca).
• Setup: Intel Arc Pro B70, torch 2.14.0+xpu, diffusers 0.40.0. Wan2.2-T2V-A14B VAE (the same Wan2.1-style VAE as I2V-A14B), spatial_shard_height, xccl gather. Random input, 1280x720, 81 frames.
• Timing: 1 warm-up, then best of 3. The time is the slowest rank (an all-reduce max).
• Spread: not recorded. The benchmark keeps only the minimum of the 3 runs, so I can't give a mean ± std or per-rank times from this run. The 4-rank raw output, for example:

  tile         4.464 s   rel_l2 9.535e-03 (frame0 1.611e-02)  max_abs 5.000e-01
  rank m=0     2.760 s   rel_l2 4.176e-02 (frame0 4.434e-02)  max_abs 3.203e+00
  rank m=1     3.011 s   rel_l2 5.882e-03 (frame0 1.979e-02)  max_abs 4.688e-01
  rank m=2     3.157 s   rel_l2 3.030e-03 (frame0 1.008e-02)  max_abs 2.812e-01
  rank m=4     3.254 s   rel_l2 2.086e-03 (frame0 3.444e-03)  max_abs 4.688e-02
  speedup vs tile: rank m=0 1.62x, rank m=1 1.48x, rank m=2 1.41x, rank m=4 1.37x
  

• Quality signal: latent rel_l2 / max_abs against a single-card un-tiled encode (full tables in the description). Also, 21 CPU unit tests pass, including an exact split → merge round trip for 2/4/8 ranks at margins 0/1/2.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on claude-code (mimo-v2.6-pro) under experiment fleet-strict-claudecode-mimo26pro-100-c10-20261015.

[vllm-omni-review-bot]

Omni ReviewBot review status

Review for commit 4d21ae4ca67b5f1374740360fd9830f34ad53e10 ended as failed (step 'fetch' (pr.fetch_diff): cannot pin PR #8966 to 4d21ae4: git fetch failed: fatal: pack has 106 unresolved deltas
fatal: fetch-pack: invalid index-pack output; no fallback: the routing table has no fallback arm).

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 4d21ae4ca67b refers to an earlier head; the pull request now points at 6b7b15c274ca.

[xuechendi]

Rebased onto current main (7cafa13cc) to resolve the merge conflict. New head: 81a1a27cc.

• The only conflict was the reviewed-snapshot counts in tests/config/test_environment_variables.py. main added one platform-external variable (29 → 30), and this PR adds the two Wan encode-split variables as model-specific/promote (95 → 97, 67 → 69). No code changes otherwise.
• Tested on Intel Arc Pro B70 (2 cards) in the XPU docker image: tests/config/test_environment_variables.py, tests/diffusion/distributed/ Wan VAE tests (incl. test_wan_encode_split.py) and tests/diffusion/models/wan2_2/. 687 passed, 4 failed. The 4 failures (test_wan_spatial_shard.py decode reference ×3, test_wan_vae_fastpath_install.py::test_omni_diffusion_config_validates_vae_fast_path) fail the same way on main in that container. They need Wan2.1-T2V-1.3B-Diffusers configs from the HF hub, which aren't cached offline there, so they are unrelated to this PR.

[hsliuustc0106]

Looks good overall: opt-in via VLLM_OMNI_WAN_VAE_ENCODE_SPLIT=rank, so default tile behavior stays unchanged on CUDA/NPU/XPU; split/merge + docs/tests are solid, and this is the only one of the recent Wan batch with fully green CI including AMD/Intel.

Please rebase onto current main (currently dirty). Nit: margin=0 is worse than tile for accuracy — consider requiring margin >= 1 or warning at 0.

Happy to approve after the rebase lands and CI stays green.