[scestola-h]

Purpose

Extend the shared Wan VAE fast path from #7056 with a lossless Ascend path for
causal Conv3D spatial padding.

Diffusers materializes temporal and spatial padding with F.pad before every
Wan causal Conv3D. On Ascend, keeping temporal/cache assembly explicit but
passing the height/width padding to F.conv3d avoids that spatial padded copy.
The first call for each module/input-shape/dtype/layout runs both formulations
and compares them bitwise. A mismatch returns the reference result and caches a
fallback verdict for that shape.

The installer intentionally binds only WanCausalConv3d on NPU. CUDA and XPU
retain their complete existing fast paths unchanged, channels_last remains
CUDA/XPU-only, compiled or training calls use the Diffusers implementation,
and spatial-shard decode remains excluded.

This adapts the Helios-local prototype by Diego Didona (cuoredicervo,
ef47d374e) to the existing shared fastpath instead of adding a model-specific
monkeypatch.

Test Plan

Hardware: eight Ascend A3 dies (physical IDs 8–15), one die per run. Runtime:
PyTorch 2.10.0 + torch-npu 2.10.0.post4, Diffusers 0.40.0, npu-smi 25.2.3.

Focused A3 coverage:

• a tiny Wan VAE verifies that NPU installation patches only
  WanCausalConv3d;
• 12 BF16/FP32 causal-Conv3D cases cover cache/no-cache and representative
  channel/spatial shapes;
• every case checks both the first-call verification path and cached-verdict
  path for exact equality.

Real-weight A/B command (run independently on dies 8–15):

python benchmarks/diffusion/bench_wan_vae_decode.py \
  --model /home/data/weights/Helios-Distilled \
  --size 640x384 --frames 33 --dtype fp32 \
  --fast-path off,lossless --warmup 1 --iters 5 --seed 1101

Each reported die time is the minimum of five synchronized measurements after
one warmup. The summary reports mean ± population standard deviation and range
across the eight dies.

Local checks: complete changed-file pre-commit suite (Ruff, mypy 3.10, test
markers, SPDX, forbidden imports, torch.cuda policy, Buildkite schema, typos)
plus Markdown lint.

The exact ready-pipeline CPU selector requested in review is
pytest -sv tests/diffusion -m 'core_model and cpu'. The added registry test is
selected by it; the local host does not have the vLLM-Omni test environment, so
the result will come from the PR's Buildkite run rather than being claimed as a
local pass.

vLLM Version: 0.29.0 runtime package with the PR's self-contained Wan VAE
fastpath overlaid; the checkout targets the current vLLM 0.31 dependency.

vLLM-Omni validation Base / Commit: 156ddb3f4 / 0080a6fa8f4cf95b4ae74d98dd165331b286e4da

Current PR head: 26185ff1b, rebased onto 3a936aee8 on 2026-10-30.
The latest conflict was documentation-only: upstream #9424 lets tolerance-based
CUDA/XPU channels_last convolutions skip the bitwise spatial-padding probe to
reduce peak memory. The resolution retains that behavior while keeping
verify_spatial_padding=True for Ascend lossless. The rebase's changed-file
hooks, mypy, Buildkite validation, Python compile, diff checks, and pinned
Markdown lint all pass.

Test Result

Focused A3 validation passed: the installer patched 16
WanCausalConv3d modules and no other module type; all 12 numerical cases were
bitwise exact on the first and cached calls.

Helios VAE decode	off	lossless	Change
Time, mean ± std	1977.8 ± 1.3 ms	1835.7 ± 0.9 ms	1.077x faster / -7.19%
Per-die range	1976.1–1979.4 ms	1834.1–1836.9 ms	1.076–1.079x
Peak allocated memory	5.85 GiB	5.91 GiB	+0.06 GiB
Output on all 8 dies	reference	torch_equal=True, bit-equal	max abs diff 0

The original model-local prototype also measured 12.099 s → 11.883 s (1.8%)
for a 264-frame Helios end-to-end run with byte-identical output. That result is
supporting context only; the table above is the validation of the shared
fastpath implementation in this PR.

The eight-die table was collected before rebasing over the encoder fast-path
work in #8067. On the rebased head, a fresh die-15 A/B measured 1977.0 ms →
1835.1 ms (1.077x) with torch_equal=True, bit-equal output, and maximum
absolute difference 0. The current-head focused run also passed all 12 numerical
cases and confirmed that only WanCausalConv3d modules were patched.

AI assistance: Used AI coding assistants to help adapt the existing Wan VAE
fast-path architecture, draft tests and documentation, and prepare the PR text.
I reviewed the changes and validated them on Ascend A3.

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

Module owners: @princepride @david6666666 @Isotr0py

Routing: @princepride via module of the changed files, CODEOWNERS; @david6666666 via module of the changed files; @Isotr0py via module of the changed files

@scestola-h, 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: three questions on the performance claim

@scestola-h this PR reads as a performance or value claim:

• title: [Hardware][Ascend] Optimize Wan VAE spatial padding
• claim: | Time, mean ± std | 1977.8 ± 1.3 ms | 1835.7 ± 0.9 ms | 1.077x faster / -7.19% |

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.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 120s (try 2 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 600s (try 3 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

2 actionable finding(s).

CI at 72f91143cee1 (2026-10-10T16:04:28.873339+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/cursor/cursor-grok-4.6-high could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

The change adds an Ascend branch to the shared Wan VAE lossless fast path. When decoder parameters are already on NPU, the installer rebinds only WanCausalConv3d: temporal cache assembly stays an explicit concat and pad, and height/width padding is passed into F.conv3d so the spatial padded activation is not materialized. The first call for each module, input shape, dtype, and layout runs that formulation against the pre-padded convolution, keeps the result only when the two are bitwise equal, and otherwise caches a fallback to the Diffusers forward. CUDA still installs the existing fused decoder path, and the A3 ready step plus the decode benchmark now cover this NPU path.

Change flow

flowchart TD
  A["[EXISTING] Wan VAE decode"]:::existing --> B["[CHANGED] Registry allows Ascend"]:::changed
  B --> C["[CHANGED] Installer binds causal Conv3D"]:::changed
  C --> D["[CHANGED] Spatial pad inside F.conv3d"]:::changed
  D --> E["[EXISTING] Bitwise verdict or Diffusers fallback"]:::existing
  D --> F["[NEW] A3 causal-conv exactness tests"]:::new
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

See inline comments below.

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

[P1] CPU diffusion guard for the registry change is not in the reported runs

Evidence and suggested fix

The modified registry test is selected by the ready-pipeline step "Simple · Diffusion Test" in .buildkite/cuda/test-ready.yml, whose command is pytest -sv tests/diffusion -m 'core_model and cpu' (no --run-level). test_wan_vae_fastpath_install.py is marked core_model and cpu, so that step is the guard for this change. The PR Test Plan and Test Result report the new A3 causal-conv cases and the changed-file pre-commit hooks (Ruff, mypy, markers, SPDX, import and torch.cuda policy, Buildkite schema, typos). They do not report this CPU selector, and they do not say the CPU diffusion suite was left unrun. Please run that command, or state the gap explicitly.

[scestola-h]

The registry coverage is marked core_model + cpu, so it is selected by pytest -sv tests/diffusion -m 'core_model and cpu'. I attempted that exact selector locally, but the available container has a stale vLLM dependency and collection stops before selection (MultiModalProcessingResult is absent). I have documented that limitation in the PR body rather than claiming a pass; the ready-pipeline Buildkite result is still required before merge.

[vllm-omni-review-bot]

[P2] Ascend-only binding depends on parameter device, not the NPU platform

Evidence and suggested fix

The registry turns the fast path on when current_omni_platform.is_npu() is true, but install_wan_vae_fastpath sets npu_spatial_padding_only only when the first decoder parameter's device type is already npu. A CPU-resident VAE on an NPU platform therefore receives the full CUDA replacement map (WanDecoder3d, residual blocks, resample, upsample, and RMSNorm), and channels_last is not refused. Those replacements are not a no-op on NPU: rms_norm_fastpath still returns its restructured vector_norm path when the Triton kernels decline, and a channels_last install also converts convolution weights. The new CPU test calls the registry hook with is_npu=True on a default CPU VAE and only asserts is_installed, so it accepts that full bind. The A3 installer test does check the causal-conv-only set, but only after .to("npu"). Pipeline construction is not in this sparse checkout, so it is not proven that serving always installs after the VAE is on NPU. Gate the limited bind and the channels_last refusal on the NPU platform as well as parameter device, and assert report.patched in the CPU registry test for a module that is not on NPU.

[scestola-h]

Fixed in 0080a6fa8. The registry now passes an explicit NPU-only mode based on current_omni_platform.is_npu(), so installation is safe even while the VAE parameters are still on CPU. In that mode it binds only WanCausalConv3d, never installs the CUDA-oriented encoder/RMSNorm/resample rewrites, and rejects channels_last. The CPU registry test now asserts the exact patched set is {WanCausalConv3d} and that the encoder installer is not called.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] CPU diffusion guard for the registry change is not in the reported runs — tests/diffusion/distributed/test_wan_vae_fastpath_install.py:636

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[scestola-h]

Performance evidence after rebasing onto current main (156ddb3f4), head 0080a6fa8:

• Bottleneck: Diffusers materializes the H/W padding tensor before each Wan causal Conv3D. This change leaves temporal/cache assembly explicit but passes spatial padding to F.conv3d, avoiding that copy. The measurements are synchronized VAE-stage timings, not a kernel-profiler trace.
• Isolated value: Helios-Distilled, 640×384, 33 frames, FP32, one A3 die: 1977.8 ± 1.3 ms → 1835.7 ± 0.9 ms across dies 8–15 (−7.19%, 1.077×), one warmup + five timed iterations per die. Peak allocation changed 5.85 → 5.91 GiB.
• A/B control: same weights, input, seed, process, and fast-path implementation; only --fast-path off vs lossless changes. Outputs were bit-exact on every die (max absolute difference 0). A fresh current-head run on die 15 reproduced 1977.0 → 1835.1 ms (1.077×) with bit-exact output.
• End-to-end context: the earlier model-local prototype measured 12.099 → 11.883 s (−1.8%) on a 264-frame Helios run with byte-identical output. I treat that only as supporting context, not the primary shared-fastpath result.

The focused current-head A3 run also passed all 12 BF16/FP32 cache/no-cache cases on both first-verification and cached-verdict calls.

[scestola-h]

Author self-review of 0080a6fa8:

• Rebased over [optimization] Added Wan 2.2 VAE encode optimizations #8067 and preserved its complete CUDA encoder/decoder fast paths.
• Found and fixed a platform-coupling issue: NPU selection no longer depends only on the VAE parameter device during registry setup. The registry explicitly requests the NPU subset, which can patch only WanCausalConv3d; NPU encoder/RMSNorm/resample rewrites and channels_last are excluded.
• Added a registry assertion for the exact NPU patched-module set and for no encoder installation.
• Re-ran the full changed-file checks, Markdown lint, focused A3 numerics, and real-weight Helios VAE A/B; all passed. The exact CPU diffusion selector cannot run in the available local container because its vLLM dependency is stale and collection fails before selection, so the ready-pipeline Buildkite result remains required.
• Residual risk is limited to CI integration/collection on the current dependency set; CUDA behavior is inherited from current main and was not performance-tested here.

AI coding assistants helped with the rebase review, tests, and wording; I reviewed all changes and ran the reported validation.

[psv666]

Ran this PR's own benchmarks/diffusion/bench_wan_vae_decode.py on 8×Ascend 910B1, single card (torch_npu 2.10.0.post4, CANN 9.1.0, vllm-omni @ 0080a6f), local Wan2.2-TI2V-5B-Diffusers VAE, bf16:

size	off	lossless	speedup
480x480, 33 frames	1038.1 ms	986.3 ms	1.05x
832x480, 81 frames (480p)	4315.2 ms	4055.4 ms	1.06x

installed=True, torch_equal=True, bits_equal=True, max_abs_diff=0, PSNR=inf at both sizes — the lossless claim holds exactly, including the built-in first-call bitwise self-check.

On this stack (CANN 9.1.0) the win is ~6% of VAE decode: the spatial F.pad copy it removes is a small share of the causal Conv3D pipeline at these shapes. Fine as a free and exact win; just flagging that users shouldn't expect this fastpath to move end-to-end Wan latency on 910B1 by much.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 10 days

@scestola-h this pull request has had no human commit, comment or review since 2026-10-16. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[scestola-h]

Rebased onto current main after #8968 added XPU support to the shared Wan VAE fast path. The conflict resolution preserves the complete CUDA/XPU decoder and encoder paths and keeps this PR limited to the lossless causal-Conv3D spatial-padding path on NPU. New head: 63f23797a, base: 13c1e8228. Changed-file hooks, mypy, Buildkite schema validation, Python compilation, diff checks, and pinned Markdown lint all pass locally; fresh GitHub CI is now running. This remains ready for maintainer review.

[psv666]

Status on head 63f23797a (base 13c1e8228): no new code findings; the Ascend binding scope is unchanged and correct.

• Ascend path (causal_conv_forward NPU branch, _conv_verdict_key, _conv_with_spatial_padding, the NPU install mode and the registry npu_spatial_padding_only=is_npu call) is byte-identical to 0080a6fa8, the tree I measured on 8x910B1 on 2026-10-16; #8968 only widened the registry's full-path branch to XPU, so that result still applies here.
• Merge state is CONFLICTING/DIRTY: the only textual conflict is the _conv_with_spatial_padding docstring in forwards.py (main's #9424 wording vs this PR), so the rebase is trivial, not semantic.
• When resolving, keep verify_spatial_padding=True for Ascend lossless; #9424 skips the first-call bitwise probe only under channels_last, which Ascend already rejects.
• One confirmation run after rebase: Wan2.2 TI2V VAE, bf16, --fast-path off,lossless, 480x480/33 and 832x480/81 frames, single 910B1 (the base moved through #9424 memory changes).

[scestola-h]

Thanks for the exact conflict analysis. Rebased onto current main (3a936aee8); new head is 26185ff1b. The sole conflict was the _conv_with_spatial_padding docstring. The resolution keeps #9424’s tolerance-based channels_last behavior and its generalized device-convolution wording, while Ascend still rejects channels_last and installs lossless with verify_spatial_padding=True, so its first-call bitwise probe is unchanged. Range-diff shows no production-code delta beyond that upstream integration. Changed-file hooks, mypy, Buildkite schema validation, Python compilation, diff checks, and pinned Markdown lint all pass; fresh CI is running.

For the requested post-rebase hardware confirmation: our available pool is Atlas A3 (Ascend910_9362), not 910B1, and dies 8–15 are currently occupied by a live eight-die Helios service, so I have not disrupted it. The cached Wan2.2-TI2V-5B checkpoint is present; I can run the corresponding A3 off,lossless confirmation once those dies are released. An exact 910B1 confirmation would need to run on the 910B1 environment used for the earlier independent measurement.

[psv666]

Rebase verified on 26185ff1b (base 3a936aee8).

• Patch-level compare against the head I previously reviewed (63f2379): the only content delta is the _conv_with_spatial_padding docstring resolution; every other hunk difference is a base line-number offset.
• verify_spatial_padding=True is retained for Ascend: NPU installs at level="lossless", so FastPathConfig.channels_last is False and the new channels_last and not verify_spatial_padding early return in forwards.py:358 is never taken. The bitwise probe still runs.
• NPU binding is still WanCausalConv3d only (35 modules patched in this run), channels_last is still refused on NPU (install.py:169), and CUDA/XPU keep _REPLACEMENT_FORWARDS unchanged.

Local re-run, one 910B1, ab_snap:v31 (torch 2.10.0+cpu, torch-npu 2.10.0.post4), Wan2.2-TI2V-5B VAE, bf16, warmup 1 / iters 5, arm order alternating per round:

size / frames	off ms	lossless ms	ratio	repeats
480x480 / 33	1042.9	990.3	1.053x	3
832x480 / 81	4311.6	4055.7	1.063x	2

torch_equal=True with max abs diff 0 at both sizes. Peak allocated on lossless is +0.02 GiB (4.045 -> 4.062 and 6.255 -> 6.275), the same sign as your +0.06 GiB.

Versus my 2026-10-16 run on 0080a6fa8 (1038.1 -> 986.3 ms, 1.052x): +4.8 ms off and +4.0 ms lossless, i.e. +0.4-0.5% in both arms with the ratio unchanged. A uniform shift in both arms is host-level drift, not an arm-specific regression, and it is within what I see between sessions on this shared machine. The new 832x480/81 point shows the benefit growing slightly with workload size (6.3% vs 5.3% wall time).

Remaining gaps are unchanged from my previous review and are not new blockers: Ascend exposes only the decoder lossless level (channels_last and the encoder fast path stay CUDA/XPU-only), spatial-shard decode stays excluded, the first call per (module, shape) still pays the bitwise probe, and this evidence is a single 910B1 die rather than the eight A3 dies in your table.