feat(diffusion): add QwenImage spatial-shard decode

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
feat(diffusion): select Wan VAE decode mode automatically
2026-08-11 10:49:10 +08:00 · 2026-08-11 09:41:05 +08:00 · 2026-08-11 05:10:50 +08:00 · 2026-08-11 05:08:31 +08:00 · 2026-08-10 22:20:37 +08:00 · 2026-08-10 21:55:01 +08:00
47 changed files with 2117 additions and 364 deletions
@@ -1,6 +1,11 @@
 # CODEOWNERS is used to automatically assign PR reviewers
 # Format: path pattern @username or @team
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 #
 # The vllm_omni/** source-tree rules below are derived from the `owners` and
 # `primary_code_paths` front matter in docs/design/module/*.md. When updating
 # ownership, update the design doc first, then regenerate the matching rule here.
 # Last-match-wins: more specific rules must come AFTER broader ones.
 # Default owners (fallback reviewers for all files)
 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao
-# Engine and model_executor
+# ComfyUI / external apps
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/apps/ @fhfuih @alex-jw-brooks
 /vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
 /vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/worker/ @tzhouam @gcanlin
-# Diffusion
+# Benchmarks (out-of-tree)
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
 # Diffusion CPU offloading
 /vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
 # Omni and TTS support
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
 /vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
 # Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
 /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
-# Usability and Observability
+# =====================================================================
-/vllm_omni/deploy/ @tzhouam @lishunyang12
+# Module design ownership (docs/design/module/*.md)
-/vllm_omni/metrics/ @lishunyang12
+# Each section header cites the source design doc.
-/vllm_omni/profiler/ @gcanlin
+# =====================================================================
-# ComfyUI
+# Autoregressive runtime — docs/design/module/ar_runtime.md
-/apps/ @fhfuih @alex-jw-brooks
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
 /vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
-# Large Scale Serving
+# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
 /vllm_omni/engine/ @tzhouam @fake0fan
 # File-level overrides: input/output modality contracts own these engine files
 # — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Entrypoints and serving boundaries — docs/design/module/entrypoints.md
 /vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
 # File-level overrides within entrypoints/
 /vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
 /vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
 # Error contracts — docs/design/module/error_contracts.md
 /vllm_omni/errors.py @alex-jw-brooks @NickCao
 # OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
-# Quantization and configuration
+# Model integration — docs/design/module/model_integration.md
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
 /vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
 /vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
 # Config override within model_executor — docs/design/module/vllm_omni_config.md
 /vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
 # Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
 /vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
 /vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
 # Configuration & deploy — docs/design/module/vllm_omni_config.md
 /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
 # Quantization — docs/design/module/quantization.md
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
-# Hardware plugins and accelerators
+# Observability — docs/design/module/observability.md
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+/vllm_omni/metrics/ @lishunyang12 @vraiti
 # Profiling — docs/design/module/profiling.md
 /vllm_omni/profiler/ @gcanlin
 # Benchmarking (in-tree) — docs/design/module/benchmarking.md
 /vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
 # LoRA (no dedicated design doc; ownership intentionally narrowed)
 /vllm_omni/lora/ @SamitHuang @wtomin
 # Transformers utils (related path only in design docs; preserves prior ownership)
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
 # =====================================================================
 # Diffusion module family — docs/design/module/diffusion/
 # Umbrella rule first; sub-module rules after it win via last-match.
 # =====================================================================
 # Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
 # — docs/design/module/diffusion/index.md
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 # Cache management — docs/design/module/cache_management.md
 /vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
 # Diffusion quantization — docs/design/module/quantization.md
 /vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
 # Diffusion profiling — docs/design/module/profiling.md
 /vllm_omni/diffusion/profiler/ @gcanlin
 # Diffusion runtime (scheduler, executor, worker)
 # — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
 /vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
 # Diffusion model integration (pipelines, registration, checkpoint loading)
 # — docs/design/module/diffusion/diffusion_model_integration.md
 /vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
 /vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 /vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 # Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
 # — docs/design/module/diffusion/parallelism.md
 /vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 /vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 # Diffusion offloader (CPU offloading, DLO)
 # — docs/design/module/diffusion/offloader.md
 /vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
@@ -66,7 +66,10 @@ repos:
      entry: python tools/pre_commit/check_tts_adapter.py
      language: python
      pass_filenames: false
-      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
      # The budget-is-now-slack reminder is printed on a PASSING run, and
      # pre-commit hides output from passing hooks unless the hook is verbose.
      verbose: true
    # Keep `suggestion` last
    - id: suggestion
@@ -449,20 +449,21 @@ Complete examples in the codebase:
 ---
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode (Wan and QwenImage)
-The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
+The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE and distributed QwenImage backend additionally support a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available.
 ### How it differs from tile parallel
-| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|-----------------------------------------------|
+|--------|--------------------------|----------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts |
-| Scope | Decode + encode | Decode only |
+| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks |
 | Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only |
-Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
+Spatial-shard decode swaps the decoder's spatial convolutions/padding for context-aware halo-exchanging variants so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Wan and QwenImage share the padding, 2D convolution, attention, halo, and trim primitives in `wan_spatial_shard.py`; `qwenimage_spatial_shard.py` supplies a causal-convolution subclass that remains discoverable by Diffusers' QwenImage cache accounting.
 ### Wiring
@@ -473,15 +474,17 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
  -> DiffusionParallelConfig.vae_parallel_mode
  -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
  -> DistributedVaeExecutor.parallel_mode
-  -> DistributedAutoencoderKLWan.tiled_decode dispatch
+  -> capability-based pipeline preparation
  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
 ```
-`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
+The shared spatial policy gates the path: it requires distributed decode to be enabled, a 5D latent, a supported VAE capability, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request.
 ### Notes
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
+- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
 - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
 - QwenImage additionally covers cache discovery, compile compatibility, mixed tile/spatial requests, both axes, and all-rank result broadcast in `tests/diffusion/distributed/test_qwenimage_spatial_shard.py`.
 ---
@@ -7,6 +7,7 @@ owners:
  - "@yinpeiqi"
  - "@fake0fan"
  - "@Sy0307"
  - "@Gaohan123"
 primary_code_paths:
  - vllm_omni/core/**
  - vllm_omni/worker/**
@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
  - "@alex-jw-brooks"
  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/benchmarks/**
  - vllm_omni/entrypoints/cli/benchmark/**
@@ -5,6 +5,7 @@ status: draft
 owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
 primary_code_paths:
  - vllm_omni/diffusion/cache/**
 related_code_paths:
@@ -3,9 +3,10 @@ title: Diffusion Model Integration
 kind: module
 status: draft
 owners:
-  - "@Isotr0py"
+  - "@fhfuih"
  - "@Bounty-hunter"
  - "@wtomin"
-  - "@ZJY0516"
+  - "@RuixiangMa"
 primary_code_paths:
  - vllm_omni/diffusion/models/**
  - vllm_omni/diffusion/model_loader/**
@@ -6,6 +6,7 @@ owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
  - "@fhfuih"
 primary_code_paths:
  - vllm_omni/diffusion/sched/**
  - vllm_omni/diffusion/executor/**
@@ -4,7 +4,8 @@ kind: module
 status: draft
 owners:
  - "@yuanheng-zhao"
-  - "@Isotr0py"
+  - "@david6666666"
  - "@lishunyang12"
 primary_code_paths:
  - vllm_omni/diffusion/offloader/**
 related_code_paths:
@@ -6,6 +6,7 @@ owners:
  - "@princepride"
  - "@yuanheng-zhao"
  - "@xuechendi"
  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/diffusion/distributed/**
  - vllm_omni/diffusion/attention/parallel/**
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
 owners:
  - "@alex-jw-brooks"
  - "@linyueqian"
  - "@NickCao"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -8,8 +8,15 @@ owners:
  - "@xuechendi"
 primary_code_paths:
  - vllm_omni/platforms/**
 related_code_paths:
  - vllm_omni/attention/**
 primary_code_path_owners:
  # Additional owners scoped to attention kernels (quantized/custom attention)
  - paths:
      - vllm_omni/attention/**
    owners:
      - "@david6666666"
      - "@lishunyang12"
 related_code_paths: []
 depends_on:
  - ar_runtime.md
  - diffusion/index.md
@@ -6,6 +6,8 @@ architecture_state: current
 owners:
  - "@Sy0307"
  - "@amy-why-3459"
  - "@Gaohan123"
  - "@alex-jw-brooks"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -9,6 +9,14 @@ primary_code_paths:
  - vllm_omni/model_executor/**
  - vllm_omni/model_extras/**
  - vllm_omni/plugins/**
 primary_code_path_owners:
  # Additional owners scoped to a sub-path of model_executor/**
  - paths:
      - vllm_omni/model_executor/models/**
    owners:
      - "@Sy0307"
      - "@amy-why-3459"
      - "@linyueqian"
 related_code_paths:
  - vllm_omni/transformers_utils/**
  - vllm_omni/tokenizers/**
@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
  - "@lishunyang12"
  - "@vraiti"
 primary_code_paths:
  - vllm_omni/metrics/**
 related_code_paths:
@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).
 ### Which geometries run sparse
-Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
-is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
-tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
-that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+promoted rows remain visible to every query and key.
 attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
 speedup, so pick an aligned geometry when you want one.
-Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
-spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
-patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
-are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+and kernel tiling aligned, at the cost of retaining more blocks.
 distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
 to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
-"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
-resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
-which you can buy back with `start_step`.
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
 that implements protected video tails for rf_v2.
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
@@ -15,7 +15,7 @@
 ## Overview
-VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan/QwenImage spatial-shard decode, which shards decoder feature maps along height or width.
 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
+| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). |
 Additional requirements:
@@ -126,9 +126,9 @@ Additional requirements:
 ---
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode (Wan and QwenImage)
-The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
+The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis.
 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
@@ -141,7 +141,7 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(
        tensor_parallel_size=2,
        vae_patch_parallel_size=2,               # must match the DiT group size
-        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
+        vae_parallel_mode="auto",                # default
    ),
 )
 ```
@@ -152,14 +152,18 @@ Or from the CLI / serving entrypoint:
 vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
    --tensor-parallel-size 2 \
    --vae-patch-parallel-size 2 \
-    --vae-parallel-mode spatial_shard_width
+    --vae-parallel-mode auto
 ```
 **Constraints and behavior:**
- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- Spatial-shard decode is **decode-only** and implemented for the **Wan** VAE and the distributed **QwenImage** backend used by Qwen-Image and Krea2. Pipelines using a different VAE backend keep their existing behavior.
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
+- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
 - Wan preserves its rank-0-only decode result, while QwenImage preserves its existing all-rank result by broadcasting after rank-0 assembly.
 - Spatial sharding trades activation memory for latency as resolution grows. Use `vae_parallel_mode="tile"` when bounded VAE memory is more important than decode speed.
 The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends.
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends
 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 ### Multi-NPU: 768P combined-service configuration
 Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
-degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
+degree 8, native tiled VAE patch parallelism degree 8, and distributed
 layerwise offload:
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
  --usp 8 \
  --ring 1 \
  --text-encoder-tp-size 8 \
-  --enable-layerwise-offload \
+  --enable-distributed-layerwise-offload \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --vae-patch-parallel-size 8 \
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
 partition restart is required. Layerwise offload applies to both DiTs.
 On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
 768P without offloading or sharding: use distributed layerwise offload (as
 above) or HSDP — see
 [§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
 ### Optional optimizations
 Two independent optimizations may be enabled on top of the configuration
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
 over the whole sequence, so ring parallelism would split away the keys it
 needs. Scale with `--usp` instead.
 ## Memory and attention optimizations (A3)
 ### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
 Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
 combined MiniMax-H3 service at 768P. Enable **one** of the following at
 server startup:
 **Distributed layerwise offload** — DiT layers are offloaded with parameters
 gathered across the parallel group instead of replicated per rank:
 ```bash
  --enable-distributed-layerwise-offload
 ```
 Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
 input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
 usage is high with this option.
 **HSDP** — hybrid sharded data parallelism for the DiT parameters. The
 multi-stream memory-reuse knob is mandatory with this flag:
 ```bash
 export MULTI_STREAM_MEMORY_REUSE=2
 ```
 ```bash
  --use-hsdp
 ```
 Host memory usage is small, but large shapes may still OOM; HBM usage
 optimizations for this path are ongoing.
 ### FLASH_ATTN backend with MindIE-SD
 Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
 [§ Environment](#environment)). On NPU this backend carries most of the
 memory-reduction and performance work: a mask-free packed varlen path that
 never materializes the quadratic `full_qk` padding mask, and K/V prefix
 slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
 transformer.
 ### Optional: LaserAttention fused kernel
 For an additional attention speedup, select the Ascend LaserAttention fused
 kernel before starting the server:
 ```bash
 export MINDIE_SD_FA_TYPE="ascend_laser_attention"
 ```
 This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
 exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
 kernel's fp16 workspace cannot overflow on outlier activations. Measured
 speedup numbers will be added here.
 ## HTTP API examples
 The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:
 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
 These measurements describe the validated shapes rather than a general
 throughput guarantee.
@@ -579,6 +579,7 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
    assert cfg.data_parallel_size == 3
    assert cfg.cfg_parallel_size == 3
    assert cfg.mask_sp_padding is True
    assert cfg.vae_parallel_mode == "auto"
    assert cfg.world_size == 72
@@ -9,6 +9,10 @@ This script tests two main scenarios:
 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
 """
 import sys
 from types import SimpleNamespace
 from unittest.mock import Mock
 import pytest
 import torch
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
        impl.forward_cuda(query, query, query, metadata)
 # ---------------------------------------------------------------------------
 # NPU mask-free packed attention (PR #5891).
 #
 # These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
 #   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
 #     sync, no MindIE-SD import);
 #   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
 #   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
 #
 # They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
 # into ``sys.modules`` (same pattern as ``test_paged_attention`` /
 # ``benchmarks/conftest``). No production code is exercised for real kernels.
 def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
    return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
 def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
    """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
    Restored automatically by monkeypatch. Defaults are no-op callables so the
    top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
    instance methods that would otherwise call them.
    """
    fake = SimpleNamespace(
        attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
        attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
    )
    monkeypatch.setitem(sys.modules, "mindiesd", fake)
    return fake
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 def test_resolve_packed_seq_accepts_real_plus_pad_contract():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
 def test_resolve_packed_seq_single_document_no_padding():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
@pytest.mark.parametrize(
    "case",
    [
        "causal",
        "missing_key",
        "batch_ne_1",
        "cu_too_long",
        "cu_qk_mismatch",
        "max_seqlen_not_int",
        "used_gt_total",
        "used_zero",
        "pad_longer_than_real",
    ],
 )
 def test_resolve_packed_seq_rejects_malformed_metadata(case):
    """Any metadata outside the exact [real, pad] contract returns None so the
    caller falls back to the masked path."""
    cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
    cu2 = torch.tensor([0, 8], dtype=torch.int32)
    cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
    cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
    base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
    q = torch.randn(1, 8, 2, 4)
    impl = _npu_impl(causal=(case == "causal"))
    if case == "missing_key":
        extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5}  # 3 of 4 keys absent
    elif case == "batch_ne_1":
        q = torch.randn(2, 8, 2, 4)
        extra = base
    elif case == "cu_too_long":
        extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
    elif case == "cu_qk_mismatch":
        extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
    elif case == "max_seqlen_not_int":
        extra = {**base, "max_seqlen_q": 5.0}
    elif case == "used_gt_total":
        extra = {**base, "max_seqlen_q": 9}  # 9 > total length 8
    elif case == "used_zero":
        extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
    elif case == "pad_longer_than_real":
        # real=3, pad=5: the real document must be the longer one.
        extra = {
            "cu_seqlens_q": cu_short,
            "cu_seqlens_k": cu_short,
            "max_seqlen_q": 3,
            "max_seqlen_k": 3,
        }
    else:
        extra = base
    assert impl._resolve_packed_seq_npu(q, q, extra) is None
 # --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
 def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
    """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
    Lets a test assert which branch ``forward_fa_npu`` dispatched to without
    running any real kernel.
    """
    impl = _npu_impl()
    sentinel = torch.tensor([1.0])
    impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
    impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
    return impl, sentinel
 def test_npu_env_unset_uses_varlen(monkeypatch):
    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_env_laser_uses_slice(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_prefix_kv_slice_npu.assert_called_once()
    impl._forward_varlen_packed_npu.assert_not_called()
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
 def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
    """Locks in the CURRENT dispatch behavior (PR #5891).
    Any value other than ``ascend_laser_attention`` — including the valid
    ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
    the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
    comment #1 on the PR and is intentionally captured here as a regression
    net: update this test when the dispatch starts honoring/rejecting these.
    """
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
    """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
    the existing masked attention_forward path; the new branches are skipped."""
    fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
    _fake_mindiesd(monkeypatch, attention_forward=fake_forward)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(
        attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
        extra={},  # no npu_attn_varlen opt-in
    )
    q = torch.randn(1, 8, 2, 4)
    out = impl.forward_fa_npu(q, q, q, metadata)
    impl._forward_varlen_packed_npu.assert_not_called()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
    fake_forward.assert_called_once()
    assert out is fake_forward.return_value
 # --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
 def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured["key"] = key
        captured["value"] = value
        captured.update(kwargs)
        return torch.ones_like(query)  # deterministic stand-in kernel output
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()  # softmax_scale == 0.5
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        "laser_input_scale": 256.0,
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    # K/V sliced to the valid prefix (used_k == 5); q keeps full length.
    assert captured["key"].shape == (1, 5, 2, 4)
    assert captured["value"].shape == (1, 5, 2, 4)
    assert captured["query"].shape == (1, 8, 2, 4)
    assert captured["attn_mask"] is None
    # Pre-divide q/k/v by the factor and compensate the kernel scale by its
    # square (exact for power-of-two); output scaled back by the factor.
    torch.testing.assert_close(captured["query"], q / 256.0)
    torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
    assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
 def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured.update(kwargs)
        return torch.ones_like(query)
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        # no laser_input_scale → no pre-scaling
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    torch.testing.assert_close(captured["query"], q)
    assert captured["scale"] == pytest.approx(0.5)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
 if __name__ == "__main__":
    print("Running FlashAttention Padding Tests...")
    print("=" * 60)
@@ -5,10 +5,9 @@
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
-mask from the video and prefix segments separately, so the two only agree on the
+to the always-kept prefix segment before generating the block mask.
 block count when the video segment ends on a boundary.
 """
 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
    assert plan.latent_shape == list(ALIGNED_GRID)
-def test_misaligned_video_segment_stays_dense():
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
-    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
-    # pad keys would take a share of every real query's softmax denominator.
+
-    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+    assert plan is not None
    assert plan.prefix_len == PREFIX_ROWS
    assert plan.used_len == PREFIX_ROWS + 62496
    assert plan.latent_shape == list(MISALIGNED_GRID)
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
 def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
    torch.manual_seed(1)
    prefix_len, grid = 384, (4, 16, 10)
    video_len = grid[0] * grid[1] * grid[2]
    used = prefix_len + video_len
    heads, head_dim = 4, 128
    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
    impl = make_impl()
    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
    out = impl._forward_sparse_npu(q, k, v, plan)
    reference = torch.nn.functional.scaled_dot_product_attention(
        q.transpose(1, 2),
        k.transpose(1, 2),
        v.transpose(1, 2),
        scale=impl.softmax_scale,
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
    class FakeBackend:
        supports_prefix_kv_slicing = False
        # Mirror AttentionBackend.supports_packed_mask_free: H3's packed
        # attention path calls attn_backend.supports_packed_mask_free() to decide
        # whether to build the padding mask. Return False so the fake exercises
        # the masked path (the abstract backend default).
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    def __init__(self, **kwargs):
        del kwargs
        super().__init__()
@@ -0,0 +1,233 @@
 import os
 from dataclasses import dataclass
 import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import get_open_port
 from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import (
    DistributedAutoencoderKLQwenImage,
 )
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
    prepare_pipeline_spatial_shard_decode,
 )
 from vllm_omni.diffusion.offloader.module_collector import PipelineModules
 from vllm_omni.platforms import current_omni_platform
@dataclass
 class _FakeExecutor:
    group: object
    parallel_size: int
    parallel_mode: str
 def _tiny_qwenimage_vae() -> DistributedAutoencoderKLQwenImage:
    return DistributedAutoencoderKLQwenImage(
        base_dim=8,
        z_dim=4,
        dim_mult=[1, 2, 2, 2],
        num_res_blocks=1,
        attn_scales=[],
        temperal_downsample=[False, True, True],
        input_channels=3,
        latents_mean=[0.0] * 4,
        latents_std=[1.0] * 4,
    )
@pytest.mark.core_model
@pytest.mark.cpu
 def test_qwenimage_auto_spatial_shard_gate_is_shape_and_topology_aware(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.get_world_size",
        lambda group=None: 4,
    )
    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
    vae.use_tiling = True
    vae.distributed_executor = _FakeExecutor(group=object(), parallel_size=4, parallel_mode="auto")
    vae.is_distributed_enabled = lambda: True
    landscape = torch.zeros((1, 16, 1, 60, 104))
    portrait = torch.zeros((1, 16, 1, 104, 60))
    assert vae._spatial_shard_decode_split_dim(landscape, 4) == "width"
    assert vae._spatial_shard_decode_split_dim(portrait, 4) == "height"
    assert vae._spatial_shard_decode_enabled(landscape) is True
    vae.distributed_executor.parallel_size = 2
    assert vae._spatial_shard_decode_enabled(landscape) is False
@pytest.mark.core_model
@pytest.mark.cpu
 def test_pipeline_dispatches_qwenimage_spatial_preparation(monkeypatch: pytest.MonkeyPatch):
    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
    prepared = []
    vae._prepare_spatial_shard_decode = lambda: prepared.append(vae)
    monkeypatch.setattr(
        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
        lambda pipeline: PipelineModules(
            dits=[],
            dit_names=[],
            encoders=[],
            encoder_names=[],
            vaes=[torch.nn.Identity(), vae],
        ),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.is_initialized",
        lambda: True,
    )
    prepare_pipeline_spatial_shard_decode(torch.nn.Identity())
    assert prepared == [vae]
@pytest.mark.core_model
@pytest.mark.cpu
@torch.no_grad()
 def test_qwenimage_installed_decoder_preserves_cache_discovery_state_and_direct_path():
    torch.manual_seed(0)
    vae = _tiny_qwenimage_vae().eval()
    latents = torch.randn(1, 4, 2, 4, 5)
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents, return_dict=False)[0]
    causal_conv_count = sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules())
    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="height")
    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="width")
    vae.clear_cache()
    assert causal_conv_count > 0
    assert len(vae._feat_map) == causal_conv_count
    assert sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules()) == causal_conv_count
    assert any(
        isinstance(module, qwenimage_spatial_shard.QwenImageDistCausalConv3d) for module in vae.decoder.modules()
    )
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
@pytest.mark.core_model
@pytest.mark.cpu
@torch.no_grad()
 def test_qwenimage_installed_decoder_direct_path_can_be_compiled():
    torch.manual_seed(1)
    vae = _tiny_qwenimage_vae().eval()
    latents = torch.randn(1, 4, 1, 3, 5)
    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object())
    eager = vae.decode(latents, return_dict=False)[0]
    compiled_decode = torch.compile(vae.decode, backend="eager", fullgraph=False)
    compiled = compiled_decode(latents, return_dict=False)[0]
    torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
 def _qwenimage_spatial_shard_worker(rank: int, master_port: str, results) -> None:
    from vllm_omni.diffusion.distributed.parallel_state import (
        destroy_model_parallel,
        init_distributed_environment,
        initialize_model_parallel,
    )
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = master_port
    device = current_omni_platform.get_torch_device(rank)
    current_omni_platform.set_device(device)
    backend = current_omni_platform.dist_backend
    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
    try:
        torch.manual_seed(2)
        vae = _tiny_qwenimage_vae().to(device=device, dtype=torch.float32).eval()
        vae.init_distributed()
        generator = torch.Generator(device=device).manual_seed(3)
        landscape = torch.randn((1, 4, 2, 5, 7), generator=generator, device=device)
        portrait = torch.randn((1, 4, 2, 7, 5), generator=generator, device=device)
        with torch.inference_mode():
            vae.use_tiling = False
            vae.set_parallel_size(1, mode="tile")
            landscape_reference = vae.decode(landscape, return_dict=False)[0]
            portrait_reference = vae.decode(portrait, return_dict=False)[0]
            vae.enable_tiling(
                tile_sample_min_height=32,
                tile_sample_min_width=32,
                tile_sample_stride_height=24,
                tile_sample_stride_width=24,
            )
            vae.set_parallel_size(2, mode="tile")
            tile_before_install = vae.decode(landscape, return_dict=False)[0]
            vae.set_parallel_size(2, mode="auto")
            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
            landscape_auto = vae.decode(landscape, return_dict=False)[0]
            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
            portrait_auto = vae.decode(portrait).sample
            vae.set_parallel_size(2, mode="spatial_shard_height")
            landscape_height = vae.decode(landscape, return_dict=False)[0]
            vae.set_parallel_size(2, mode="spatial_shard_width")
            portrait_width = vae.decode(portrait, return_dict=False)[0]
            vae.set_parallel_size(2, mode="tile")
            tile_after_install = vae.decode(landscape, return_dict=False)[0]
        assert tile_after_install.shape == tile_before_install.shape
        assert torch.equal(tile_after_install, tile_before_install)
        assert all(
            buffer is None
            for module in vae.decoder.modules()
            for name, buffer in module._buffers.items()
            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
        )
        diffs = {
            "landscape_auto": (landscape_auto - landscape_reference).abs().max().item(),
            "portrait_auto": (portrait_auto - portrait_reference).abs().max().item(),
            "landscape_height": (landscape_height - landscape_reference).abs().max().item(),
            "portrait_width": (portrait_width - portrait_reference).abs().max().item(),
        }
        results[rank] = {
            "diffs": diffs,
            "landscape_shape": tuple(landscape_auto.shape),
            "portrait_shape": tuple(portrait_auto.shape),
        }
    finally:
        destroy_model_parallel()
        if dist.is_initialized():
            dist.destroy_process_group()
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "L4"}, num_cards=2)
 def test_qwenimage_spatial_shard_matches_reference_broadcasts_and_switches_paths():
    manager = mp.get_context("spawn").Manager()
    results = manager.dict()
    mp.spawn(
        _qwenimage_spatial_shard_worker,
        args=(str(get_open_port()), results),
        nprocs=2,
        join=True,
    )
    assert set(results) == {0, 1}
    assert results[0]["landscape_shape"] == results[1]["landscape_shape"] == (1, 3, 5, 40, 56)
    assert results[0]["portrait_shape"] == results[1]["portrait_shape"] == (1, 3, 5, 56, 40)
    for rank_results in results.values():
        assert max(rank_results["diffs"].values()) <= 2e-5
@@ -1,4 +1,5 @@
 import os
 from contextlib import contextmanager
 from types import SimpleNamespace
 import pytest
@@ -10,6 +11,8 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
    DistributedAutoencoderKLWan,
    prepare_pipeline_wan_spatial_shard_decode,
    select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform
@@ -17,6 +20,23 @@ from vllm_omni.platforms import current_omni_platform
 # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
@contextmanager
 def _active_spatial_shard(split_dim: str):
    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
        wan_spatial_shard.SpatialShardContext(
            input_extent=2,
            local_input_extent=2,
            split_dim=split_dim,
            rank=0,
            world_size=3,
        )
    )
    try:
        yield
    finally:
        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
@@ -255,11 +275,13 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
    with _active_spatial_shard("height"):
        mid = mid_rank_pad(x)
    assert mid.shape == (1, 1, 2, 3)
    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
    with _active_spatial_shard("height"):
        last = last_rank_pad(x)
    assert last.shape == (1, 1, 3, 3)
@@ -270,12 +292,14 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
    x = torch.ones((1, 1, 2, 2))
    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
-    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
    with _active_spatial_shard("width"):
        mid = mid_rank_pad(x)
    assert mid.shape == (1, 1, 2, 2)
    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
-    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
    with _active_spatial_shard("width"):
        last = last_rank_pad(x)
    assert last.shape == (1, 1, 2, 3)
@@ -299,6 +323,113 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
    assert vae._spatial_shard_decode_enabled(z) is False
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
    "z_shape,world_size,expected",
    [
        ((1, 16, 1, 60, 104), 4, "width"),
        ((1, 16, 5, 60, 104), 4, "width"),
        ((1, 16, 5, 104, 60), 4, "height"),
        ((1, 16, 5, 2, 2), 4, None),
        ((1, 16, 5, 64, 64), 4, "width"),
    ],
 )
 def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
    assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
@pytest.mark.core_model
@pytest.mark.cpu
 def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
        lambda group=None: 4,
    )
    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
    vae.use_tiling = True
    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
    vae.is_distributed_enabled = lambda: True
    large = torch.zeros((1, 16, 5, 60, 104))
    unshardable = torch.zeros((1, 16, 1, 2, 2))
    assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
    assert vae._spatial_shard_decode_enabled(large) is True
    assert vae._spatial_shard_decode_enabled(unshardable) is False
    vae.distributed_executor.parallel_size = 2
    assert vae._spatial_shard_decode_enabled(large) is False
@pytest.mark.core_model
@pytest.mark.cpu
 def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
    object.__setattr__(
        vae,
        "distributed_executor",
        SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"),
    )
    installed = []
    monkeypatch.setattr(
        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
        lambda pipeline: SimpleNamespace(vaes=[vae]),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized",
        lambda: True,
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
        lambda group=None: 4,
    )
    monkeypatch.setattr(
        wan_spatial_shard,
        "install_wan_spatial_shard_decode",
        lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
    )
    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
    assert installed == []
    vae.distributed_executor.parallel_size = 4
    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
    assert installed == [(vae, vae.distributed_executor.group, "height")]
@pytest.mark.core_model
@pytest.mark.cpu
@torch.no_grad()
 def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
    vae = DistributedAutoencoderKLWan(
        base_dim=8,
        decoder_base_dim=8,
        z_dim=4,
        dim_mult=[1, 2],
        num_res_blocks=1,
        attn_scales=[],
        temperal_downsample=[False],
        latents_mean=[0.0] * 4,
        latents_std=[1.0] * 4,
        scale_factor_temporal=1,
        scale_factor_spatial=2,
    ).eval()
    latents = torch.randn(1, 4, 1, 4, 5)
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents, return_dict=False)[0]
    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
    assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -444,3 +575,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
    assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
        f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
    )
 def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
    from vllm_omni.diffusion.distributed.parallel_state import (
        destroy_model_parallel,
        init_distributed_environment,
        initialize_model_parallel,
    )
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = master_port
    device = current_omni_platform.get_torch_device(rank)
    current_omni_platform.set_device(device)
    backend = current_omni_platform.dist_backend
    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
    try:
        vae = DistributedAutoencoderKLWan.from_pretrained(
            _SPATIAL_SHARD_MODEL,
            subfolder=_SPATIAL_SHARD_SUBFOLDER,
            torch_dtype=torch.float32,
        ).to(device=device, dtype=torch.float32)
        vae.eval()
        generator = torch.Generator(device=device).manual_seed(1)
        small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
        landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
        portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
        with torch.inference_mode():
            vae.use_tiling = False
            vae.set_parallel_size(1, mode="tile")
            small_reference = vae.decode(small, return_dict=False)[0]
            vae.use_tiling = True
            vae.set_parallel_size(2, mode="tile")
            tile_before_install = vae.decode(small, return_dict=False)[0]
            vae.set_parallel_size(2, mode="auto")
            assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
            small_sharded = vae.decode(small, return_dict=False)[0]
            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
            vae.decode(landscape, return_dict=False)
            vae.set_parallel_size(2, mode="tile")
            tile_after_install = vae.decode(small, return_dict=False)[0]
            vae.use_tiling = False
            vae.set_parallel_size(1, mode="tile")
            portrait_reference = vae.decode(portrait, return_dict=False)[0]
            vae.use_tiling = True
            vae.set_parallel_size(2, mode="auto")
            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
            portrait_sharded = vae.decode(portrait, return_dict=False)[0]
        assert all(
            buffer is None
            for module in vae.decoder.modules()
            for name, buffer in module._buffers.items()
            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
        )
        if rank == 0:
            small_diff = (small_sharded.float() - small_reference.float()).abs()
            portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
            return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
            return_dict["small_max_abs_diff"] = small_diff.max().item()
            return_dict["small_mean_abs_diff"] = small_diff.mean().item()
            return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
            return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
    finally:
        destroy_model_parallel()
        if dist.is_initialized():
            dist.destroy_process_group()
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_auto_spatial_shard_switches_shape_paths_safely():
    manager = mp.get_context("spawn").Manager()
    return_dict = manager.dict()
    mp.spawn(
        _auto_spatial_shard_decode_worker,
        args=(return_dict, "29502"),
        nprocs=2,
        join=True,
    )
    assert return_dict.get("tile_fallback_exact") is True
    assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
    assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
    assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
    assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
@@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    class FakeBackend:
        supports_prefix_kv_slicing = True
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
@@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    assert attention.attention.metadata.extra["valid_kv_length"] == 5
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
    )
    class FakeBackend:
        supports_prefix_kv_slicing = False
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return True
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
        def __init__(self):
            super().__init__()
            self.metadata = None
        def forward(self, query, key, value, metadata):
            self.metadata = metadata
            return query
    attention = object.__new__(MiniMaxH3Attention)
    torch.nn.Module.__init__(attention)
    attention.attention = FakeAttention()
    q = torch.randn(8, 2, 4)
    attention._run_packed_attention(
        q,
        q,
        q,
        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
        max_seqlen=5,
        packed_total=8,
    )
    assert attention.attention.metadata.attn_mask is None
    assert attention.attention.metadata.extra["valid_kv_length"] == 5
    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
 def test_packed_attention_keeps_padding_mask_for_other_backends():
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
    class FakeBackend:
        supports_prefix_kv_slicing = False
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
@@ -120,3 +120,34 @@ def test_tp_accepts_checkpoint_supported_sizes():
    arch = MiniMaxH3DiTArchConfig()
    for tp_size in (1, 2, 4, 7):
        model._validate_tp_config(arch=arch, tp_size=tp_size)
@pytest.mark.parametrize(("parallel_size", "enabled"), [(1, False), (2, True)])
 def test_video_vae_auto_mode_preserves_native_tile_parallelism(monkeypatch, parallel_size, enabled):
    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
    parallel_state = {}
    video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
    nn.Module.__init__(video_vae)
    video_vae.remote = nn.Identity()
    video_vae.model = type("FakeModel", (), {"parallel_tiling": False})()
    monkeypatch.setattr(vae_module, "get_dit_group", lambda: object())
    monkeypatch.setattr(vae_module.dist, "get_world_size", lambda group: 2)
    monkeypatch.setattr(vae_module.dist, "get_rank", lambda group: 0)
    monkeypatch.setattr(
        vae_module.importlib,
        "import_module",
        lambda package: type(
            "FakeParallelModule",
            (),
            {"get_parallel_state": staticmethod(lambda: parallel_state)},
        )(),
    )
    video_vae.set_parallel_size(parallel_size, mode="auto")
    assert video_vae.parallel_size == parallel_size
    assert video_vae.model.parallel_tiling is enabled
    assert parallel_state["group_size"] == parallel_size
    assert parallel_state["sp_enabled"] is enabled
@@ -728,6 +728,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.core_model
@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    memory_pool_active = False
    preparation_pool_states = []
    @contextmanager
    def _memory_pool_context(tag):
        nonlocal memory_pool_active
        assert tag == "weights"
        memory_pool_active = True
        try:
            yield
        finally:
            memory_pool_active = False
    class _DummyLoader:
        def __init__(self, load_config, od_config=None):
            del load_config, od_config
@@ -775,16 +788,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
    monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
    monkeypatch.setattr(
        model_runner_module,
        "prepare_pipeline_spatial_shard_decode",
        lambda pipeline: preparation_pool_states.append(memory_pool_active),
    )
    monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
    monkeypatch.setattr(
        model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
    )
-    DiffusionModelRunner.load_model(runner)
+    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
    assert runner.cache_backend is None
    assert runner.od_config.cache_backend is None
    assert dummy_cache_backend.enabled is False
    assert preparation_pool_states == [True]
@pytest.mark.core_model
@@ -76,6 +76,18 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
    assert args.cfg_parallel_size == 4
 def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None:
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    OmniServeCommand().subparser_init(subparsers)
    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
    tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"])
    assert default_args.vae_parallel_mode == "auto"
    assert tile_args.vae_parallel_mode == "tile"
 def _make_headless_args(**kwargs) -> TrackingNamespace:
    defaults = {
        "model": "fake-model",
@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
    return found
-def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
+def _check(
    label: str,
    path: Path,
    actual: int,
    budget: int,
    sample: list[str],
    budget_name: str,
    notices: list[str],
 ) -> list[str]:
    if actual > budget:
        listing = "\n".join(f"      {line}" for line in sample)
        return [
            f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
        ]
    if actual < budget:
-        return [
+        # Progress is never a build failure. Removing branches is the goal, so a
-            f"{path}: {actual} {label}, down from a budget of {budget}. "
+        # PR that does it must not hand-edit this file to stay green -- nor break
-            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
+        # `main` for everyone else when it merges from a stale base.
-            f"so the ratchet holds the new ground.",
+        notices.append(
-        ]
+            f"{path}: {actual} {label}, below the budget of {budget}. "
            f"Lower {budget_name} to {actual} in "
            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
        )
    return []
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
            print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
            return 1
    notices: list[str] = []
    branches = _model_type_branches(ast.parse(serving_path.read_text()))
    detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
        MAX_MODEL_TYPE_BRANCHES,
        [f"line {lineno}: {src}" for lineno, src in branches],
        "MAX_MODEL_TYPE_BRANCHES",
        notices,
    )
    errors += _check(
        "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
        MAX_LEGACY_DETECTORS,
        [f"line {lineno}: {name}" for lineno, name in detectors],
        "MAX_LEGACY_DETECTORS",
        notices,
    )
    for notice in notices:
        print(f"check_tts_adapter: {notice}", file=sys.stderr)
    for error in errors:
        print(f"check_tts_adapter: {error}", file=sys.stderr)
    return 1 if errors else 0
@@ -452,7 +452,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
    cfg_parallel_size: int = Field(default=1, ge=1)
    vae_patch_parallel_size: int = Field(default=1, ge=1)
    text_encoder_tp_size: int = Field(default=1, ge=1)
-    vae_parallel_mode: str = "tile"
+    vae_parallel_mode: str = "auto"
    use_hsdp: bool = False
    mask_sp_padding: bool = False
    hsdp_shard_size: int = -1
@@ -466,9 +466,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
            raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
        if self.ulysses_mode not in {"strict", "advanced_uaa"}:
            raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
            raise ValueError(
-                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of "
                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                f"but got {self.vae_parallel_mode!r}."
            )
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
    # avoid a slower masked-attention plan when tail padding is not semantic.
    supports_prefix_kv_slicing: bool = False
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        """Whether packed attention never reads attn_mask on this platform.
        When True, models that pack a [real, pad] two-document layout and
        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
        constructing the padding mask entirely. Backends whose mask-free
        behavior is platform-dependent must check current_omni_platform.
        """
        return False
    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
    # Platform resolution rejects an explicit selection outside this set, so a
    # hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
    #     the packed cu_seqlens tensors.
    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
    #     backend that advertises supports_prefix_kv_slicing.
    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
    #     Requires the [real, pad] two-document packing contract; see
    #     FlashAttentionImpl._forward_varlen_packed_npu.
    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
    #     with this set (>1), q/k/v are divided by the factor before the op,
    #     the kernel scale_value is multiplied by its square, and the output
    #     is scaled back (exact for power-of-two factors). Absent means no
    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
    # Piecewise attention metadata (mixed causal/full masks).
    # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from functools import partial
 import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
    piecewise_attn,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supports_piecewise_spans: bool = True
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
        # before ever reading attn_mask. NPU: forward_fa_npu honors the
        # npu_attn_varlen opt-in and rebuilds the padding mask itself if the
        # packed contract fails. XPU reads attn_mask, so models must keep
        # constructing it there.
        return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
    @classmethod
    def supports_attention_mask(cls) -> bool:
        return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
        # model marks extra["npu_attn_varlen"] and carries packed metadata, so
        # the padding document is excluded without reading or materializing
        # the quadratic full_qk mask.
        #   - default: packed varlen via mindiesd attention_forward_varlen;
        #   - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
        #     slicing via mindiesd attention_forward (so the env-selected op
        #     type is honored; the TND varlen op cannot be).
        extra = attn_metadata.extra if attn_metadata else {}
        if extra.get("npu_attn_varlen", False):
            if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
            else:
                out = self._forward_varlen_packed_npu(query, key, value, extra)
            if out is not None:
                return out
        attention_mask = attn_metadata.attn_mask if attn_metadata else None
        if attention_mask is None and extra.get("npu_attn_varlen", False):
            # Models skip mask construction when this opt-in is set
            # (FlashAttentionBackend.supports_packed_mask_free). The packed
            # paths above declined (contract mismatch), so rebuild the padding
            # mask here to keep the masked fallback correct.
            used = extra.get("valid_kv_length")
            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
                raise ValueError(
                    "npu_attn_varlen packed metadata is unusable and no attn_mask "
                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
                    "refusing to run unmasked attention over padding rows."
                )
            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
        # NPU aclnnFlashAttentionScore requires mask shape to be one of:
        # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
            op_type="fused_attn_score",
            layout=layout,
        )
    def _resolve_packed_seq_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        extra: dict,
    ) -> tuple[list[int], list[int]] | None:
        """Resolve packed document boundaries (cumulative end offsets per doc).
        Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
        or None unless the metadata matches the exact contract MiniMax-H3
        produces:
          - q/k/v are [1, T, N, D] (BSND, single packed batch);
          - cu_seqlens describe a [real, pad] two-document packing, with the
            padding document as a strict suffix;
          - max_seqlen_q/k are Python ints equal to the real document length
            (so boundaries are derived without any device sync).
        """
        if self.causal:
            return None
        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
        if any(k not in extra for k in packed_keys):
            return None
        cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
        used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
        if query.shape[0] != 1 or key.shape[0] != 1:
            return None
        if not isinstance(used_q, int) or not isinstance(used_k, int):
            return None
        total_q, total_k = query.shape[1], key.shape[1]
        # .shape is host-side metadata: counting documents never syncs.
        if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
            return None
        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
            return None
        if used_q == total_q and used_k == total_k:
            # No padding document: one full-length document.
            return [total_q], [total_k]
        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
            # [real, pad] packing: the real document must be the longer one
            # (consistent with the max_seqlen naming).
            return [used_q, total_q], [used_k, total_k]
        return None
    def _forward_varlen_packed_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
        Returns None (caller falls back to the mask path) when the packed
        contract does not hold; see _resolve_packed_seq_npu.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        seq_q, seq_k = resolved
        try:
            from mindiesd import attention_forward_varlen
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        q = query.squeeze(0)  # [T, N, D] == TND
        k = key.squeeze(0)
        v = value.squeeze(0)
        # attention_forward_varlen wants cu_seqlens as host lists of cumulative
        # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
        out = attention_forward_varlen(
            q,
            k,
            v,
            [0, *seq_q],
            [0, *seq_k],
            softmax_scale=self.softmax_scale,
        )
        return out.unsqueeze(0)
    def _forward_prefix_kv_slice_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Mask-free attention by slicing K/V to the valid prefix (zero-copy
        views), then mindiesd attention_forward without a mask.
        Same numerical contract as the varlen path: the padding document is a
        strict suffix, so dropping it from K/V is identical to masking it out.
        Query keeps full length; outputs on padding rows are never consumed
        downstream. Returns None when the packed contract does not hold.
        Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
        cannot honor the env-selected op type, while attention_forward can.
        The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
        bf16 activations with large outliers overflow 65504 into ±inf and the
        affected rows turn NaN. Models may opt into exact power-of-two input
        pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
        are divided by the factor, the kernel scale_value is multiplied by its
        square, and the output is scaled back. Absent or invalid factor means
        no pre-scaling.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        _, seq_k = resolved
        used_k = seq_k[0]  # real document length (first cumulative end)
        try:
            from mindiesd import attention_forward
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # mindiesd always takes BSND input regardless of `layout`; the arg
        # selects the op-internal layout and laser only supports BNSD, so do
        # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
        layout = "BNSD"
        key = key[:, :used_k]
        value = value[:, :used_k]
        input_scale = extra.get("laser_input_scale")
        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
        if preserve_input_range:
            query = query / input_scale
            key = key / input_scale
            value = value / input_scale
        scale = self.softmax_scale
        if preserve_input_range:
            scale *= input_scale**2
        def _postprocess(out: torch.Tensor) -> torch.Tensor:
            return out * input_scale if preserve_input_range else out
        return _postprocess(
            attention_forward(
                query,
                key,
                value,
                attn_mask=None,
                opt_mode="manual",
                op_type="fused_attn_score",  # MINDIE_SD_FA_TYPE env overrides this
                layout=layout,
                scale=scale,
            )
        )
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
    other case — warmup denoise steps, exempt layers, a layer that does not declare
    ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection or not aligned to the kernel's
+    segments too short to pay for block selection — delegates to FlashAttention,
-    block size — delegates to FlashAttention, so a model can select this backend
+    so a model can select this backend unconditionally. MindIE-SD handles an
-    unconditionally.
+    irregular video tail internally, retaining it outside the sparse blocks.
    """
    def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
                _MIN_VIDEO_BLOCKS,
            )
            return None
        if video_len % _BLOCK_SIZE != 0:
            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
            # builds its block mask from the two segments pooled apart. The counts
            # only agree when the video segment ends on a block boundary; otherwise
            # a prefix block is forced on out of range, and the block straddling the
            # seam holds real prefix rows that selection may drop. Padding the
            # sequence to realign it is not an option because rf_v2 ignores
            # attn_mask, so pad keys would take a share of every softmax denominator.
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
                "multiple of %d to enable sparsity.",
                tuple(latent_shape),
                video_len,
                _BLOCK_SIZE,
                _BLOCK_SIZE,
            )
            return None
        logger.info_once(
            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
@@ -206,10 +206,13 @@ class DiffusionParallelConfig:
    process group.
    """
-    vae_parallel_mode: str = "tile"
+    vae_parallel_mode: str = "auto"
    """VAE parallel decode strategy.
-    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
+    - "auto": Select the longer spatial-shard axis for validated Wan and
      distributed QwenImage tiled decodes on a full DiT-group topology;
      otherwise use patch/tile parallel decode.
    - "tile": Patch/tile parallel decode. Each rank decodes a subset
      of spatial tiles and the results are stitched on rank 0.
    - "spatial_shard_height": Spatially-sharded decode that splits decoder
      feature maps along height and exchanges halo rows around spatial
@@ -250,8 +253,9 @@ class DiffusionParallelConfig:
        assert self.allgather_degree > 0, "AllGather degree must be > 0"
        assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
        assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
-        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
+        assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+            "vae_parallel_mode must be one of "
            "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
            f"but got {self.vae_parallel_mode!r}."
        )
        if self.allgather_degree > 1:
@@ -4,16 +4,23 @@
 from typing import Any
 import torch
 import torch.distributed as dist
 from diffusers.models.autoencoders import AutoencoderKLQwenImage
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
    DistributedOperator,
    DistributedVaeMixin,
    GridSpec,
    TileTask,
 )
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
    prepare_spatial_shard_decode,
    resolve_spatial_shard_split_dim,
    spatial_shard_decode_enabled,
 )
 logger = init_logger(__name__)
@@ -104,10 +111,39 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
        ]
        return dec
    def _spatial_shard_decode_split_dim(
        self,
        z: torch.Tensor | None = None,
        world_size: int | None = None,
    ) -> str | None:
        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
        return spatial_shard_decode_enabled(self, z, model_name="QwenImage")
    def _prepare_spatial_shard_decode(self) -> None:
        prepare_spatial_shard_decode(
            self,
            install=qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode,
        )
    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)
        if self._spatial_shard_decode_enabled(z):
            world_size = dist.get_world_size(group=self.distributed_executor.group)
            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
            assert split_dim is not None
            logger.debug("Decode running with QwenImage VAE spatial_shard_%s mode", split_dim)
            return qwenimage_spatial_shard.spatial_shard_decode(
                self,
                z,
                group=self.distributed_executor.group,
                return_dict=return_dict,
                split_dim=split_dim,
            )
        logger.debug("Decode running with distributed executor")
        result = self.distributed_executor.execute(
            z,
@@ -18,11 +18,25 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
    GridSpec,
    TileTask,
 )
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
    prepare_pipeline_spatial_shard_decode,
    prepare_spatial_shard_decode,
    resolve_spatial_shard_split_dim,
    spatial_shard_decode_enabled,
 )
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
    select_auto_spatial_shard_split_dim as select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
    """Backward-compatible alias for the model-agnostic pipeline dispatcher."""
    prepare_pipeline_spatial_shard_decode(pipeline)
 class OmniAutoencoderKLWan(AutoencoderKLWan):
    def _execution_context(self):
        try:
@@ -270,45 +284,30 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
        dec = torch.clamp(dec, min=-1.0, max=1.0)
        return dec
-    def _spatial_shard_decode_split_dim(self) -> str | None:
+    def _spatial_shard_decode_split_dim(
-        mode = self.distributed_executor.parallel_mode
+        self,
-        if mode == "spatial_shard_width":
+        z: torch.Tensor | None = None,
-            return "width"
+        world_size: int | None = None,
-        if mode == "spatial_shard_height":
+    ) -> str | None:
-            return "height"
+        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
        return None
    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        split_dim = self._spatial_shard_decode_split_dim()
+        return spatial_shard_decode_enabled(self, z, model_name="Wan")
        if split_dim is None:
            return False
        if z.ndim != 5:
            logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
            return False
        if not self.is_distributed_enabled():
            return False
-        group = self.distributed_executor.group
+    def _prepare_spatial_shard_decode(self) -> None:
-        world_size = dist.get_world_size(group=group)
+        prepare_spatial_shard_decode(
-        requested_size = int(self.distributed_executor.parallel_size)
+            self,
-        if requested_size != world_size:
+            install=wan_spatial_shard.install_wan_spatial_shard_decode,
            logger.warning(
                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
                "to match the DIT group size; falling back to tiled decode. "
                "requested=%s dit_group=%s split_dim=%s",
                requested_size,
                world_size,
                split_dim,
        )
            return False
        return True
    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)
        if self._spatial_shard_decode_enabled(z):
-            split_dim = self._spatial_shard_decode_split_dim()
+            world_size = dist.get_world_size(group=self.distributed_executor.group)
            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
            assert split_dim is not None
            logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
            return wan_spatial_shard.spatial_shard_decode(
                self,
@@ -0,0 +1,86 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """QwenImage adapter for the shared spatial-shard VAE decoder kernels."""
 from __future__ import annotations
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
 from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
    SpatialShardCausalConv3dMixin,
    install_spatial_shard_decode,
    spatial_shard_decode_impl,
 )
 class QwenImageDistCausalConv3d(SpatialShardCausalConv3dMixin, QwenImageCausalConv3d):
    """Context-aware QwenImage causal conv that remains cache-discoverable."""
    def __init__(self, source: nn.Conv3d, group: dist.ProcessGroup) -> None:
        # Initialize Conv3d directly so the replacement preserves source device,
        # dtype, groups, and bias while remaining an isinstance() match for
        # QwenImageCausalConv3d. Diffusers clear_cache() relies on that match.
        nn.Conv3d.__init__(
            self,
            source.in_channels,
            source.out_channels,
            source.kernel_size,
            stride=source.stride,
            padding=0,
            dilation=source.dilation,
            groups=source.groups,
            bias=source.bias is not None,
            padding_mode=source.padding_mode,
            device=source.weight.device,
            dtype=source.weight.dtype,
        )
        self._init_spatial_shard_causal_conv(source, group)
    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
        return self._spatial_shard_causal_conv_forward(x, cache_x)
 def install_qwenimage_spatial_shard_decode(
    vae: SpatialShardVAE,
    group: dist.ProcessGroup,
    split_dim: str = "height",
 ) -> None:
    install_spatial_shard_decode(
        vae,
        group,
        split_dim,
        causal_conv_class_name="QwenImageCausalConv3d",
        causal_conv_factory=QwenImageDistCausalConv3d,
        attention_block_class_name="QwenImageAttentionBlock",
        installed_attr="_vllm_omni_qwenimage_spatial_shard_installed",
        model_name="QwenImage",
    )
 def spatial_shard_decode(
    vae: SpatialShardVAE,
    z: torch.Tensor,
    *,
    group: dist.ProcessGroup,
    return_dict: bool = True,
    split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
    """Decode QwenImage latents and preserve its all-rank output contract."""
    return spatial_shard_decode_impl(
        vae,
        z,
        group=group,
        install=install_qwenimage_spatial_shard_decode,
        model_name="QwenImage",
        pass_first_chunk=False,
        broadcast_result=True,
        unpatchify_patch_size=None,
        return_dict=return_dict,
        split_dim=split_dim,
    )
@@ -0,0 +1,142 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 from collections.abc import Callable
 from typing import Protocol
 import torch
 import torch.distributed as dist
 from vllm.logger import init_logger
 logger = init_logger(__name__)
 class SpatialShardExecutor(Protocol):
    group: dist.ProcessGroup
    parallel_mode: str
    parallel_size: int
    def _sync_final_result(
        self,
        rank0_result: torch.Tensor,
        output_ndim: int,
        output_device: torch.device,
        output_dtype: torch.dtype,
    ) -> torch.Tensor: ...
 class SpatialShardVAE(Protocol):
    distributed_executor: SpatialShardExecutor
    decoder: torch.nn.Module
    post_quant_conv: torch.nn.Module
    dtype: torch.dtype
    _conv_idx: list[int]
    _feat_map: list[object]
    def is_distributed_enabled(self) -> bool: ...
    def clear_cache(self) -> None: ...
 def select_auto_spatial_shard_split_dim(
    z_shape: tuple[int, ...],
    world_size: int,
 ) -> str | None:
    """Select the longer spatial axis that can be split across every rank."""
    if len(z_shape) != 5 or world_size <= 1:
        return None
    _, _, _, height, width = z_shape
    candidates = [
        (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
    ]
    if not candidates:
        return None
    # Prefer width for square latents. It was modestly faster in the original
    # four-GPU Wan benchmark and keeps the choice deterministic across models.
    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
 def resolve_spatial_shard_split_dim(
    mode: str,
    z: torch.Tensor | None = None,
    world_size: int | None = None,
 ) -> str | None:
    if mode == "spatial_shard_width":
        return "width"
    if mode == "spatial_shard_height":
        return "height"
    if mode == "auto" and z is not None and world_size is not None:
        return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
    return None
 def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool:
    """Apply the shared request- and topology-level spatial decode policy."""
    executor = vae.distributed_executor
    if executor.parallel_mode == "tile":
        return False
    if z.ndim != 5:
        logger.warning(
            "%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name
        )
        return False
    if not vae.is_distributed_enabled():
        return False
    group = executor.group
    world_size = dist.get_world_size(group=group)
    requested_size = int(executor.parallel_size)
    requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode)
    if requested_size != world_size:
        if executor.parallel_mode == "auto":
            logger.debug(
                "%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
                model_name,
                requested_size,
                world_size,
            )
        elif requested_split_dim is not None:
            logger.warning(
                "%s VAE spatial sharded decode currently requires vae_patch_parallel_size "
                "to match the DIT group size; falling back to tiled decode. "
                "requested=%s dit_group=%s split_dim=%s",
                model_name,
                requested_size,
                world_size,
                requested_split_dim,
            )
        return False
    return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None
 def prepare_spatial_shard_decode(
    vae: SpatialShardVAE,
    *,
    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
 ) -> None:
    """Install one VAE's dynamic wrappers while allocations are still tagged."""
    executor = vae.distributed_executor
    mode = executor.parallel_mode
    if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
        return
    requested_size = int(executor.parallel_size)
    world_size = dist.get_world_size(group=executor.group)
    if requested_size <= 1 or requested_size != world_size:
        return
    split_dim = "width" if mode == "spatial_shard_width" else "height"
    install(vae, executor.group, split_dim)
 def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
    """Prepare every VAE that advertises spatial-shard decode capability."""
    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
    if not dist.is_initialized():
        return
    for vae in ModuleDiscovery.discover(pipeline).vaes:
        prepare = getattr(vae, "_prepare_spatial_shard_decode", None)
        if callable(prepare):
            prepare()
@@ -7,19 +7,21 @@
 #   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
 # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
 # This version generalizes the height-only sharding to shard along height or
-# width and adds the Wan causal-conv ``feat_cache`` handling.
+# width and supports the causal-conv ``feat_cache`` handling used by Wan and
-"""Spatially-sharded Wan VAE decode.
+# QwenImage.
 """Shared spatial-shard kernels plus the Wan VAE adapter.
-The existing distributed Wan VAE path shards *tiles*.  This module adds an
+The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
-opt-in decode backend that shards decoder feature maps along height or width and
+are also reused by QwenImage's thin adapter. The backend shards decoder feature
-exchanges boundary rows/columns before spatial convolutions.  It is
+maps along height or width and exchanges boundary rows/columns before spatial
-intentionally decode-only and keeps checkpoint loading unchanged by patching the
+convolutions. It is decode-only and keeps checkpoint loading unchanged by
-already-loaded decoder.
+patching the already-loaded decoder.
 """
 from __future__ import annotations
 import math
 from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -34,6 +36,8 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
 logger = init_logger(__name__)
@@ -52,12 +56,17 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
 )
 def _active_spatial_shard_split_dim() -> str | None:
    context = _SPATIAL_SHARD_CONTEXT.get()
    return context.split_dim if context is not None else None
 def _spatial_dim(split_dim: str) -> int:
    if split_dim == "height":
        return -2
    if split_dim == "width":
        return -1
-    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
 def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
@@ -155,9 +164,9 @@ def split_for_parallel_decode(
    rank = 0 if rank is None else int(rank)
    world_size = 1 if world_size is None else int(world_size)
    if world_size < 1:
-        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
    if not 0 <= rank < world_size:
-        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
    dim = _spatial_dim(split_dim)
    expected_extent = x.shape[dim] * (2**upsample_count)
@@ -323,21 +332,27 @@ class WanDistZeroPad2d(nn.Module):
        padding: tuple[int, int, int, int],
        group: dist.ProcessGroup,
        *,
-        split_dim: str = "height",
+        defer_split_after_padding: bool = False,
        split_padding: tuple[int, int] | None = None,
    ) -> None:
        super().__init__()
        self.padding = padding
-        self.split_dim = split_dim
+        self.defer_split_after_padding = defer_split_after_padding
        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
        self.split_padding = split_padding or default_split_padding
        self.group = group
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        split_dim = _active_spatial_shard_split_dim()
        if split_dim is None:
            return F.pad(x, self.padding)
        rank, world_size = _rank_world(self.group)
        left, right, top, bottom = self.padding
        if self.defer_split_after_padding:
            if split_dim == "height":
                bottom = 0
            else:
                right = 0
        if world_size > 1:
-            if self.split_dim == "height":
+            if split_dim == "height":
                top = top if rank == 0 else 0
                bottom = bottom if rank == world_size - 1 else 0
            else:
@@ -351,8 +366,7 @@ class WanDistConv2d(nn.Conv2d):
        self,
        source: nn.Conv2d,
        group: dist.ProcessGroup,
-        split_dim: str = "height",
+        deferred_padding: tuple[int, int, int, int] | None = None,
        split_padding: tuple[int, int] | None = None,
    ):
        super().__init__(
            source.in_channels,
@@ -369,176 +383,253 @@ class WanDistConv2d(nn.Conv2d):
        )
        self.load_state_dict(source.state_dict())
        self.group = group
-        self.split_dim = split_dim
+        self._direct_padding = source.padding
-        self.split_tensor_dim = _spatial_dim(split_dim)
+        self._direct_padding_mode = source.padding_mode
-        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self.halo_size = (kernel_extent - 1) // 2
+        self._deferred_padding = deferred_padding
-        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-        if split_dim == "height":
+            raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-            if split_padding is None:
+        self._spatial_pad_h = int(source_padding[-2])
-                split_padding = (pad_h, pad_h)
+        self._spatial_pad_w = int(source_padding[-1])
-            self._non_split_padding = (pad_w, pad_w, 0, 0)
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-            self.kernel_extent = self.kernel_size[-2]
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-            self.stride_extent = self.stride[-2]
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
        else:
            if split_padding is None:
                split_padding = (pad_w, pad_w)
            self._non_split_padding = (0, 0, pad_h, pad_h)
            self.kernel_extent = self.kernel_size[-1]
            self.stride_extent = self.stride[-1]
        self.split_pad_left, self.split_pad_right = split_padding
        self._halo_recv_top_buf: torch.Tensor | None = None
        self._halo_recv_bottom_buf: torch.Tensor | None = None
        self._trim_cache: dict[int, tuple[int, int, int]] = {}
    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.pad(x, self._non_split_padding)
+        split_dim = _active_spatial_shard_split_dim()
        if split_dim is None:
            if self._direct_padding_mode != "zeros":
                x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
                padding: str | tuple[int, int] = (0, 0)
            else:
                padding = self._direct_padding
            return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
        split_tensor_dim = _spatial_dim(split_dim)
        if split_dim == "height":
            kernel_extent = self.kernel_size[-2]
            stride_extent = self.stride[-2]
            non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
            default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
            deferred_split_padding = (
                (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
            )
        else:
            kernel_extent = self.kernel_size[-1]
            stride_extent = self.stride[-1]
            non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
            default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
            deferred_split_padding = (
                (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
            )
        split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
        halo_size = (kernel_extent - 1) // 2
        x = F.pad(x, non_split_padding)
        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
            x,
            group=self.group,
-            halo_size=self.halo_size,
+            halo_size=halo_size,
-            split_dim=self.split_dim,
+            split_dim=split_dim,
            recv_top_buf=self._halo_recv_top_buf,
            recv_bottom_buf=self._halo_recv_bottom_buf,
        )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        shift, start, upper_bound = self._get_trim_params(
            x.shape[split_tensor_dim],
            split_dim=split_dim,
            halo_size=halo_size,
            split_pad_left=split_pad_left,
            split_pad_right=split_pad_right,
            kernel_extent=kernel_extent,
            stride_extent=stride_extent,
        )
        if shift:
            x_padded = _narrow_along_dim(
                x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
            )
        out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=split_dim)
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+    def _get_trim_params(
-        trim_params = self._trim_cache.get(local_extent)
+        self,
        local_extent: int,
        *,
        split_dim: str,
        halo_size: int,
        split_pad_left: int,
        split_pad_right: int,
        kernel_extent: int,
        stride_extent: int,
    ) -> tuple[int, int, int]:
        cache_key = (split_dim, local_extent)
        trim_params = self._trim_cache.get(cache_key)
        if trim_params is None:
            rank, world_size = _rank_world(self.group)
            trim_params = _compute_conv_trim_params(
                local_extent=local_extent,
                rank=rank,
                world_size=world_size,
-                halo_size=self.halo_size,
+                halo_size=halo_size,
-                pad_before=self.split_pad_left,
+                pad_before=split_pad_left,
-                pad_after=self.split_pad_right,
+                pad_after=split_pad_right,
-                kernel_extent=self.kernel_extent,
+                kernel_extent=kernel_extent,
-                stride_extent=self.stride_extent,
+                stride_extent=stride_extent,
            )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = trim_params
        return trim_params
-class WanDistCausalConv3d(nn.Conv3d):
+# Model-neutral aliases let other causal VAEs reuse the implementation while
-    def __init__(
+# preserving the original Wan class identities for stacked PRs and downstreams.
 SpatialShardZeroPad2d = WanDistZeroPad2d
 SpatialShardConv2d = WanDistConv2d
 class SpatialShardCausalConv3dMixin:
    """Shared causal-convolution behavior for supported Diffusers VAE classes."""
    def _init_spatial_shard_causal_conv(
        self,
        source: nn.Conv3d,
        group: dist.ProcessGroup,
-        split_dim: str = "height",
+    ) -> None:
    ):
        super().__init__(
            source.in_channels,
            source.out_channels,
            source.kernel_size,
            stride=source.stride,
            padding=0,
            dilation=source.dilation,
            groups=source.groups,
            bias=source.bias is not None,
            padding_mode=source.padding_mode,
            device=source.weight.device,
            dtype=source.weight.dtype,
        )
        self.load_state_dict(source.state_dict())
        self.group = group
        self.split_dim = split_dim
        self.split_tensor_dim = _spatial_dim(split_dim)
        source_padding = getattr(source, "_padding", None)
        if source_padding is None:
            p_t, p_h, p_w = source.padding
            source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
        self._source_padding = tuple(source_padding)
-        if split_dim == "height":
+        self._padding = tuple(source_padding)
-            self.split_pad_left = int(self._source_padding[2])
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-            self.split_pad_right = int(self._source_padding[3])
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-            self.kernel_extent = self.kernel_size[-2]
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
            self.stride_extent = self.stride[-2]
            padding = (
                self._source_padding[0],
                self._source_padding[1],
                0,
                0,
                self._source_padding[4],
                self._source_padding[5],
            )
        else:
            self.split_pad_left = int(self._source_padding[0])
            self.split_pad_right = int(self._source_padding[1])
            self.kernel_extent = self.kernel_size[-1]
            self.stride_extent = self.stride[-1]
            padding = (
                0,
                0,
                self._source_padding[2],
                self._source_padding[3],
                self._source_padding[4],
                self._source_padding[5],
            )
        self.halo_size = (self.kernel_extent - 1) // 2
        self._padding = padding if self.halo_size > 0 else self._source_padding
        self._halo_recv_top_buf: torch.Tensor | None = None
        self._halo_recv_bottom_buf: torch.Tensor | None = None
        self._trim_cache: dict[int, tuple[int, int, int]] = {}
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+    def _spatial_shard_causal_conv_forward(
-        padding = list(self._padding)
+        self,
        x: torch.Tensor,
        cache_x: torch.Tensor | None = None,
    ) -> torch.Tensor:
        split_dim = _active_spatial_shard_split_dim()
        padding = list(self._source_padding)
        if cache_x is not None and padding[4] > 0:
            cache_x = cache_x.to(x.device)
            x = torch.cat([cache_x, x], dim=2)
            padding[4] -= cache_x.shape[2]
        if split_dim is None:
            return nn.Conv3d.forward(self, F.pad(x, padding))
        split_tensor_dim = _spatial_dim(split_dim)
        if split_dim == "height":
            split_pad_left = int(padding[2])
            split_pad_right = int(padding[3])
            kernel_extent = self.kernel_size[-2]
            stride_extent = self.stride[-2]
            padding[2] = 0
            padding[3] = 0
        else:
            split_pad_left = int(padding[0])
            split_pad_right = int(padding[1])
            kernel_extent = self.kernel_size[-1]
            stride_extent = self.stride[-1]
            padding[0] = 0
            padding[1] = 0
        halo_size = (kernel_extent - 1) // 2
        x = F.pad(x, padding)
        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
            x,
            group=self.group,
-            halo_size=self.halo_size,
+            halo_size=halo_size,
-            split_dim=self.split_dim,
+            split_dim=split_dim,
            recv_top_buf=self._halo_recv_top_buf,
            recv_bottom_buf=self._halo_recv_bottom_buf,
        )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        shift, start, upper_bound = self._get_trim_params(
            x.shape[split_tensor_dim],
            split_dim=split_dim,
            halo_size=halo_size,
            split_pad_left=split_pad_left,
            split_pad_right=split_pad_right,
            kernel_extent=kernel_extent,
            stride_extent=stride_extent,
        )
        if shift:
            x_padded = _narrow_along_dim(
                x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
            )
-        out = super().forward(x_padded)
+        out = nn.Conv3d.forward(self, x_padded)
-        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=split_dim)
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+    def _get_trim_params(
-        trim_params = self._trim_cache.get(local_extent)
+        self,
        local_extent: int,
        *,
        split_dim: str,
        halo_size: int,
        split_pad_left: int,
        split_pad_right: int,
        kernel_extent: int,
        stride_extent: int,
    ) -> tuple[int, int, int]:
        cache_key = (split_dim, local_extent)
        trim_params = self._trim_cache.get(cache_key)
        if trim_params is None:
            rank, world_size = _rank_world(self.group)
            trim_params = _compute_conv_trim_params(
                local_extent=local_extent,
                rank=rank,
                world_size=world_size,
-                halo_size=self.halo_size,
+                halo_size=halo_size,
-                pad_before=self.split_pad_left,
+                pad_before=split_pad_left,
-                pad_after=self.split_pad_right,
+                pad_after=split_pad_right,
-                kernel_extent=self.kernel_extent,
+                kernel_extent=kernel_extent,
-                stride_extent=self.stride_extent,
+                stride_extent=stride_extent,
            )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = trim_params
        return trim_params
 class WanDistCausalConv3d(SpatialShardCausalConv3dMixin, nn.Conv3d):
    # Cross-PR compatibility contract for the lossless Wan data-movement
    # installer: it may fuse the no-context direct path, but must leave this
    # wrapper in control whenever a spatial-shard context is active.
    _vllm_omni_dynamic_spatial_shard_conv = True
    def __init__(
        self,
        source: nn.Conv3d,
        group: dist.ProcessGroup,
    ):
        super().__init__(
            source.in_channels,
            source.out_channels,
            source.kernel_size,
            stride=source.stride,
            padding=0,
            dilation=source.dilation,
            groups=source.groups,
            bias=source.bias is not None,
            padding_mode=source.padding_mode,
            device=source.weight.device,
            dtype=source.weight.dtype,
        )
        self._init_spatial_shard_causal_conv(source, group)
    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
        return self._spatial_shard_causal_conv_forward(x, cache_x)
 def _compute_conv_trim_params(
    *,
    local_extent: int,
@@ -582,12 +673,15 @@ def _trim_local_conv_output(
    return out
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
    if getattr(module, "_vllm_omni_spatial_shard_attention", False):
        return
    orig_forward = module.forward
    def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
        split_dim = _active_spatial_shard_split_dim()
        if split_dim is None:
            return orig_forward(x, *args, **kwargs)
        _, world_size = _rank_world(group)
        if world_size <= 1:
            return orig_forward(x, *args, **kwargs)
@@ -609,60 +703,43 @@ def _replace_child(
    name: str,
    child: nn.Module,
    group: dist.ProcessGroup,
-    split_dim: str,
+    *,
    causal_conv_class_name: str,
    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
 ) -> None:
-    if child.__class__.__name__ == "WanCausalConv3d":
+    if child.__class__.__name__ == causal_conv_class_name:
-        setattr(
+        setattr(parent, name, causal_conv_factory(child, group))
            parent,
            name,
            WanDistCausalConv3d(
                child,
                group,
                split_dim=split_dim,
            ),
        )
        return
    if isinstance(child, nn.ZeroPad2d):
        padding = tuple(int(p) for p in child.padding)
        module_padding = padding
        if parent.__class__.__name__ == "Sequential":
            # Let the following WanDistConv2d account for global after-edge
            # padding; this module handles the non-split dimension and before edge.
            if split_dim == "height":
                module_padding = (padding[0], padding[1], padding[2], 0)
            else:
                module_padding = (padding[0], 0, padding[2], padding[3])
        setattr(
            parent,
            name,
-            WanDistZeroPad2d(
+            SpatialShardZeroPad2d(
-                module_padding,
+                padding,
                group,
-                split_dim=split_dim,
+                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
            ),
        )
        return
    if isinstance(child, nn.Conv2d):
-        split_padding = None
+        deferred_padding = None
        if name == "1" and parent.__class__.__name__ == "Sequential":
-            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
+            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
            # stride-2 conv with padding=0.  Only the last rank should see the
            # bottom/right global padding, which is approximated by split_padding.
            prev = getattr(parent, "0", None)
-            if isinstance(prev, WanDistZeroPad2d):
+            if isinstance(prev, SpatialShardZeroPad2d):
-                split_padding = prev.split_padding
+                deferred_padding = prev.padding
            elif isinstance(prev, nn.ZeroPad2d):
-                pad = prev.padding
+                deferred_padding = tuple(int(value) for value in prev.padding)
                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
        setattr(
            parent,
            name,
-            WanDistConv2d(
+            SpatialShardConv2d(
                child,
                group,
-                split_dim=split_dim,
+                deferred_padding=deferred_padding,
                split_padding=split_padding,
            ),
        )
@@ -670,21 +747,45 @@ def _replace_child(
 def _patch_decoder_modules(
    module: nn.Module,
    group: dist.ProcessGroup,
-    split_dim: str,
+    *,
    causal_conv_class_name: str,
    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
    attention_block_class_name: str,
    inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == "WanAttentionBlock":
+    if module.__class__.__name__ == attention_block_class_name:
-        _patch_attention_block(module, group, split_dim)
+        _patch_attention_block(module, group)
        inside_attention = True
    for name, child in list(module.named_children()):
-        if child.__class__.__name__ == "WanCausalConv3d":
+        if child.__class__.__name__ == causal_conv_class_name:
-            _replace_child(module, name, child, group, split_dim)
+            _replace_child(
                module,
                name,
                child,
                group,
                causal_conv_class_name=causal_conv_class_name,
                causal_conv_factory=causal_conv_factory,
            )
            continue
        if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(module, name, child, group, split_dim)
+            _replace_child(
                module,
                name,
                child,
                group,
                causal_conv_class_name=causal_conv_class_name,
                causal_conv_factory=causal_conv_factory,
            )
            continue
-        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
+        _patch_decoder_modules(
            child,
            group,
            causal_conv_class_name=causal_conv_class_name,
            causal_conv_factory=causal_conv_factory,
            attention_block_class_name=attention_block_class_name,
            inside_attention=inside_attention,
        )
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -695,47 +796,58 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
    return count
-def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
+def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
    decoder = getattr(vae, "decoder", None)
    if decoder is None:
        return
    for module in decoder.modules():
        if "_halo_recv_top_buf" in module._buffers:
            module._halo_recv_top_buf = None
        if "_halo_recv_bottom_buf" in module._buffers:
            module._halo_recv_bottom_buf = None
 def install_spatial_shard_decode(
    vae: SpatialShardVAE,
    group: dist.ProcessGroup,
    split_dim: str,
    *,
    causal_conv_class_name: str,
    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
    attention_block_class_name: str,
    installed_attr: str,
    model_name: str,
 ) -> None:
    """Patch ``vae.decoder`` once for spatially-sharded decode.
    This mutates the already-loaded decoder in place by swapping its spatial
-    convolutions/padding for halo-exchanging variants and wrapping
+    convolutions/padding for context-aware variants and attaching a dedicated
-    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    spatial forward. Outside that forward the replacements execute their exact
-    instance and is applied only once (subsequent calls are no-ops). A given
+    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    instance is bound to a single ``split_dim``; switching between
+    The same installed decoder can alternate height and width sharding.
    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
    otherwise.
-    Only group-relative rank 0 assembles the final decoded frame, mirroring the
+    Group-relative rank 0 assembles the final decoded frame. The model adapter
-    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
+    then preserves its prior contract: Wan keeps the result on rank 0, while
-    take part in the collectives but return an empty placeholder.
+    QwenImage broadcasts it to every rank.
    """
    _spatial_dim(split_dim)
-    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+    if getattr(vae, installed_attr, False):
        installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
        if installed_split_dim != split_dim:
            raise ValueError(
                "Wan spatial-shard VAE decoder was already patched for "
                f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
            )
        return
    decoder = getattr(vae, "decoder", None)
    if decoder is None:
-        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
+        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
-    _patch_decoder_modules(decoder, group, split_dim)
+    _patch_decoder_modules(
        decoder,
        group,
        causal_conv_class_name=causal_conv_class_name,
        causal_conv_factory=causal_conv_factory,
        attention_block_class_name=attention_block_class_name,
    )
    upsample_count = _decoder_upsample_count(decoder)
    orig_forward = decoder.forward
-    def _forward(
+    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
        self: nn.Module,
        x: torch.Tensor,
        feat_cache: list[torch.Tensor] | None = None,
        feat_idx: list[int] | None = None,
        first_chunk: bool = False,
    ) -> torch.Tensor:
        if feat_idx is None:
            feat_idx = [0]
        tensor_dim = _spatial_dim(split_dim)
        input_extent = x.shape[tensor_dim]
        x, expected_extent = split_for_parallel_decode(
@@ -755,29 +867,50 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
            )
        )
        try:
-            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+            out = orig_forward(x, *args, **kwargs)
        finally:
            _SPATIAL_SHARD_CONTEXT.reset(token)
        return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
-    decoder.forward = MethodType(_forward, decoder)
+    decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
-    vae._vllm_omni_wan_spatial_shard_installed = True
+    setattr(vae, installed_attr, True)
-    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
-def spatial_shard_decode(
+def install_wan_spatial_shard_decode(
-    vae: Any,
+    vae: SpatialShardVAE,
    group: dist.ProcessGroup,
    split_dim: str = "height",
 ) -> None:
    install_spatial_shard_decode(
        vae,
        group,
        split_dim,
        causal_conv_class_name="WanCausalConv3d",
        causal_conv_factory=WanDistCausalConv3d,
        attention_block_class_name="WanAttentionBlock",
        installed_attr="_vllm_omni_wan_spatial_shard_installed",
        model_name="Wan",
    )
 def spatial_shard_decode_impl(
    vae: SpatialShardVAE,
    z: torch.Tensor,
    *,
    group: dist.ProcessGroup,
    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
    model_name: str,
    pass_first_chunk: bool,
    broadcast_result: bool,
    unpatchify_patch_size: int | None,
    return_dict: bool = True,
    split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
-    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
+    install(vae, group, split_dim)
    if z.shape[2] == 0:
-        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
    # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
    # the halo/all-gather collectives; they just skip keeping/assembling the output.
@@ -786,31 +919,72 @@ def spatial_shard_decode(
    vae.clear_cache()
    try:
-        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
+        context_factory = getattr(vae, "_execution_context", None)
        context = context_factory() if callable(context_factory) else nullcontext()
        with context:
            x = vae.post_quant_conv(z)
            decoded_chunks = []
            spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward")
            for i in range(z.shape[2]):
                vae._conv_idx = [0]
-                chunk = vae.decoder(
+                if pass_first_chunk:
                    chunk = spatial_forward(
                        x[:, :, i : i + 1, :, :],
                        feat_cache=vae._feat_map,
                        feat_idx=vae._conv_idx,
-                    first_chunk=(i == 0),
+                        first_chunk=i == 0,
                        split_dim=split_dim,
                    )
                else:
                    chunk = spatial_forward(
                        x[:, :, i : i + 1, :, :],
                        feat_cache=vae._feat_map,
                        feat_idx=vae._conv_idx,
                        split_dim=split_dim,
                    )
                if produce_output:
                    decoded_chunks.append(chunk)
            if produce_output:
                out = torch.cat(decoded_chunks, dim=2)
-                if vae.config.patch_size is not None:
+                if unpatchify_patch_size is not None:
-                    out = unpatchify(out, patch_size=vae.config.patch_size)
+                    out = unpatchify(out, patch_size=unpatchify_patch_size)
                out = torch.clamp(out, min=-1.0, max=1.0)
            else:
                out = z.new_zeros(0)
    finally:
        vae.clear_cache()
        # Halo buffers are request-shape scratch space allocated outside the
        # tagged weight pool. Drop live references so sleep/offload never has
        # to preserve or discard them as model state.
        clear_spatial_shard_runtime_buffers(vae)
    if broadcast_result:
        out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype)
    if not return_dict:
        return (out,)
    return DecoderOutput(sample=out)
 def spatial_shard_decode(
    vae: SpatialShardVAE,
    z: torch.Tensor,
    *,
    group: dist.ProcessGroup,
    return_dict: bool = True,
    split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
    """Decode a Wan latent with spatial activation sharding."""
    return spatial_shard_decode_impl(
        vae,
        z,
        group=group,
        install=install_wan_spatial_shard_decode,
        model_name="Wan",
        pass_first_chunk=True,
        broadcast_result=False,
        unpatchify_patch_size=vae.config.patch_size,
        return_dict=return_dict,
        split_dim=split_dim,
    )
@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
 # lookup and masked out afterwards).
 MINIMAX_H3_ADALN_MODALITY_NUM = 3
 # Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
 # (consumed only via the "laser_input_scale" extra key; other backends and
 # platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
 # workspace, and H3's outlier activations (per-element amax in the hundreds)
 # push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
 # 256 is a power of two, so pre-dividing q/k/v and the compensating
 # kernel-scale/output multiplies are exact in floating point.
 MINIMAX_H3_LASER_INPUT_SCALE = 256.0
 def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
    if key not in kwargs or kwargs[key] is None:
@@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
        used = min(max_seqlen, packed_total)
        attn_mask = None
        # Ring attention can dispatch to a different implementation from the
-        # configured backend, so this no-mask fast path is local-only.
+        # configured backend, so the no-mask fast paths are local-only.
-        prefix_slice = (
+        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
-            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
+        # supports_packed_mask_free: backend consumes the packed metadata
        # without ever reading attn_mask (CUDA packed varlen, NPU
        # npu_attn_varlen opt-in with its own fallback rebuild).
        no_mask = not getattr(self.attention, "use_ring", False) and (
            self.attention.attn_backend.supports_prefix_kv_slicing
            or self.attention.attn_backend.supports_packed_mask_free()
        )
-        if used < packed_total and not prefix_slice:
+        if used < packed_total and not no_mask:
            attn_mask = torch.arange(packed_total, device=q.device)[None] < used
        metadata = AttentionMetadata(
            attn_mask=attn_mask,
@@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
                "max_seqlen_q": max_seqlen,
                "max_seqlen_k": max_seqlen,
                "valid_kv_length": used,
                # Opt the NPU flash backend into the packed varlen path so the
                # quadratic full_qk mask is never materialized. Ring attention
                # is excluded: it keeps the aligned padding rows for its
                # fixed-size P2P buffers and still needs the mask.
                "npu_attn_varlen": not getattr(self.attention, "use_ring", False),
                # fp16-range protection for the ascend_laser_attention kernel
                # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
                # backend/path.
                "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
            },
            video_layout=video_layout,
        )
@@ -155,6 +155,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
        parallel_size: int,
        mode: str = "tile",
    ) -> None:
        # ``auto`` is the cross-model default. MiniMax H3 has no spatial-shard
        # capability, so preserve its native tile path for that default while
        # continuing to reject explicitly requested spatial modes.
        if mode == "auto":
            mode = "tile"
        if mode != "tile":
            raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
        group = get_dit_group()
@@ -30,6 +30,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
    prepare_pipeline_spatial_shard_decode,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                    custom_pipeline_name=custom_pipeline_name,
                    device=self.device,
                )
                prepare_pipeline_spatial_shard_decode(self.pipeline)
        time_after_load = time.perf_counter()
        logger.info(
@@ -505,7 +505,7 @@ class OrchestratorArgs:
    diffusion_kv_cache_skip_layers: str | None = None
    cfg_parallel_size: int = 1
    vae_patch_parallel_size: int = 1
-    vae_parallel_mode: str = "tile"
+    vae_parallel_mode: str = "auto"
    text_encoder_tp_size: int = 1
    default_sampling_params: str | None = None
    max_generated_image_size: int | None = None
@@ -972,7 +972,7 @@ class AsyncOmniEngine:
            cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
            vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
-            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "auto"
            text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
            enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
            use_hsdp = normalized_kwargs.get("use_hsdp", False)
@@ -750,10 +750,12 @@ class OmniServeCommand(CLISubcommand):
        omni_config_group.add_argument(
            "--vae-parallel-mode",
            type=str,
-            default="tile",
+            default="auto",
-            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
+            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
            help="VAE parallel decode strategy for diffusion models. "
-            "'tile' (default) uses patch/tile parallel decode; "
+            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
            "tiled decodes "
            "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; "
            "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
            "decoder feature maps along height/width and exchanges halo regions. The "
            "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "
@@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
                )
                backend_upper = "FLASH_ATTN"
            if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
                # The NPU FLASH_ATTN backend imports mindiesd lazily at first
                # forward, but CANN snapshots the custom-op registry at the
                # first custom-op regInfo lookup in the process (e.g. a
                # vllm-ascend custom op during model load/warmup). Import
                # mindiesd here so its env.py prepends the mindiesd vendor
                # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
                # that snapshot; otherwise aclnnLaserAttention /
                # FusedAttentionScore fail with EZ1001 "does not support
                # opType" for the rest of the process.
                import mindiesd  # noqa: F401
            backend = DiffusionAttentionBackendEnum[backend_upper]
            logger.debug("Using diffusion attention backend '%s'", backend_upper)
            return backend.get_path()
Author	SHA1	Message	Date
Hongsheng Liu	fd82e41fff	feat(diffusion): add QwenImage spatial-shard decode Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 10:49:10 +08:00
Hongsheng Liu	810112189b	feat(diffusion): select Wan VAE decode mode automatically Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 09:41:05 +08:00
fan2956andFan	c27623c29c	[Bugfix] rainfusion support all shape video (#6000 ) Signed-off-by: Fan <fan@FandeMacBook-Air.local> Co-authored-by: Fan <fan@FandeMacBook-Air.local>	2026-08-11 05:10:50 +08:00
Yueqian Lin	f8fa86fb1f	[BugFix] Don't fail the TTS ratchet when the branch count goes down (#6008 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-08-11 05:08:31 +08:00
ooooooyeandHongsheng Liu	422d8fd8f9	[Bugfix][Diffusion] Restore supports_packed_mask_free on teacache FakeBackend (#5997 ) Signed-off-by: brandneway <gyuan4892@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 22:20:37 +08:00
Hongsheng Liu	62e239bf1e	[Docs] Fix broken attention backend link (#5998 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 21:55:01 +08:00
ooooooye	740f45abcd	[Perf][Diffusion] MiniMax-H3: opt into packed varlen attention on NPU to eliminate quadratic mask materialization (#5891 ) Signed-off-by: brandneway <gyuan4892@gmail.com>	2026-08-10 19:47:03 +08:00
Hongsheng Liu hsliu_ustc Copilot Autofix powered by AI	5e4b9f8c6e	[Misc] Align CODEOWNERS with module/feature design docs (#5958 ) Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>	2026-08-10 19:25:03 +08:00
Hongsheng Liu	ae811043e9	docs: keep shared task examples in navigation (#5987 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 19:13:52 +08:00