Merge origin/main to resolve TTS ratchet conflict

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.github/CODEOWNERS

@@ -1,11 +1,6 @@
 # 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
@@ -26,10 +21,25 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao
 
-# ComfyUI / external apps
+# Engine and model_executor
-/apps/ @fhfuih @alex-jw-brooks
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/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
 
-# Benchmarks (out-of-tree)
+# Diffusion
+/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
@@ -37,112 +47,20 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
 
-# =====================================================================
+# Usability and Observability
-# Module design ownership (docs/design/module/*.md)
+/vllm_omni/deploy/ @tzhouam @lishunyang12
-# Each section header cites the source design doc.
+/vllm_omni/metrics/ @lishunyang12
-# =====================================================================
-
-# Autoregressive runtime — docs/design/module/ar_runtime.md
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
-/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
-
-# 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
-
-# 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
-
-# Observability — docs/design/module/observability.md
-/vllm_omni/metrics/ @lishunyang12 @vraiti
-
-# Profiling — docs/design/module/profiling.md
 /vllm_omni/profiler/ @gcanlin
 
-# Benchmarking (in-tree) — docs/design/module/benchmarking.md
+# ComfyUI
-/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
+/apps/ @fhfuih @alex-jw-brooks
 
-# LoRA (no dedicated design doc; ownership intentionally narrowed)
+# Large Scale Serving
-/vllm_omni/lora/ @SamitHuang @wtomin
+/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
 
-# Transformers utils (related path only in design docs; preserves prior ownership)
+# Quantization and configuration
-/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
+/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
+/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
 
-# =====================================================================
+# Hardware plugins and accelerators
-# Diffusion module family — docs/design/module/diffusion/
+/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
-# 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

.pre-commit-config.yaml

@@ -66,10 +66,7 @@ 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)|tools/pre_commit/check_tts_adapter\.py)$
+      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.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

docs/.nav.yml

@@ -16,21 +16,31 @@ nav:
   - Examples:
     - examples/README.md
     - Offline Inference:
+      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
+      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
+      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
+      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
-      - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
       - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
+      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
+      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
-      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
+      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
+      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
       - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*

docs/design/feature/vae_parallel.md

@@ -449,21 +449,20 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-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.
+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"`).
 
 ### How it differs from tile parallel
 
-| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|----------------------|-----------------------------------------------|
+|--------|--------------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
+| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
-| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks |
+| Scope | Decode + encode | Decode only |
-| Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only |
 
-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.
+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`.
 
 ### Wiring
 
@@ -474,17 +473,15 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
   -> DiffusionParallelConfig.vae_parallel_mode
   -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
   -> DistributedVaeExecutor.parallel_mode
-  -> capability-based pipeline preparation
+  -> DistributedAutoencoderKLWan.tiled_decode dispatch
-  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
 ```
 
-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.
+`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.
 
 ### Notes
 
-- 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.
+- 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.
 - 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`.
 
 ---
 

docs/design/module/ar_runtime.md

@@ -7,7 +7,6 @@ owners:
   - "@yinpeiqi"
   - "@fake0fan"
   - "@Sy0307"
-  - "@Gaohan123"
 primary_code_paths:
   - vllm_omni/core/**
   - vllm_omni/worker/**

docs/design/module/benchmarking.md

@@ -4,7 +4,6 @@ kind: module
 status: draft
 owners:
   - "@alex-jw-brooks"
-  - "@Bounty-hunter"
 primary_code_paths:
   - vllm_omni/benchmarks/**
   - vllm_omni/entrypoints/cli/benchmark/**

docs/design/module/cache_management.md

@@ -5,7 +5,6 @@ status: draft
 owners:
   - "@Isotr0py"
   - "@princepride"
-  - "@SamitHuang"
 primary_code_paths:
   - vllm_omni/diffusion/cache/**
 related_code_paths:

docs/design/module/diffusion/diffusion_model_integration.md

@@ -3,10 +3,9 @@ title: Diffusion Model Integration
 kind: module
 status: draft
 owners:
-  - "@fhfuih"
+  - "@Isotr0py"
-  - "@Bounty-hunter"
   - "@wtomin"
-  - "@RuixiangMa"
+  - "@ZJY0516"
 primary_code_paths:
   - vllm_omni/diffusion/models/**
   - vllm_omni/diffusion/model_loader/**

docs/design/module/diffusion/diffusion_runtime.md

@@ -6,7 +6,6 @@ owners:
   - "@Isotr0py"
   - "@princepride"
   - "@SamitHuang"
-  - "@fhfuih"
 primary_code_paths:
   - vllm_omni/diffusion/sched/**
   - vllm_omni/diffusion/executor/**

docs/design/module/diffusion/offloader.md

@@ -4,8 +4,7 @@ kind: module
 status: draft
 owners:
   - "@yuanheng-zhao"
-  - "@david6666666"
+  - "@Isotr0py"
-  - "@lishunyang12"
 primary_code_paths:
   - vllm_omni/diffusion/offloader/**
 related_code_paths:

docs/design/module/diffusion/parallelism.md

@@ -6,7 +6,6 @@ owners:
   - "@princepride"
   - "@yuanheng-zhao"
   - "@xuechendi"
-  - "@Bounty-hunter"
 primary_code_paths:
   - vllm_omni/diffusion/distributed/**
   - vllm_omni/diffusion/attention/parallel/**

docs/design/module/entrypoints.md

@@ -6,7 +6,6 @@ architecture_state: current-plus-in-flight
 owners:
   - "@alex-jw-brooks"
   - "@linyueqian"
-  - "@NickCao"
 document_stewards:
   - "@hsliuustc0106"
   - "@Gaohan123"

docs/design/module/execution_platforms.md

@@ -8,15 +8,8 @@ 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

docs/design/module/input_output_modality_contracts.md

@@ -6,8 +6,6 @@ architecture_state: current
 owners:
   - "@Sy0307"
   - "@amy-why-3459"
-  - "@Gaohan123"
-  - "@alex-jw-brooks"
 document_stewards:
   - "@hsliuustc0106"
   - "@Gaohan123"

docs/design/module/model_integration.md

@@ -9,14 +9,6 @@ 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/**

docs/design/module/observability.md

@@ -4,7 +4,6 @@ kind: module
 status: draft
 owners:
   - "@lishunyang12"
-  - "@vraiti"
 primary_code_paths:
   - vllm_omni/metrics/**
 related_code_paths:

docs/mkdocs/hooks/generate_examples.py

@@ -24,22 +24,15 @@ SERVING_MODE_TITLES = {
     "online_serving": "Online serving",
 }
 
-# Keep the Examples navigation focused on reusable modality/task entry points.
+# Generated example pages keep their source-derived URLs even when their
-# Model-specific scripts remain available under ``examples/`` and are surfaced
+# navigation category changes. Quantization helpers are feature documentation,
-# from these shared pages instead of getting one page per model.
+# not user-serving examples, so keep them under the existing Features /
-GENERAL_EXAMPLE_SLUGS = frozenset(
+# Quantization section instead of creating a second top-level Quantization
-    {
+# category under Examples.
-        "image_to_image",
+EXAMPLE_NAV_TARGETS = {
-        "image_to_video",
+    "quantization": "quantization",
-        "speech_to_video",
+}
-        "text_to_audio",
+GENERATED_QUANTIZATION_EXAMPLES_TITLE = "Generated Examples"
-        "text_to_image",
-        "text_to_speech",
-        "text_to_video",
-        "x_to_text",
-        "x_to_video_audio",
-    }
-)
 
 
 def load_model_display_names() -> dict[str, str]:
@@ -63,13 +56,6 @@ def load_model_display_names() -> dict[str, str]:
 MODEL_DISPLAY_NAMES = load_model_display_names()
 
 
-def is_general_example(example: "Example") -> bool:
-    """Return whether an example should get a generated serving page."""
-    if example.category not in SERVING_MODE_TITLES:
-        return True
-    return example.path.stem in GENERAL_EXAMPLE_SLUGS
-
-
 def fix_case(text: str) -> str:
     subs = {
         "api": "API",
@@ -372,9 +358,12 @@ def update_nav_file(examples: list[Example]):
             items.append({example.nav_title: rel_path.as_posix()})
         return items
 
-    # Build the new Examples section - start with preserved items.
+    # Build the new Examples section - start with preserved items. Categories
+    # not explicitly routed elsewhere retain the historical Examples behavior.
     examples_section = preserved_items.copy()
     for category in sorted(examples_by_category.keys()):
+        if EXAMPLE_NAV_TARGETS.get(category, "examples") != "examples":
+            continue
         items = category_items(category)
         if items:
             category_title = fix_case(category.replace("_", " ").title())
@@ -383,6 +372,36 @@ def update_nav_file(examples: list[Example]):
     # Update the nav structure
     nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
 
+    features_content = next(
+        (item["Features"] for item in user_guide_content if isinstance(item, dict) and "Features" in item),
+        None,
+    )
+    quantization_content = None
+    if isinstance(features_content, list):
+        quantization_content = next(
+            (item["Quantization"] for item in features_content if isinstance(item, dict) and "Quantization" in item),
+            None,
+        )
+
+    if isinstance(quantization_content, list):
+        # Remove only the group owned by this hook so repeated builds are
+        # idempotent and hand-authored Quantization navigation is retained.
+        quantization_content[:] = [
+            item
+            for item in quantization_content
+            if not (isinstance(item, dict) and GENERATED_QUANTIZATION_EXAMPLES_TITLE in item)
+        ]
+
+    quantization_examples = examples_by_category.get("quantization")
+    if quantization_examples:
+        if not isinstance(quantization_content, list):
+            logger.warning(
+                "Could not find 'User Guide' -> 'Features' -> 'Quantization' section; "
+                "omitting generated quantization examples from navigation"
+            )
+        else:
+            quantization_content.append({GENERATED_QUANTIZATION_EXAMPLES_TITLE: category_items("quantization")})
+
     # Write back to file
     nav_data["nav"] = nav_list
     with open(NAV_FILE, "w", encoding="utf-8") as f:
@@ -413,8 +432,6 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
         for path in category.glob("*/*.md"):
             examples.append(Example(path.parent, category.stem))
 
-    examples = [example for example in examples if is_general_example(example)]
-
     # Generate the example documentation
     for example in sorted(examples, key=lambda e: e.path.stem):
         example_name = f"{example.path.stem}.md"

docs/user_guide/diffusion/attention_backends.md

@@ -283,20 +283,23 @@ blocks — use Ulysses SP (`ring_degree=1`).
 
 ### Which geometries run sparse
 
-MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
+Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
-necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
+is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
-does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
+tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
-promoted rows remain visible to every query and key.
+that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+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.
 
-For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
+Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
-128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
+spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
-MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
+patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
-and kernel tiling aligned, at the cost of retaining more blocks.
+are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+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.
 
-8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
+"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
-creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
+resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
-grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+which you can buy back with `start_step`.
-that implements protected video tails for rf_v2.
 
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -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/QwenImage 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 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 | `"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). |
+| `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). |
 
 Additional requirements:
 
@@ -126,9 +126,9 @@ Additional requirements:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-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.
+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"`.
 
 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="auto",                # default
+        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
     ),
 )
 ```
@@ -152,18 +152,14 @@ 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 auto
+    --vae-parallel-mode spatial_shard_width
 ```
 
 **Constraints and behavior:**
 
-- 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.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
-- 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.
+- 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.
-- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
+- `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).
-- 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:
 

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -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](../../diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/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.

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -79,8 +79,7 @@ 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 distributed
+degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
-layerwise offload:
 
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -99,7 +98,7 @@ vllm serve "${MODEL}" \
   --usp 8 \
   --ring 1 \
   --text-encoder-tp-size 8 \
-  --enable-distributed-layerwise-offload \
+  --enable-layerwise-offload \
   --vae-parallel-mode tile \
   --vae-use-tiling \
   --vae-patch-parallel-size 8 \
@@ -113,11 +112,6 @@ 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
@@ -140,63 +134,6 @@ 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
@@ -217,8 +154,8 @@ configuration above:
 
 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 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 |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
 
 These measurements describe the validated shapes rather than a general
 throughput guarantee.

recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md

@@ -4,17 +4,6 @@ This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
 resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
 recipes it uses **no offload of any kind** — see the capacity note below.
 
-Validated on:
-
-- Host: DGX Spark (GB10), aarch64
-- GPUs: 1 (unified memory)
-- Driver: <FILL_DRIVER_VERSION>
-- vLLM: 0.26.0
-- vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
-- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
-  `seed=1101`
-
-
 ## Capacity requirements
 
 | Resource | DGX Spark (GB10) |
@@ -22,8 +11,7 @@ Validated on:
 | Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
 | GPUs | 1 |
 | Checkpoint storage | 135 GiB per partition |
-| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
+| Measured peak allocator high-water (FP8) | 97.7 GiB |
-| Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
 
 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
@@ -41,21 +29,9 @@ math relative to the RTX and datacenter recipes:
   fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
   Qwen3-VL text encoder and both VAEs stay BF16.
 
-With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
+With ~97.7 GiB of the pool held by the allocator at peak, there is roughly 23 GiB
-23 GiB of headroom for the OS, page cache, and any other process on the box.
+of headroom for the OS, page cache, and any other process on the box. Treat this
-Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
+platform as single-tenant while a request is in flight.
-and the allocator figure understates true pool usage.
-
-Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
-`MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
-a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
-roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
-which is the CUDA context, non-PyTorch allocations, and page cache the header
-does not count.
-
-Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
-second process, and adding reference images, reference videos, or a larger output
-shape will run into the OOM killer.
 
 ## One GB10: 960x576, 8 seconds
 
@@ -118,52 +94,6 @@ ffprobe -v error -show_entries \
 Both requests return `200 OK` with an MP4 body: H.264 video at the requested
 geometry plus a 32 kHz stereo AAC track.
 
-## One GB10: Ref2VA, 960x576, 8 seconds
-
-`FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
-Stop the FL2VA server before starting this one.
-
-```bash
-CUDA_VISIBLE_DEVICES=0 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
-VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
-vllm serve /path/to/MiniMax-H3/Ref2VA \
-  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
-  --init-timeout 3600 \
-  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
-  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
-  --vae-parallel-mode tile --vae-use-tiling \
-  --quantization fp8 --enforce-eager \
-  --diffusion-attention-backend CUDNN_ATTN
-```
-
-The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
-takes over an hour on this platform. Use 14400 s.
-
-Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
-omitted when `width` and `height` are supplied.
-
-```bash
-curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
-  -D ref2va-headers.txt \
-  -w 'client_e2e_s=%{time_total}\n' \
-  --max-time 14400 \
-  -F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
-  -F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
-  -F 'width=960' \
-  -F 'height=576' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=50' \
-  -F 'flow_shift=12' \
-  -F 'seed=1101' \
-  -F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
-  -o h3-gb10-ref2va.mp4
-```
-
-The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
-32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
-
 ## Measured latency
 
 Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
@@ -172,69 +102,32 @@ Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
 | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
 | T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
 | T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
-| Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
-| Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
 
 How these were measured:
 
 - Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
-  `X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
+  `X-Stage-Durations` response header of `/v1/videos/sync`.
-  `MiniMaxH3Pipeline.encode_prompt` (text encode),
-  `MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
-  (VAE decode).
-- MP4 mux is not a header field. It is the
-  `Video response encoding (MP4 bytes)` line in the server log.
 - Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
   high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
   the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
   higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
-  memory has no discrete VRAM to report and the field returns `[N/A]`. See the
+  memory has no discrete VRAM to report and the field returns `[N/A]`.
-  capacity section for a `/proc/meminfo` cross-check.
+- Text encode and VAE decode were measured on a separate 10-step run at the same
-- For T2VA, text encode and VAE decode were measured on a separate 10-step run at
+  960x576 / 8 s geometry. Both are independent of step count, so they transfer to
-  the same 960x576 / 8 s geometry. Both are independent of step count, so they
+  the 50-step row unchanged.
-  transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
-  same way, from the 10-step run at identical output geometry and duration.
-- Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
-  (`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
-  geometry and reference input rather than step count.
 - Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
   engine bookkeeping sit outside the profiled stages.
 
 Denoise dominates at 96% of end-to-end, so step count and geometry are the only
 knobs that materially change the wall clock on this platform.
 
-Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
+Repeated identical 10-step requests varied by up to 23% in denoise time
-4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
-versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
-The cost is in denoise. A single reference image expands the Qwen3-VL
-presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
-6883 tokens, 1 reference images`), and those tokens sit in the attention context
-of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
-1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
-expected for an identical output shape.
-
-Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
 (397 s to 490 s) on a machine that had already been under sustained load. Treat
 these as single-run measurements on one machine, not a throughput guarantee, and
 expect the slower end of that range for back-to-back jobs.
 
-Ref2VA per-step time is load-dependent in the same way. On an idle machine the
-first 10-step request ran at 66.55 s per step. The 50-step request that followed
-occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
-further 10-step requests taken immediately afterwards settled at 74.88 s and
-74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
-is a stable thermal state rather than measurement noise.
-
-Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
-only reachable on a cold machine and does not hold for the length of one request.
-
-Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
-8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
-VAE decode was stable at 67.9-69.8 s throughout.
-
 Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
-default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
+default 1800 s is shorter than a 50-step run on this platform.
-shorter than a 50-step Ref2VA run.
 
 ## Known limitations
 
@@ -246,17 +139,11 @@ shorter than a 50-step Ref2VA run.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
   `height` are both supplied they take precedence and `aspect_ratio` only has to
   name one of the supported ratios.
-- Sustained load degrades throughput. See the variance notes above.
+- Sustained load degrades throughput. See the variance note above.
 - Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
-  with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
+  the same command with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video
-  above; reference videos and additional images push the token count higher and
+  count and prompt length increase activation memory; begin with one request at
-  were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
+  a time and keep the short edge at 576.
-  single-image request, there is very little room to grow — re-measure before
-  adding references or raising the output shape.
-- Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
-  under **Validated on**: FP8 weight loading and image-only Ref2VA are both
-  broken on `release/v0.26.0`.suggestion is  `v0.26.1`.
 - Online FP8 is incompatible with layerwise offload — the offload path produces a
   weight stride the Cutlass FP8 kernel rejects. This is not a practical
   restriction here since offload is unusable on GB10 anyway.
-```

tests/config/test_omni_config.py

@@ -579,7 +579,6 @@ 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
 
 

tests/diffusion/attention/test_flash_attn.py

@@ -9,10 +9,6 @@ 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
 
@@ -375,264 +371,6 @@ 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)

tests/diffusion/attention/test_rainfusion_plan.py

@@ -5,9 +5,10 @@
 
 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 do not. Both are
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
-handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
+first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
-to the always-kept prefix segment before generating the block mask.
+mask from the video and prefix segments separately, so the two only agree on the
+block count when the video segment ends on a boundary.
 """
 
 import dataclasses
@@ -56,13 +57,10 @@ def test_block_aligned_video_segment_runs_sparse():
     assert plan.latent_shape == list(ALIGNED_GRID)
 
 
-def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
+def test_misaligned_video_segment_stays_dense():
-    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
+    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
-
+    # pad keys would take a share of every real query's softmax denominator.
-    assert plan is not None
+    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is 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])
@@ -157,28 +155,3 @@ 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"

tests/diffusion/cache/test_teacache_extractors.py

@@ -488,14 +488,6 @@ 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__()

tests/diffusion/distributed/test_qwenimage_spatial_shard.py

@@ -1,233 +0,0 @@
-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

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -11,8 +10,6 @@ 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
 
@@ -20,23 +17,6 @@ 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():
@@ -275,14 +255,12 @@ 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)
-        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)
-        last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -292,15 +270,13 @@ 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())
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
-    with _active_spatial_shard("width"):
+    mid = mid_rank_pad(x)
-        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())
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
-    with _active_spatial_shard("width"):
+    last = last_rank_pad(x)
-        last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -323,113 +299,6 @@ 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):
@@ -575,99 +444,3 @@ 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

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -492,10 +492,6 @@ 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
 
@@ -530,48 +526,6 @@ 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,
@@ -580,10 +534,6 @@ 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
 

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -69,7 +69,6 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
 
     attention = object.__new__(MiniMaxH3Attention)
     nn.Module.__init__(attention)
-    attention.rot_dim = 96
     attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
     attention.rope._forward_method = attention.rope.forward_native
 
@@ -120,34 +119,3 @@ 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

tests/diffusion/test_diffusion_model_runner.py

@@ -728,19 +728,6 @@ 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
@@ -788,22 +775,16 @@ 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, memory_pool_context_fn=_memory_pool_context)
+    DiffusionModelRunner.load_model(runner)
 
     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

tests/docs/test_generate_examples.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+import yaml
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -46,9 +47,111 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
         generate_examples.load_model_display_names()
 
 
-def test_only_general_serving_examples_get_generated_pages():
+def test_quantization_examples_keep_source_urls_under_feature_navigation(tmp_path, monkeypatch):
-    general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
+    nav_file = tmp_path / "docs" / ".nav.yml"
-    model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))
+    nav_file.parent.mkdir(parents=True)
+    nav_file.write_text(
+        yaml.safe_dump(
+            {
+                "nav": [
+                    {
+                        "User Guide": [
+                            {
+                                "Examples": [
+                                    "examples/README.md",
+                                    {"Offline Inference": [{"Old": "user_guide/examples/offline_inference/old.md"}]},
+                                ]
+                            },
+                            {"Features": [{"Quantization": [{"Overview": "user_guide/quantization/overview.md"}]}]},
+                        ]
+                    }
+                ]
+            },
+            sort_keys=False,
+        ),
+        encoding="utf-8",
+    )
 
-    assert generate_examples.is_general_example(general_example)
+    monkeypatch.setattr(generate_examples, "ROOT_DIR", tmp_path)
-    assert not generate_examples.is_general_example(model_example)
+    monkeypatch.setattr(generate_examples, "NAV_FILE", nav_file)
+    monkeypatch.setattr(generate_examples, "EXAMPLE_DOC_DIR", tmp_path / "docs/user_guide/examples")
+
+    quantization_example = SimpleNamespace(
+        category="quantization",
+        path=tmp_path / "examples/quantization/check_modelopt_fp8_export.py",
+        nav_title="Check Modelopt FP8 Export",
+    )
+    offline_example = SimpleNamespace(
+        category="offline_inference",
+        path=tmp_path / "examples/offline_inference/demo.py",
+        nav_title="Demo",
+    )
+
+    generate_examples.update_nav_file([quantization_example, offline_example])
+    generate_examples.update_nav_file([quantization_example, offline_example])
+
+    nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
+    user_guide = nav_data["nav"][0]["User Guide"]
+    examples = next(item["Examples"] for item in user_guide if "Examples" in item)
+    features = next(item["Features"] for item in user_guide if "Features" in item)
+    quantization = next(item["Quantization"] for item in features if "Quantization" in item)
+
+    assert not any(isinstance(item, dict) and "Quantization" in item for item in examples)
+    generated_examples = [
+        item[generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE]
+        for item in quantization
+        if generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE in item
+    ]
+    assert len(generated_examples) == 1
+    assert generated_examples[0] == [
+        {"Check Modelopt FP8 Export": "user_guide/examples/quantization/check_modelopt_fp8_export.md"}
+    ]
+
+    # A later build without quantization examples must not leave a stale
+    # generated group in the hand-authored Quantization section.
+    generate_examples.update_nav_file([offline_example])
+    nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
+    user_guide = nav_data["nav"][0]["User Guide"]
+    features = next(item["Features"] for item in user_guide if "Features" in item)
+    quantization = next(item["Quantization"] for item in features if "Quantization" in item)
+    assert not any(
+        isinstance(item, dict) and generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE in item
+        for item in quantization
+    )
+
+
+def test_quantization_examples_are_not_added_to_examples_without_feature_target(tmp_path, monkeypatch, caplog):
+    nav_file = tmp_path / "docs" / ".nav.yml"
+    nav_file.parent.mkdir(parents=True)
+    nav_file.write_text(
+        yaml.safe_dump(
+            {
+                "nav": [
+                    {
+                        "User Guide": [
+                            {"Examples": ["examples/README.md"]},
+                        ]
+                    }
+                ]
+            },
+            sort_keys=False,
+        ),
+        encoding="utf-8",
+    )
+
+    monkeypatch.setattr(generate_examples, "ROOT_DIR", tmp_path)
+    monkeypatch.setattr(generate_examples, "NAV_FILE", nav_file)
+    monkeypatch.setattr(generate_examples, "EXAMPLE_DOC_DIR", tmp_path / "docs/user_guide/examples")
+
+    quantization_example = SimpleNamespace(
+        category="quantization",
+        path=tmp_path / "examples/quantization/check_modelopt_fp8_export.py",
+        nav_title="Check Modelopt FP8 Export",
+    )
+
+    generate_examples.update_nav_file([quantization_example])
+
+    nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
+    examples = nav_data["nav"][0]["User Guide"][0]["Examples"]
+    assert not any(isinstance(item, dict) and "Quantization" in item for item in examples)
+    assert "omitting generated quantization examples from navigation" in caplog.text

tests/entrypoints/test_serve.py

@@ -76,18 +76,6 @@ 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",

tools/pre_commit/check_tts_adapter.py

@@ -123,29 +123,18 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
     return found
 
 
-def _check(
+def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
-    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:
-        # Progress is never a build failure. Removing branches is the goal, so a
+        return [
-        # PR that does it must not hand-edit this file to stay green -- nor break
+            f"{path}: {actual} {label}, down from a budget of {budget}. "
-        # `main` for everyone else when it merges from a stale base.
+            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
-        notices.append(
+            f"so the ratchet holds the new ground.",
-            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 []
 
 
@@ -163,7 +152,6 @@ 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()))
 
@@ -174,7 +162,6 @@ 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",
@@ -183,11 +170,8 @@ 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

vllm_omni/config/omni_config.py

@@ -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 = "auto"
+    vae_parallel_mode: str = "tile"
     use_hsdp: bool = False
     mask_sp_padding: bool = False
     hsdp_shard_size: int = -1
@@ -466,10 +466,9 @@ 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 {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of "
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
-                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/diffusion/attention/backends/abstract.py

@@ -20,17 +20,6 @@ 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.
@@ -132,17 +121,6 @@ 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.

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 from functools import partial
 
 import torch
@@ -17,7 +16,6 @@ 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__)
 
@@ -26,15 +24,6 @@ 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
@@ -426,36 +415,7 @@ 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]
@@ -473,161 +433,3 @@ 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,
-            )
-        )

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -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 — delegates to FlashAttention,
+    segments too short to pay for block selection or not aligned to the kernel's
-    so a model can select this backend unconditionally. MindIE-SD handles an
+    block size — delegates to FlashAttention, so a model can select this backend
-    irregular video tail internally, retaining it outside the sparse blocks.
+    unconditionally.
     """
 
     def __init__(
@@ -292,6 +292,25 @@ 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 "

vllm_omni/diffusion/data.py

@@ -206,13 +206,10 @@ class DiffusionParallelConfig:
     process group.
     """
 
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
     """VAE parallel decode strategy.
 
-    - "auto": Select the longer spatial-shard axis for validated Wan and
+    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
-      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
@@ -253,9 +250,8 @@ 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 {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
+        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of "
+            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
-            "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
         if self.allgather_degree > 1:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -4,23 +4,16 @@
 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__)
 
@@ -111,39 +104,10 @@ 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,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,25 +18,11 @@ 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:
@@ -284,30 +270,45 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    def _spatial_shard_decode_split_dim(
+    def _spatial_shard_decode_split_dim(self) -> str | None:
-        self,
+        mode = self.distributed_executor.parallel_mode
-        z: torch.Tensor | None = None,
+        if mode == "spatial_shard_width":
-        world_size: int | None = None,
+            return "width"
-    ) -> str | None:
+        if mode == "spatial_shard_height":
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
+            return "height"
+        return None
 
     def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="Wan")
+        split_dim = self._spatial_shard_decode_split_dim()
+        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
 
-    def _prepare_spatial_shard_decode(self) -> None:
+        group = self.distributed_executor.group
-        prepare_spatial_shard_decode(
+        world_size = dist.get_world_size(group=group)
-            self,
+        requested_size = int(self.distributed_executor.parallel_size)
-            install=wan_spatial_shard.install_wan_spatial_shard_decode,
+        if requested_size != world_size:
-        )
+            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):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
+            split_dim = self._spatial_shard_decode_split_dim()
-            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,

vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py

@@ -1,86 +0,0 @@
-# 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,
-    )

vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py

@@ -1,142 +0,0 @@
-# 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()

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -7,21 +7,19 @@
 #   (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 supports the causal-conv ``feat_cache`` handling used by Wan and
+# width and adds the Wan causal-conv ``feat_cache`` handling.
-# QwenImage.
+"""Spatially-sharded Wan VAE decode.
-"""Shared spatial-shard kernels plus the Wan VAE adapter.
 
-The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
+The existing distributed Wan VAE path shards *tiles*.  This module adds an
-are also reused by QwenImage's thin adapter. The backend shards decoder feature
+opt-in decode backend that shards decoder feature maps along height or width and
-maps along height or width and exchanges boundary rows/columns before spatial
+exchanges boundary rows/columns before spatial convolutions.  It is
-convolutions. It is decode-only and keeps checkpoint loading unchanged by
+intentionally decode-only and keeps checkpoint loading unchanged by patching the
-patching the already-loaded decoder.
+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
@@ -36,8 +34,6 @@ 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__)
 
 
@@ -56,17 +52,12 @@ _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 VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported Wan 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:
@@ -164,9 +155,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"VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
     if not 0 <= rank < world_size:
-        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"Wan 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)
@@ -332,27 +323,21 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        defer_split_after_padding: bool = False,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ) -> None:
         super().__init__()
         self.padding = padding
-        self.defer_split_after_padding = defer_split_after_padding
+        self.split_dim = split_dim
+        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 split_dim == "height":
+            if self.split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -366,7 +351,8 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        deferred_padding: tuple[int, int, int, int] | None = None,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -383,253 +369,176 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self._direct_padding = source.padding
+        self.split_dim = split_dim
-        self._direct_padding_mode = source.padding_mode
+        self.split_tensor_dim = _spatial_dim(split_dim)
-        self._direct_reversed_padding = source._reversed_padding_repeated_twice
+        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
-        self._deferred_padding = deferred_padding
+        self.halo_size = (kernel_extent - 1) // 2
-        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
+        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
-        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
+        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
-            raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
+        if split_dim == "height":
-        self._spatial_pad_h = int(source_padding[-2])
+            if split_padding is None:
-        self._spatial_pad_w = int(source_padding[-1])
+                split_padding = (pad_h, pad_h)
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
+            self._non_split_padding = (pad_w, pad_w, 0, 0)
-        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]
+        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:
-        split_dim = _active_spatial_shard_split_dim()
+        x = F.pad(x, self._non_split_padding)
-        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=halo_size,
+            halo_size=self.halo_size,
-            split_dim=split_dim,
+            split_dim=self.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(
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
-            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,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        self,
+        trim_params = self._trim_cache.get(local_extent)
-        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=halo_size,
+                halo_size=self.halo_size,
-                pad_before=split_pad_left,
+                pad_before=self.split_pad_left,
-                pad_after=split_pad_right,
+                pad_after=self.split_pad_right,
-                kernel_extent=kernel_extent,
+                kernel_extent=self.kernel_extent,
-                stride_extent=stride_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
-# Model-neutral aliases let other causal VAEs reuse the implementation while
+class WanDistCausalConv3d(nn.Conv3d):
-# preserving the original Wan class identities for stacked PRs and downstreams.
+    def __init__(
-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,
-    ) -> None:
+        split_dim: str = "height",
+    ):
+        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)
-        self._padding = tuple(source_padding)
+        if split_dim == "height":
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
+            self.split_pad_left = int(self._source_padding[2])
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
+            self.split_pad_right = int(self._source_padding[3])
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+            self.kernel_extent = self.kernel_size[-2]
+            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 _spatial_shard_causal_conv_forward(
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        self,
+        padding = list(self._padding)
-        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=halo_size,
+            halo_size=self.halo_size,
-            split_dim=split_dim,
+            split_dim=self.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(
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
-            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,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
-        out = nn.Conv3d.forward(self, x_padded)
+        out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        self,
+        trim_params = self._trim_cache.get(local_extent)
-        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=halo_size,
+                halo_size=self.halo_size,
-                pad_before=split_pad_left,
+                pad_before=self.split_pad_left,
-                pad_after=split_pad_right,
+                pad_after=self.split_pad_right,
-                kernel_extent=kernel_extent,
+                kernel_extent=self.kernel_extent,
-                stride_extent=stride_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = 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,
@@ -673,15 +582,12 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> 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)
@@ -703,43 +609,60 @@ 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__ == causal_conv_class_name:
+    if child.__class__.__name__ == "WanCausalConv3d":
-        setattr(parent, name, causal_conv_factory(child, group))
-        return
-    if isinstance(child, nn.ZeroPad2d):
-        padding = tuple(int(p) for p in child.padding)
         setattr(
             parent,
             name,
-            SpatialShardZeroPad2d(
+            WanDistCausalConv3d(
-                padding,
+                child,
                 group,
-                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
+                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(
+                module_padding,
+                group,
+                split_dim=split_dim,
+                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        deferred_padding = None
+        split_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
-            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
+            # WanResample downsample 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, SpatialShardZeroPad2d):
+            if isinstance(prev, WanDistZeroPad2d):
-                deferred_padding = prev.padding
+                split_padding = prev.split_padding
             elif isinstance(prev, nn.ZeroPad2d):
-                deferred_padding = tuple(int(value) for value in prev.padding)
+                pad = prev.padding
+                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
         setattr(
             parent,
             name,
-            SpatialShardConv2d(
+            WanDistConv2d(
                 child,
                 group,
-                deferred_padding=deferred_padding,
+                split_dim=split_dim,
+                split_padding=split_padding,
             ),
         )
 
@@ -747,45 +670,21 @@ 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__ == attention_block_class_name:
+    if module.__class__.__name__ == "WanAttentionBlock":
-        _patch_attention_block(module, group)
+        _patch_attention_block(module, group, split_dim)
         inside_attention = True
 
     for name, child in list(module.named_children()):
-        if child.__class__.__name__ == causal_conv_class_name:
+        if child.__class__.__name__ == "WanCausalConv3d":
-            _replace_child(
+            _replace_child(module, name, child, group, split_dim)
-                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(
+            _replace_child(module, name, child, group, split_dim)
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
             continue
-        _patch_decoder_modules(
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
-            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:
@@ -796,58 +695,47 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
+def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> 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 context-aware variants and attaching a dedicated
+    convolutions/padding for halo-exchanging variants and wrapping
-    spatial forward. Outside that forward the replacements execute their exact
+    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
-    local PyTorch behavior, so a later auto-selected tile request remains safe.
+    instance and is applied only once (subsequent calls are no-ops). A given
-    The same installed decoder can alternate height and width sharding.
+    instance is bound to a single ``split_dim``; switching between
+    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
+    otherwise.
 
-    Group-relative rank 0 assembles the final decoded frame. The model adapter
+    Only group-relative rank 0 assembles the final decoded frame, mirroring the
-    then preserves its prior contract: Wan keeps the result on rank 0, while
+    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
-    QwenImage broadcasts it to every rank.
+    take part in the collectives but return an empty placeholder.
     """
     _spatial_dim(split_dim)
-    if getattr(vae, installed_attr, False):
+    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", 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(f"{model_name} spatial-shard VAE decode requires a decoder module.")
+        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(
+    _patch_decoder_modules(decoder, group, split_dim)
-        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(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
+    def _forward(
+        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(
@@ -867,50 +755,29 @@ def install_spatial_shard_decode(
             )
         )
         try:
-            out = orig_forward(x, *args, **kwargs)
+            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
         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._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
+    decoder.forward = MethodType(_forward, decoder)
-    setattr(vae, installed_attr, True)
+    vae._vllm_omni_wan_spatial_shard_installed = True
-    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
+    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 
-def install_wan_spatial_shard_decode(
+def spatial_shard_decode(
-    vae: SpatialShardVAE,
+    vae: Any,
-    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(vae, group, split_dim)
+    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
 
     if z.shape[2] == 0:
-        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError("Wan 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.
@@ -919,72 +786,31 @@ def spatial_shard_decode_impl(
 
     vae.clear_cache()
     try:
-        context_factory = getattr(vae, "_execution_context", None)
+        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
-        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]
-                if pass_first_chunk:
+                chunk = vae.decoder(
-                    chunk = spatial_forward(
+                    x[:, :, i : i + 1, :, :],
-                        x[:, :, i : i + 1, :, :],
+                    feat_cache=vae._feat_map,
-                        feat_cache=vae._feat_map,
+                    feat_idx=vae._conv_idx,
-                        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 unpatchify_patch_size is not None:
+                if vae.config.patch_size is not None:
-                    out = unpatchify(out, patch_size=unpatchify_patch_size)
+                    out = unpatchify(out, patch_size=vae.config.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,
-    )

vllm_omni/diffusion/layers/norm.py

@@ -68,7 +68,6 @@ class RMSNorm(CustomOp):
     def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
-        self.hidden_size = hidden_size
         self.variance_epsilon = eps
 
     def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
@@ -108,19 +107,6 @@ class RMSNorm(CustomOp):
         except Exception:
             return self.forward_native(x)
 
-    def forward_musa(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
-        # H3 Q/K norm with its inline RoPE path on MUSA.
-        return F.rms_norm(
-            x,
-            (self.hidden_size,),
-            self.weight,
-            self.variance_epsilon,
-        )
-
     def forward_npu(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/layers/rope.py

@@ -193,16 +193,7 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        # Keep the full-dimension NeoX form inline for MUSA compilation.
+        return self.forward_native(x, cos, sin)
-        # Other layouts use the shared native implementation.
-        if self.half_head_dim or self.interleaved:
-            return self.forward_native(x, cos, sin)
-        if cos.dim() == 3:
-            cos, sin = cos[0], sin[0]
-        cos = cos.unsqueeze(-2)
-        sin = sin.unsqueeze(-2)
-        x1, x2 = x.chunk(2, dim=-1)
-        return x * cos + torch.cat((-x2, x1), dim=-1) * sin
 
     def forward_native(
         self,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -107,15 +107,6 @@ 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:
@@ -342,7 +333,6 @@ class MiniMaxH3Attention(nn.Module):
         )
         self.num_heads = self.qkv_proj.num_heads
         self.num_kv_heads = self.qkv_proj.num_kv_heads
-        self.rot_dim = 6 * arch.rope_inv_freq_len
         self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
@@ -375,7 +365,7 @@ class MiniMaxH3Attention(nn.Module):
         x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
         are cast to the activation dtype before the elementwise math.
         """
-        rot_dim = self.rot_dim
+        rot_dim = freqs.shape[-1]
         x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
         cos = torch.cos(freqs).to(x.dtype)  # [T, rot_dim]
         sin = torch.sin(freqs).to(x.dtype)
@@ -411,16 +401,11 @@ 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 the no-mask fast paths are local-only.
+        # configured backend, so this no-mask fast path is local-only.
-        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
+        prefix_slice = (
-        # supports_packed_mask_free: backend consumes the packed metadata
+            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
-        # 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 no_mask:
+        if used < packed_total and not prefix_slice:
             attn_mask = torch.arange(packed_total, device=q.device)[None] < used
         metadata = AttentionMetadata(
             attn_mask=attn_mask,
@@ -430,15 +415,6 @@ 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,
         )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -155,11 +155,6 @@ 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()

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -30,9 +30,6 @@ 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 (
@@ -246,7 +243,6 @@ 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(

vllm_omni/engine/arg_utils.py

@@ -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 = "auto"
+    vae_parallel_mode: str = "tile"
     text_encoder_tp_size: int = 1
     default_sampling_params: str | None = None
     max_generated_image_size: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -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 "auto"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
             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)

vllm_omni/entrypoints/cli/serve.py

@@ -750,12 +750,10 @@ class OmniServeCommand(CLISubcommand):
         omni_config_group.add_argument(
             "--vae-parallel-mode",
             type=str,
-            default="auto",
+            default="tile",
-            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
+            "'tile' (default) uses patch/tile parallel decode; "
-            "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. "

vllm_omni/platforms/npu/platform.py

@@ -139,18 +139,6 @@ 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()