[MaciejBalaNV]

This PR optimizes Wan 2.2 VAE encode, mirroring #7056 which optimized decode. There are 2 new modes added: lossless and channels_last. lossless is bit-identical to the baseline. channels_last can have minimal differences. lossless achieves 1.76x speedup over the baseline while channels_last is 2.39x faster than the baseline for GB200 GPU.

Purpose

This PR adds VAE encoder optimizations for Cosmos3’s residual Wan2.2 architecture, building on the decoder optimizations in #7056.

The lossless mode implements the following optimizations:

• RMSNorm epilogue: ATen vector_norm retains the fp32 reduction; divide, scale, gamma and + 0.0 are fused while preserving ATen’s intermediate rounding. SiLU is folded in only for dtypes that pass the installation-time bitwise exactness check.
• Causal-convolution input assembly: cache cloning, temporal concatenation, padding and cache refresh are fused. Spatial padding is delegated to cuDNN after a per-shape bitwise check; rejected shapes use explicitly padded inputs. Zero-padding operations are skipped when all padding sizes are zero.
• Normalization-to-cache fusion: eligible RMSNorm+SiLU operations write directly into the temporally assembled convolution input and the next feature-cache entry, avoiding a separate normalized activation tensor. Already-normalized history is copied unchanged.
• Residual bias fusion: conv2 and conv_shortcut biases are folded into the residual-add kernel while preserving ATen’s rounding points. Unnecessary shortcut clones are removed when the modules cannot mutate their inputs.
• Downsampling input assembly: spatial padding and frame-layout conversion are fused into one kernel. Temporal downsampling reuses fused concatenation and cache refresh while preserving the first-chunk behavior. Shortcut averaging retains the reference ATen reduction.
• Chunked input and output assembly: untiled encoding patchifies one temporal chunk at a time and writes encoder features into a preallocated output tensor. The reference schedule—one initial frame, then four frames per chunk—is preserved, and quant_conv runs once after assembly.

The channels_last mode additionally converts convolution weights to channels-last memory formats to reduce runtime transposes. It uses single-pass RMSNorm+SiLU with the preceding conv1 bias folded in, and fuses AvgDown3D shortcut averaging with residual addition.

These paths can change floating-point reduction order, so outputs are validated within tolerance rather than bitwise.

The encoder is configured independently from the decoder through --vae-encode-fast-path, accepting off, lossless and channels_last. lossless is the default for supported encoders. Unsupported architectures retain reference encoding.

The change supports VAE tiling and tile parallelism and can coexist with spatially sharded decoding. Encoder spatial parallelism continues to use tiling.

A dedicated encoder benchmark compares all three modes, reporting startup and warmed latency, throughput, peak memory, posterior accuracy and optional reconstruction quality. It enforces bitwise equality for lossless and numerical quality gates for channels_last.

Test Plan

Unit tests

This PR adds 55 new test functions which check for bit-identical outputs and correct shapes of the new Triton kernels.

Serving tests

We test the outputs of Cosmos3-Nano model following this notebook. All of the baseline and lossless outputs are bit-identical.

vLLM Version:
0.30
vLLM-Omni Commit:
ToT

Test Result

Unit tests

All unit tests pass.

Serving tests

The lossless and off modes are bit-identical in generated videos. The channels_last examples will be added soon.

Performance tests

Gathered with this command: python benchmarks/diffusion/bench_wan_vae_encode.py --model nvidia/Cosmos3-Nano --frames 93 --check-reconstruction --json encode-video.json
GB200:

--vae-fast-path	Encode time	Speedup	Output vs. off
off	0.972 s	1.00x	reference
lossless	0.551 s	1.76x	bit-identical
channels_last	0.406 s	2.39x	PSNR 39.77 dB, max abs diff 0.0156

H100:

--vae-fast-path	Encode time	Speedup	Output vs. off
off	2.7278 s	1.00x	reference
lossless	2.0839 s	1.31x	bit-identical
channels_last	1.7810 s	1.53x	PSNR 40.34 dB, max abs diff 0.0156

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[MaciejBalaNV]

I also fixed additional edge-case where the Wan VAE encoder in channels_last mode would break the VAE decoder in lossless mode, since the latter would blindly accept the input as a non-channels last format.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (InferMatrix Direct completion gate rejected: invalid or wrong-head finding recheck: 44bc455346f17674; missing finding recheck: 44bc455346f17674).

[vllm-omni-review-bot]

Omni ReviewBot review

PR description

This PR adds an independent Wan 2.2 / Cosmos3 residual VAE encoder fast path beside the existing decoder path, with CLI/config --vae-encode-fast-path (off / lossless / channels_last, default lossless). On supported encoders it installs fused Triton assembly/norm/downsample forwards and an optional chunked encode_frames path; unsupported shapes or architectures keep the Diffusers reference encode. User-visible effect: faster encode by default when the residual patchified Wan encoder is present, with lossless aiming for bit-identical posteriors and channels_last trading a small numeric delta for more speed; a large encode benchmark and unit tests accompany the change.

Change flow

flowchart TD
  A["[EXISTING] Cosmos3 / Wan encode callers"]:::existing --> B["[CHANGED] --vae-encode-fast-path<br/>config + CLI + registry install"]:::changed
  B --> C["[NEW] Encoder fastpath kernels<br/>and encoder_forwards"]:::new
  C --> D["[CHANGED] OmniAutoencoderKLWan._encode<br/>encode_frames when eligible"]:::changed
  D --> E["[EXISTING] Reference encode / tiling fallback"]:::existing
  B --> F["[NEW] bench_wan_vae_encode.py<br/>mode A/B + quality gates"]:::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

CI at e2193b2b4cde (2026-10-01T20:02:18.876595+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Findings

• [P2] Non-test diff exceeds the 1000-line budget — benchmarks/diffusion/bench_wan_vae_encode.py:3
  Supplied diff_size is 2270 non-test added lines against a 1000-line budget (19 counted files, 10 exempt). The PR body explains kernels, the independent flag, and the benchmark, but still gives no split plan or concrete exemption (vendored sync, generated code, etc.). Largest contributors: benchmarks/diffusion/bench_wan_vae_encode.py (969), wan_vae_fastpath/forwards.py (285), triton_norm_cache.py (260), triton_downsample.py (205), install.py (164). Please split the benchmark/docs follow-up from the install+kernel change, or record an explicit exemption for the largest file.
  Existing thread: #8067 (comment)

---

_{🤖 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!}

[hsliuustc0106]

Reviewed head e2193b2b4cde077f072daf31067f28bfde1cc2d3 against base c0331048d0e28cb66a18a278ef93e68d4ab1c878.

No additional high-confidence correctness findings from this static pass. I traced CLI/config propagation, independent encoder/decoder installation and rollback, temporal chunk/cache ownership, tiling/slicing fallbacks, the new kernels, and benchmark timing/quality reporting. The current head adds random-operand probes for the earlier degenerate-input/weight verdict cases, runtime checks for global and module hooks, and decoder layout regressions for channels-last encoder outputs.

All 28 changed Python files pass static syntax parsing and the diff passes whitespace checks; reported CI is green. This was a static review of the fork snapshot: CPU tests, the earlier RTX 5090 reproductions, CUDA bitwise parity, distributed tiling, and GB200/H100 speedups were not rerun. The probes/regressions are present, but this review does not establish universal bitwise equivalence from a finite probe or independently verify the performance claims.

[MaciejBalaNV]

Thanks @0z5a and @hsliuustc0106, can we merge this PR in this case?

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (InferMatrix Direct completion gate rejected: invalid or wrong-head finding recheck: 44bc455346f17674; missing finding recheck: 44bc455346f17674).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (InferMatrix Direct completion gate rejected: invalid or wrong-head finding recheck: 44bc455346f17674; missing finding recheck: 44bc455346f17674).

[vllm-omni-review-bot]

Omni ReviewBot review

Changes since the previous review

• 1 new inline finding(s); 0 finding(s) below.
• Still needs attention: Non-test diff exceeds the 1000-line budget — benchmarks/diffusion/bench_wan_vae_encode.py

CI at 77f6f0266e60 (2026-10-07T20:04:34.243786+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Full review analysis

PR description

This PR adds an independent Wan 2.2 / Cosmos3 residual VAE encoder fast path (off / lossless / channels_last), wired through --vae-encode-fast-path and OmniDiffusionConfig.vae_encode_fast_path (default lossless), separate from the existing decoder flag. Install-time Triton/ATen fusions cover RMSNorm epilogues, causal-conv assembly, norm-to-cache writes, residual bias folding, downsampling assembly, and chunked encode output; channels_last further converts encoder conv weights and uses tolerance-checked math. Untiled encoding can use a preallocated chunk schedule while tiling still rides the patched module forwards, with a dedicated encode benchmark and expanded unit coverage.

Change flow

flowchart TD
  A["[EXISTING] serve CLI / OmniDiffusionConfig"]:::existing
  B["[CHANGED] vae_encode_fast_path projection<br/>omni_config / data / registry"]:::changed
  C["[NEW] install_wan_vae_encoder_fastpath<br/>+ encode_frames"]:::new
  D["[CHANGED] wan_vae_fastpath forwards<br/>+ encoder Triton kernels"]:::changed
  E["[CHANGED] OmniAutoencoderKLWan._encode"]:::changed
  F["[NEW] bench_wan_vae_encode + encoder tests"]:::new
  A --> B --> C --> D --> E
  C --> F
  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]

[P2] Non-test diff exceeds the 1000-line budget

Evidence and suggested fix

Supplied diff_size is 2270 non-test added lines against a 1000-line budget (19 counted files, 10 exempt). The PR body explains kernels, the independent flag, and the benchmark, but still gives no split plan or concrete exemption (vendored sync, generated code, etc.). Largest contributors: benchmarks/diffusion/bench_wan_vae_encode.py (969), wan_vae_fastpath/forwards.py (285), triton_norm_cache.py (260), triton_downsample.py (205), install.py (164). Please split the benchmark/docs follow-up from the install+kernel change, or record an explicit exemption for the largest file.

[princepride]

LGTM

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on a41627d750cc refers to an earlier head; the pull request now points at 8d3f1fe1f115.