[Perf][Diffusion] Fuse lossless Wan VAE data movement

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Bugfix] rainfusion support all shape video (#6000 )
2026-08-11 09:43:11 +08:00 · 2026-08-11 05:10:50 +08:00 · 2026-08-11 05:08:31 +08:00 · 2026-08-10 22:20:37 +08:00 · 2026-08-10 21:55:01 +08:00 · 2026-08-10 19:47:03 +08:00
40 changed files with 1930 additions and 258 deletions
@@ -1,6 +1,11 @@
 # CODEOWNERS is used to automatically assign PR reviewers
 # Format: path pattern @username or @team
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
+#
+# The vllm_omni/** source-tree rules below are derived from the `owners` and
+# `primary_code_paths` front matter in docs/design/module/*.md. When updating
+# ownership, update the design doc first, then regenerate the matching rule here.
+# Last-match-wins: more specific rules must come AFTER broader ones.

 # Default owners (fallback reviewers for all files)
 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao

-# Engine and model_executor
-/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
+# ComfyUI / external apps
+/apps/ @fhfuih @alex-jw-brooks

-# 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 (out-of-tree)
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
 /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin

-# Usability and Observability
-/vllm_omni/deploy/ @tzhouam @lishunyang12
-/vllm_omni/metrics/ @lishunyang12
-/vllm_omni/profiler/ @gcanlin
+# =====================================================================
+# Module design ownership (docs/design/module/*.md)
+# Each section header cites the source design doc.
+# =====================================================================

-# ComfyUI
-/apps/ @fhfuih @alex-jw-brooks
+# Autoregressive runtime — docs/design/module/ar_runtime.md
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
+/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123

-# Large Scale Serving
+# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
+/vllm_omni/engine/ @tzhouam @fake0fan
+# File-level overrides: input/output modality contracts own these engine files
+# — docs/design/module/input_output_modality_contracts.md
+/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+
+# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
+/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
+
+# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
+/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
+# File-level overrides within entrypoints/
+/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
+/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
+
+# Error contracts — docs/design/module/error_contracts.md
+/vllm_omni/errors.py @alex-jw-brooks @NickCao
+
+# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan

-# Quantization and configuration
+# Model integration — docs/design/module/model_integration.md
+/vllm_omni/model_executor/ @tzhouam @gcanlin
+/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
+/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
+/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
+# Config override within model_executor — docs/design/module/vllm_omni_config.md
+/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
+
+# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
+/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
+
+# Configuration & deploy — docs/design/module/vllm_omni_config.md
 /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
+/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
+
+# Quantization — docs/design/module/quantization.md
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12

-# Hardware plugins and accelerators
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+# 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
+/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
+
+# LoRA (no dedicated design doc; ownership intentionally narrowed)
+/vllm_omni/lora/ @SamitHuang @wtomin
+
+# Transformers utils (related path only in design docs; preserves prior ownership)
+/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
+
+# =====================================================================
+# Diffusion module family — docs/design/module/diffusion/
+# Umbrella rule first; sub-module rules after it win via last-match.
+# =====================================================================
+
+# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
+# — docs/design/module/diffusion/index.md
+/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
+
+# Cache management — docs/design/module/cache_management.md
+/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
+
+# Diffusion quantization — docs/design/module/quantization.md
+/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
+
+# Diffusion profiling — docs/design/module/profiling.md
+/vllm_omni/diffusion/profiler/ @gcanlin
+
+# Diffusion runtime (scheduler, executor, worker)
+# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
+/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
+/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
+/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
+
+# Diffusion model integration (pipelines, registration, checkpoint loading)
+# — docs/design/module/diffusion/diffusion_model_integration.md
+/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
+/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
+/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
+
+# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
+# — docs/design/module/diffusion/parallelism.md
+/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
+/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
+
+# Diffusion offloader (CPU offloading, DLO)
+# — docs/design/module/diffusion/offloader.md
+/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
@@ -66,7 +66,10 @@ repos:
      entry: python tools/pre_commit/check_tts_adapter.py
      language: python
      pass_filenames: false
-      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
+      # The budget-is-now-slack reminder is printed on a PASSING run, and
+      # pre-commit hides output from passing hooks unless the hook is verbose.
+      verbose: true

    # Keep `suggestion` last
    - id: suggestion
@@ -16,31 +16,21 @@ 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
-      - 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
+      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.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
-      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.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
+      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
  - General:
    - usage/*
@@ -7,6 +7,7 @@ owners:
  - "@yinpeiqi"
  - "@fake0fan"
  - "@Sy0307"
+  - "@Gaohan123"
 primary_code_paths:
  - vllm_omni/core/**
  - vllm_omni/worker/**
@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
  - "@alex-jw-brooks"
+  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/benchmarks/**
  - vllm_omni/entrypoints/cli/benchmark/**
@@ -5,6 +5,7 @@ status: draft
 owners:
  - "@Isotr0py"
  - "@princepride"
+  - "@SamitHuang"
 primary_code_paths:
  - vllm_omni/diffusion/cache/**
 related_code_paths:
@@ -3,9 +3,10 @@ title: Diffusion Model Integration
 kind: module
 status: draft
 owners:
-  - "@Isotr0py"
+  - "@fhfuih"
+  - "@Bounty-hunter"
  - "@wtomin"
-  - "@ZJY0516"
+  - "@RuixiangMa"
 primary_code_paths:
  - vllm_omni/diffusion/models/**
  - vllm_omni/diffusion/model_loader/**
@@ -6,6 +6,7 @@ owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
+  - "@fhfuih"
 primary_code_paths:
  - vllm_omni/diffusion/sched/**
  - vllm_omni/diffusion/executor/**
@@ -4,7 +4,8 @@ kind: module
 status: draft
 owners:
  - "@yuanheng-zhao"
-  - "@Isotr0py"
+  - "@david6666666"
+  - "@lishunyang12"
 primary_code_paths:
  - vllm_omni/diffusion/offloader/**
 related_code_paths:
@@ -6,6 +6,7 @@ owners:
  - "@princepride"
  - "@yuanheng-zhao"
  - "@xuechendi"
+  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/diffusion/distributed/**
  - vllm_omni/diffusion/attention/parallel/**
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
 owners:
  - "@alex-jw-brooks"
  - "@linyueqian"
+  - "@NickCao"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -8,8 +8,15 @@ owners:
  - "@xuechendi"
 primary_code_paths:
  - vllm_omni/platforms/**
-related_code_paths:
  - vllm_omni/attention/**
+primary_code_path_owners:
+  # Additional owners scoped to attention kernels (quantized/custom attention)
+  - paths:
+      - vllm_omni/attention/**
+    owners:
+      - "@david6666666"
+      - "@lishunyang12"
+related_code_paths: []
 depends_on:
  - ar_runtime.md
  - diffusion/index.md
@@ -6,6 +6,8 @@ architecture_state: current
 owners:
  - "@Sy0307"
  - "@amy-why-3459"
+  - "@Gaohan123"
+  - "@alex-jw-brooks"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -9,6 +9,14 @@ primary_code_paths:
  - vllm_omni/model_executor/**
  - vllm_omni/model_extras/**
  - vllm_omni/plugins/**
+primary_code_path_owners:
+  # Additional owners scoped to a sub-path of model_executor/**
+  - paths:
+      - vllm_omni/model_executor/models/**
+    owners:
+      - "@Sy0307"
+      - "@amy-why-3459"
+      - "@linyueqian"
 related_code_paths:
  - vllm_omni/transformers_utils/**
  - vllm_omni/tokenizers/**
@@ -4,6 +4,7 @@ kind: module
 status: draft
 owners:
  - "@lishunyang12"
+  - "@vraiti"
 primary_code_paths:
  - vllm_omni/metrics/**
 related_code_paths:
@@ -24,15 +24,22 @@ SERVING_MODE_TITLES = {
    "online_serving": "Online serving",
 }

-# Generated example pages keep their source-derived URLs even when their
-# navigation category changes. Quantization helpers are feature documentation,
-# not user-serving examples, so keep them under the existing Features /
-# Quantization section instead of creating a second top-level Quantization
-# category under Examples.
-EXAMPLE_NAV_TARGETS = {
-    "quantization": "quantization",
-}
-GENERATED_QUANTIZATION_EXAMPLES_TITLE = "Generated Examples"
+# Keep the Examples navigation focused on reusable modality/task entry points.
+# Model-specific scripts remain available under ``examples/`` and are surfaced
+# from these shared pages instead of getting one page per model.
+GENERAL_EXAMPLE_SLUGS = frozenset(
+    {
+        "image_to_image",
+        "image_to_video",
+        "speech_to_video",
+        "text_to_audio",
+        "text_to_image",
+        "text_to_speech",
+        "text_to_video",
+        "x_to_text",
+        "x_to_video_audio",
+    }
+)


 def load_model_display_names() -> dict[str, str]:
@@ -56,6 +63,13 @@ 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",
@@ -358,12 +372,9 @@ 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. Categories
-    # not explicitly routed elsewhere retain the historical Examples behavior.
+    # Build the new Examples section - start with preserved items.
    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())
@@ -372,36 +383,6 @@ 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:
@@ -432,6 +413,8 @@ 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"
@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).

 ### Which geometries run sparse

-Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
-is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
-tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
-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.
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
+promoted rows remain visible to every query and key.

-Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
-spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
-patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
-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.
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
+and kernel tiling aligned, at the cost of retaining more blocks.

-"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
-resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
-which you can buy back with `start_step`.
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+that implements protected video tails for rf_v2.

 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)

@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends

 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).

 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 ### Multi-NPU: 768P combined-service configuration

 Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
-degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
+degree 8, native tiled VAE patch parallelism degree 8, and distributed
+layerwise offload:

 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
  --usp 8 \
  --ring 1 \
  --text-encoder-tp-size 8 \
-  --enable-layerwise-offload \
+  --enable-distributed-layerwise-offload \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --vae-patch-parallel-size 8 \
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
 partition restart is required. Layerwise offload applies to both DiTs.

+On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
+768P without offloading or sharding: use distributed layerwise offload (as
+above) or HSDP — see
+[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
+
 ### Optional optimizations

 Two independent optimizations may be enabled on top of the configuration
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
 over the whole sequence, so ring parallelism would split away the keys it
 needs. Scale with `--usp` instead.

+## Memory and attention optimizations (A3)
+
+### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
+
+Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
+combined MiniMax-H3 service at 768P. Enable **one** of the following at
+server startup:
+
+**Distributed layerwise offload** — DiT layers are offloaded with parameters
+gathered across the parallel group instead of replicated per rank:
+
+```bash
+  --enable-distributed-layerwise-offload
+```
+
+Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
+input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
+usage is high with this option.
+
+**HSDP** — hybrid sharded data parallelism for the DiT parameters. The
+multi-stream memory-reuse knob is mandatory with this flag:
+
+```bash
+export MULTI_STREAM_MEMORY_REUSE=2
+```
+
+```bash
+  --use-hsdp
+```
+
+Host memory usage is small, but large shapes may still OOM; HBM usage
+optimizations for this path are ongoing.
+
+### FLASH_ATTN backend with MindIE-SD
+
+Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
+[§ Environment](#environment)). On NPU this backend carries most of the
+memory-reduction and performance work: a mask-free packed varlen path that
+never materializes the quadratic `full_qk` padding mask, and K/V prefix
+slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
+transformer.
+
+### Optional: LaserAttention fused kernel
+
+For an additional attention speedup, select the Ascend LaserAttention fused
+kernel before starting the server:
+
+```bash
+export MINDIE_SD_FA_TYPE="ascend_laser_attention"
+```
+
+This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
+exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
+kernel's fp16 workspace cannot overflow on outlier activations. Measured
+speedup numbers will be added here.
+
+
 ## HTTP API examples

 The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:

 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |

 These measurements describe the validated shapes rather than a general
 throughput guarantee.
@@ -9,6 +9,10 @@ This script tests two main scenarios:
 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
 """

+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 import torch

@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
        impl.forward_cuda(query, query, query, metadata)


+# ---------------------------------------------------------------------------
+# NPU mask-free packed attention (PR #5891).
+#
+# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
+#   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
+#     sync, no MindIE-SD import);
+#   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
+#   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
+#
+# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
+# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
+# ``benchmarks/conftest``). No production code is exercised for real kernels.
+
+
+def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
+    return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
+
+
+def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
+    """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
+
+    Restored automatically by monkeypatch. Defaults are no-op callables so the
+    top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
+    instance methods that would otherwise call them.
+    """
+    fake = SimpleNamespace(
+        attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
+        attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
+    )
+    monkeypatch.setitem(sys.modules, "mindiesd", fake)
+    return fake
+
+
+# --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
+
+
+def test_resolve_packed_seq_accepts_real_plus_pad_contract():
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
+
+
+def test_resolve_packed_seq_single_document_no_padding():
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "causal",
+        "missing_key",
+        "batch_ne_1",
+        "cu_too_long",
+        "cu_qk_mismatch",
+        "max_seqlen_not_int",
+        "used_gt_total",
+        "used_zero",
+        "pad_longer_than_real",
+    ],
+)
+def test_resolve_packed_seq_rejects_malformed_metadata(case):
+    """Any metadata outside the exact [real, pad] contract returns None so the
+    caller falls back to the masked path."""
+    cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
+    cu2 = torch.tensor([0, 8], dtype=torch.int32)
+    cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
+    cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
+    base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
+    q = torch.randn(1, 8, 2, 4)
+
+    impl = _npu_impl(causal=(case == "causal"))
+    if case == "missing_key":
+        extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5}  # 3 of 4 keys absent
+    elif case == "batch_ne_1":
+        q = torch.randn(2, 8, 2, 4)
+        extra = base
+    elif case == "cu_too_long":
+        extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
+    elif case == "cu_qk_mismatch":
+        extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
+    elif case == "max_seqlen_not_int":
+        extra = {**base, "max_seqlen_q": 5.0}
+    elif case == "used_gt_total":
+        extra = {**base, "max_seqlen_q": 9}  # 9 > total length 8
+    elif case == "used_zero":
+        extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
+    elif case == "pad_longer_than_real":
+        # real=3, pad=5: the real document must be the longer one.
+        extra = {
+            "cu_seqlens_q": cu_short,
+            "cu_seqlens_k": cu_short,
+            "max_seqlen_q": 3,
+            "max_seqlen_k": 3,
+        }
+    else:
+        extra = base
+
+    assert impl._resolve_packed_seq_npu(q, q, extra) is None
+
+
+# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
+
+
+def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
+    """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
+
+    Lets a test assert which branch ``forward_fa_npu`` dispatched to without
+    running any real kernel.
+    """
+    impl = _npu_impl()
+    sentinel = torch.tensor([1.0])
+    impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
+    impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
+    return impl, sentinel
+
+
+def test_npu_env_unset_uses_varlen(monkeypatch):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_varlen_packed_npu.assert_called_once()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+
+
+def test_npu_env_laser_uses_slice(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_prefix_kv_slice_npu.assert_called_once()
+    impl._forward_varlen_packed_npu.assert_not_called()
+
+
+@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
+def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
+    """Locks in the CURRENT dispatch behavior (PR #5891).
+
+    Any value other than ``ascend_laser_attention`` — including the valid
+    ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
+    the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
+    comment #1 on the PR and is intentionally captured here as a regression
+    net: update this test when the dispatch starts honoring/rejecting these.
+    """
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
+    _fake_mindiesd(monkeypatch)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+
+    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
+
+    impl._forward_varlen_packed_npu.assert_called_once()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+
+
+def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
+    """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
+    the existing masked attention_forward path; the new branches are skipped."""
+    fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
+    _fake_mindiesd(monkeypatch, attention_forward=fake_forward)
+    impl, _ = _npu_impl_with_mocked_paths()
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
+        extra={},  # no npu_attn_varlen opt-in
+    )
+    q = torch.randn(1, 8, 2, 4)
+
+    out = impl.forward_fa_npu(q, q, q, metadata)
+
+    impl._forward_varlen_packed_npu.assert_not_called()
+    impl._forward_prefix_kv_slice_npu.assert_not_called()
+    fake_forward.assert_called_once()
+    assert out is fake_forward.return_value
+
+
+# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
+
+
+def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured["query"] = query
+        captured["key"] = key
+        captured["value"] = value
+        captured.update(kwargs)
+        return torch.ones_like(query)  # deterministic stand-in kernel output
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    impl = _npu_impl()  # softmax_scale == 0.5
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {
+        "cu_seqlens_q": cu,
+        "cu_seqlens_k": cu,
+        "max_seqlen_q": 5,
+        "max_seqlen_k": 5,
+        "laser_input_scale": 256.0,
+    }
+
+    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
+
+    # K/V sliced to the valid prefix (used_k == 5); q keeps full length.
+    assert captured["key"].shape == (1, 5, 2, 4)
+    assert captured["value"].shape == (1, 5, 2, 4)
+    assert captured["query"].shape == (1, 8, 2, 4)
+    assert captured["attn_mask"] is None
+    # Pre-divide q/k/v by the factor and compensate the kernel scale by its
+    # square (exact for power-of-two); output scaled back by the factor.
+    torch.testing.assert_close(captured["query"], q / 256.0)
+    torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
+    assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
+    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
+
+
+def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured["query"] = query
+        captured.update(kwargs)
+        return torch.ones_like(query)
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    impl = _npu_impl()
+    q = torch.randn(1, 8, 2, 4)
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {
+        "cu_seqlens_q": cu,
+        "cu_seqlens_k": cu,
+        "max_seqlen_q": 5,
+        "max_seqlen_k": 5,
+        # no laser_input_scale → no pre-scaling
+    }
+
+    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
+
+    torch.testing.assert_close(captured["query"], q)
+    assert captured["scale"] == pytest.approx(0.5)
+    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
+
+
 if __name__ == "__main__":
    print("Running FlashAttention Padding Tests...")
    print("=" * 60)
@@ -5,10 +5,9 @@

 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
-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.
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
+to the always-kept prefix segment before generating the block mask.
 """

 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
    assert plan.latent_shape == list(ALIGNED_GRID)


-def test_misaligned_video_segment_stays_dense():
-    # 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 make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 62496
+    assert plan.latent_shape == list(MISALIGNED_GRID)


@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
+    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
+    torch.manual_seed(1)
+    prefix_len, grid = 384, (4, 16, 10)
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    heads, head_dim = 4, 128
+    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
+
+    impl = make_impl()
+    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
+    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
+    out = impl._forward_sparse_npu(q, k, v, plan)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).transpose(1, 2)
+    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
+    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
    class FakeBackend:
        supports_prefix_kv_slicing = False

+        # Mirror AttentionBackend.supports_packed_mask_free: H3's packed
+        # attention path calls attn_backend.supports_packed_mask_free() to decide
+        # whether to build the padding mask. Return False so the fake exercises
+        # the masked path (the abstract backend default).
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
    def __init__(self, **kwargs):
        del kwargs
        super().__init__()
@@ -1,3 +1,5 @@
+from types import SimpleNamespace
+
 import pytest
 import torch

@@ -41,3 +43,21 @@ def test_wan_vae_execution_context_uses_platform_autocast(mocker):
        dtype=torch.bfloat16,
        enabled=True,
    )
+
+
+def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch):
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    installed = []
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
+        lambda pipeline: SimpleNamespace(vaes=[model]),
+    )
+    monkeypatch.setattr(
+        wan_vae_module,
+        "install_wan_vae_data_movement",
+        lambda vae: installed.append(vae),
+    )
+
+    wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity())
+
+    assert installed == [model]
@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders import AutoencoderKLWan
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_vae_data_movement
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
+    _cache_payload,
+    _causal_padding,
+    _is_dynamic_spatial_conv,
+    _run_cached_causal_conv,
+    install_wan_vae_data_movement,
+)
+from vllm_omni.diffusion.kernels.wan_vae_data_movement import cat_pad_5d, dup_up3d_add
+
+pytestmark = pytest.mark.core_model
+cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _dense_5d(shape: tuple[int, ...], dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
+    tensor = torch.randn(shape, device="cuda", dtype=dtype)
+    if channels_last:
+        return tensor.contiguous(memory_format=torch.channels_last_3d)
+    return tensor.contiguous()
+
+
+def _reference_cat_pad(
+    x: torch.Tensor,
+    cache: torch.Tensor | None,
+    padding: tuple[int, ...],
+) -> torch.Tensor:
+    reference_padding = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        reference_padding[4] -= cache.shape[2]
+    return F.pad(x, reference_padding)
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,padding",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+    ],
+)
+def test_cat_pad_5d_is_bitwise_and_layout_exact(
+    dtype: torch.dtype,
+    channels_last: bool,
+    channels: int,
+    frames: int,
+    height: int,
+    width: int,
+    cache_frames: int,
+    padding: tuple[int, ...],
+) -> None:
+    torch.cuda.manual_seed(0)
+    x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = None
+    if cache_frames:
+        pad_height, pad_width = padding[2], padding[0]
+        buffer = _dense_5d(
+            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width),
+            dtype,
+            channels_last,
+        )
+        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
+
+    reference = _reference_cat_pad(x, cache, padding)
+    output = cat_pad_5d(x, cache, padding)
+    assert output is not None
+    assert output.shape == reference.shape
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+    pair = cat_pad_5d(x, cache, padding, keep_cache_frames=2)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert torch.equal(output_with_cache, reference)
+    pad_height, pad_width = padding[2], padding[0]
+    keep_frames = min(2, reference.shape[2])
+    expected_cache = reference[
+        :,
+        :,
+        -keep_frames:,
+        pad_height : pad_height + height,
+        pad_width : pad_width + width,
+    ]
+    assert next_cache.shape == expected_cache.shape
+    assert torch.equal(next_cache, expected_cache)
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "permuted"])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
+    [
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),
+        (64, 32, 2, 10, 14, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_is_bitwise_and_layout_exact(
+    dtype: torch.dtype,
+    main_layout: str,
+    in_channels: int,
+    out_channels: int,
+    frames: int,
+    height: int,
+    width: int,
+    factor_temporal: int,
+    factor_spatial: int,
+    drop_first: bool,
+) -> None:
+    torch.cuda.manual_seed(0)
+    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
+    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
+    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
+    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
+    if main_layout == "channels_last":
+        main = _dense_5d(shape, dtype, True)
+    elif main_layout == "permuted":
+        main = torch.randn(
+            (shape[0], shape[2], shape[1], shape[3], shape[4]),
+            device="cuda",
+            dtype=dtype,
+        ).permute(0, 2, 1, 3, 4)
+    else:
+        main = _dense_5d(shape, dtype, False)
+
+    shortcut = source.repeat_interleave(repeats, dim=1)
+    shortcut = shortcut.view(
+        1,
+        out_channels,
+        factor_temporal,
+        factor_spatial,
+        factor_spatial,
+        frames,
+        height,
+        width,
+    )
+    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
+    if drop_first:
+        shortcut = shortcut[:, :, factor_temporal - 1 :]
+    reference = main + shortcut
+
+    output = dup_up3d_add(
+        main,
+        source,
+        factor_temporal,
+        factor_spatial,
+        repeats,
+        drop_first,
+    )
+    assert output is not None
+    assert output.shape == reference.shape
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize("temporal_only", [False, True])
+def test_cached_conv_chunk_loop_matches_reference(monkeypatch, channels_last: bool, temporal_only: bool) -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanCausalConv3d
+
+    torch.cuda.manual_seed(0)
+    channels = 64
+    if temporal_only:
+        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(channels, channels, 3, padding=1)
+    conv = conv.to(device="cuda", dtype=torch.float32)
+    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
+
+    def run(force_reference: bool, start):
+        cache = [start]
+        outputs = []
+        original = wan_vae_data_movement.cat_pad_5d
+        if force_reference:
+            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", None)
+        try:
+            for chunk in chunks:
+                outputs.append(_run_cached_causal_conv(conv, chunk, cache, 0))
+        finally:
+            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", original)
+        return outputs, cache[0]
+
+    for start in (None, "Rep"):
+        fused_outputs, fused_cache = run(False, start)
+        reference_outputs, reference_cache = run(True, start)
+        for output, reference in zip(fused_outputs, reference_outputs, strict=True):
+            assert torch.equal(output, reference)
+        payload = _cache_payload(fused_cache)
+        assert payload is not None and payload.shape[2] == CACHE_T
+        assert torch.equal(payload, reference_cache[:, :, -CACHE_T:])
+
+
+@pytest.mark.cpu
+def test_installer_is_idempotent_and_preserves_state_dict_keys() -> None:
+    vae = OmniAutoencoderKLWan(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 1],
+        num_res_blocks=1,
+        temperal_downsample=[False, False],
+        is_residual=True,
+    )
+    keys_before = set(vae.state_dict())
+    assert install_wan_vae_data_movement(vae)
+    assert install_wan_vae_data_movement(vae)
+    assert set(vae.state_dict()) == keys_before
+
+
+@pytest.mark.cpu
+def test_dynamic_spatial_conv_exposes_direct_path_compatibility_contract() -> None:
+    conv = torch.nn.Conv3d(2, 2, 3)
+    conv._vllm_omni_dynamic_spatial_shard_conv = True
+    conv._source_padding = (1, 1, 1, 1, 2, 0)
+
+    assert _is_dynamic_spatial_conv(conv)
+    assert _causal_padding(conv) == conv._source_padding
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+def test_tiny_wan_decoder_is_bitwise_exact(channels_last: bool) -> None:
+    config = dict(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 1],
+        num_res_blocks=1,
+        temperal_downsample=[False, True],
+        is_residual=True,
+    )
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to("cuda")
+    candidate = OmniAutoencoderKLWan(**config).eval().to("cuda")
+    candidate.load_state_dict(reference.state_dict())
+    if channels_last:
+        for model in (reference, candidate):
+            for module in model.decoder.modules():
+                if isinstance(module, torch.nn.Conv3d):
+                    module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
+
+    latents = torch.randn(1, 4, 3, 4, 4, device="cuda")
+    if channels_last:
+        latents = latents.contiguous(memory_format=torch.channels_last_3d)
+    expected = reference.decode(latents, return_dict=False)[0]
+    output = candidate.decode(latents, return_dict=False)[0]
+    assert output.stride() == expected.stride()
+    assert torch.equal(output, expected)
@@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    class FakeBackend:
        supports_prefix_kv_slicing = True

+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend

@@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    assert attention.attention.metadata.extra["valid_kv_length"] == 5


+def test_packed_attention_skips_mask_for_packed_mask_free_backend():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeBackend:
+        supports_prefix_kv_slicing = False
+
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return True
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = FakeBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    q = torch.randn(8, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
+
+
 def test_packed_attention_keeps_padding_mask_for_other_backends():
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
    class FakeBackend:
        supports_prefix_kv_slicing = False

+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend

@@ -728,6 +728,8 @@ 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):
+    prepared_pipelines = []
+
    class _DummyLoader:
        def __init__(self, load_config, od_config=None):
            del load_config, od_config
@@ -775,6 +777,11 @@ 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_wan_vae_data_movement",
+        lambda pipeline: prepared_pipelines.append(pipeline),
+    )
    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
@@ -785,6 +792,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    assert runner.cache_backend is None
    assert runner.od_config.cache_backend is None
    assert dummy_cache_backend.enabled is False
+    assert prepared_pipelines == [runner.pipeline]


@pytest.mark.core_model
@@ -6,7 +6,6 @@ from pathlib import Path
 from types import SimpleNamespace

 import pytest
-import yaml

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

@@ -47,111 +46,9 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
        generate_examples.load_model_display_names()


-def test_quantization_examples_keep_source_urls_under_feature_navigation(tmp_path, monkeypatch):
-    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",
-                                    {"Offline Inference": [{"Old": "user_guide/examples/offline_inference/old.md"}]},
-                                ]
-                            },
-                            {"Features": [{"Quantization": [{"Overview": "user_guide/quantization/overview.md"}]}]},
-                        ]
-                    }
-                ]
-            },
-            sort_keys=False,
-        ),
-        encoding="utf-8",
-    )
+def test_only_general_serving_examples_get_generated_pages():
+    general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
+    model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))

-    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",
-    )
-    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
+    assert generate_examples.is_general_example(general_example)
+    assert not generate_examples.is_general_example(model_example)
@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
    return found


-def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
+def _check(
+    label: str,
+    path: Path,
+    actual: int,
+    budget: int,
+    sample: list[str],
+    budget_name: str,
+    notices: list[str],
+) -> list[str]:
    if actual > budget:
        listing = "\n".join(f"      {line}" for line in sample)
        return [
            f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
        ]
    if actual < budget:
-        return [
-            f"{path}: {actual} {label}, down from a budget of {budget}. "
-            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
-            f"so the ratchet holds the new ground.",
-        ]
+        # Progress is never a build failure. Removing branches is the goal, so a
+        # PR that does it must not hand-edit this file to stay green -- nor break
+        # `main` for everyone else when it merges from a stale base.
+        notices.append(
+            f"{path}: {actual} {label}, below the budget of {budget}. "
+            f"Lower {budget_name} to {actual} in "
+            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
+        )
    return []


@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
            print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
            return 1

+    notices: list[str] = []
    branches = _model_type_branches(ast.parse(serving_path.read_text()))
    detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))

@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
        MAX_MODEL_TYPE_BRANCHES,
        [f"line {lineno}: {src}" for lineno, src in branches],
        "MAX_MODEL_TYPE_BRANCHES",
+        notices,
    )
    errors += _check(
        "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
        MAX_LEGACY_DETECTORS,
        [f"line {lineno}: {name}" for lineno, name in detectors],
        "MAX_LEGACY_DETECTORS",
+        notices,
    )

+    for notice in notices:
+        print(f"check_tts_adapter: {notice}", file=sys.stderr)
    for error in errors:
        print(f"check_tts_adapter: {error}", file=sys.stderr)
    return 1 if errors else 0
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
    # avoid a slower masked-attention plan when tail padding is not semantic.
    supports_prefix_kv_slicing: bool = False

+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        """Whether packed attention never reads attn_mask on this platform.
+
+        When True, models that pack a [real, pad] two-document layout and
+        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
+        constructing the padding mask entirely. Backends whose mask-free
+        behavior is platform-dependent must check current_omni_platform.
+        """
+        return False
+
    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
    # Platform resolution rejects an explicit selection outside this set, so a
    # hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
    #     the packed cu_seqlens tensors.
    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
    #     backend that advertises supports_prefix_kv_slicing.
+    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
+    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
+    #     Requires the [real, pad] two-document packing contract; see
+    #     FlashAttentionImpl._forward_varlen_packed_npu.
+    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
+    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
+    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
+    #     with this set (>1), q/k/v are divided by the factor before the op,
+    #     the kernel scale_value is multiplied by its square, and the output
+    #     is scaled back (exact for power-of-two factors). Absent means no
+    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.

    # Piecewise attention metadata (mixed causal/full masks).
    # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import os
 from functools import partial

 import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
    piecewise_attn,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.platforms import current_omni_platform

 logger = init_logger(__name__)

@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supports_piecewise_spans: bool = True

+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
+        # before ever reading attn_mask. NPU: forward_fa_npu honors the
+        # npu_attn_varlen opt-in and rebuilds the padding mask itself if the
+        # packed contract fails. XPU reads attn_mask, so models must keep
+        # constructing it there.
+        return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
+
    @classmethod
    def supports_attention_mask(cls) -> bool:
        return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
+        # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
+        # model marks extra["npu_attn_varlen"] and carries packed metadata, so
+        # the padding document is excluded without reading or materializing
+        # the quadratic full_qk mask.
+        #   - default: packed varlen via mindiesd attention_forward_varlen;
+        #   - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
+        #     slicing via mindiesd attention_forward (so the env-selected op
+        #     type is honored; the TND varlen op cannot be).
+        extra = attn_metadata.extra if attn_metadata else {}
+        if extra.get("npu_attn_varlen", False):
+            if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
+                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
+            else:
+                out = self._forward_varlen_packed_npu(query, key, value, extra)
+            if out is not None:
+                return out
        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        if attention_mask is None and extra.get("npu_attn_varlen", False):
+            # Models skip mask construction when this opt-in is set
+            # (FlashAttentionBackend.supports_packed_mask_free). The packed
+            # paths above declined (contract mismatch), so rebuild the padding
+            # mask here to keep the masked fallback correct.
+            used = extra.get("valid_kv_length")
+            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
+                raise ValueError(
+                    "npu_attn_varlen packed metadata is unusable and no attn_mask "
+                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
+                    "refusing to run unmasked attention over padding rows."
+                )
+            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used

        # NPU aclnnFlashAttentionScore requires mask shape to be one of:
        # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
            op_type="fused_attn_score",
            layout=layout,
        )
+
+    def _resolve_packed_seq_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        extra: dict,
+    ) -> tuple[list[int], list[int]] | None:
+        """Resolve packed document boundaries (cumulative end offsets per doc).
+
+        Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
+        or None unless the metadata matches the exact contract MiniMax-H3
+        produces:
+          - q/k/v are [1, T, N, D] (BSND, single packed batch);
+          - cu_seqlens describe a [real, pad] two-document packing, with the
+            padding document as a strict suffix;
+          - max_seqlen_q/k are Python ints equal to the real document length
+            (so boundaries are derived without any device sync).
+        """
+        if self.causal:
+            return None
+        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+        if any(k not in extra for k in packed_keys):
+            return None
+        cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
+        used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
+        if query.shape[0] != 1 or key.shape[0] != 1:
+            return None
+        if not isinstance(used_q, int) or not isinstance(used_k, int):
+            return None
+        total_q, total_k = query.shape[1], key.shape[1]
+        # .shape is host-side metadata: counting documents never syncs.
+        if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
+            return None
+        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
+            return None
+        if used_q == total_q and used_k == total_k:
+            # No padding document: one full-length document.
+            return [total_q], [total_k]
+        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
+            # [real, pad] packing: the real document must be the longer one
+            # (consistent with the max_seqlen naming).
+            return [used_q, total_q], [used_k, total_k]
+        return None
+
+    def _forward_varlen_packed_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        extra: dict,
+    ) -> torch.Tensor | None:
+        """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
+
+        Returns None (caller falls back to the mask path) when the packed
+        contract does not hold; see _resolve_packed_seq_npu.
+        """
+        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        if resolved is None:
+            return None
+        seq_q, seq_k = resolved
+
+        try:
+            from mindiesd import attention_forward_varlen
+        except ImportError:
+            raise ImportError(
+                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
+                "Please install MindIE-SD to enable NPU attention support. "
+                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+        q = query.squeeze(0)  # [T, N, D] == TND
+        k = key.squeeze(0)
+        v = value.squeeze(0)
+        # attention_forward_varlen wants cu_seqlens as host lists of cumulative
+        # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
+        out = attention_forward_varlen(
+            q,
+            k,
+            v,
+            [0, *seq_q],
+            [0, *seq_k],
+            softmax_scale=self.softmax_scale,
+        )
+        return out.unsqueeze(0)
+
+    def _forward_prefix_kv_slice_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        extra: dict,
+    ) -> torch.Tensor | None:
+        """Mask-free attention by slicing K/V to the valid prefix (zero-copy
+        views), then mindiesd attention_forward without a mask.
+
+        Same numerical contract as the varlen path: the padding document is a
+        strict suffix, so dropping it from K/V is identical to masking it out.
+        Query keeps full length; outputs on padding rows are never consumed
+        downstream. Returns None when the packed contract does not hold.
+
+        Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
+        cannot honor the env-selected op type, while attention_forward can.
+
+        The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
+        bf16 activations with large outliers overflow 65504 into ±inf and the
+        affected rows turn NaN. Models may opt into exact power-of-two input
+        pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
+        are divided by the factor, the kernel scale_value is multiplied by its
+        square, and the output is scaled back. Absent or invalid factor means
+        no pre-scaling.
+        """
+        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        if resolved is None:
+            return None
+        _, seq_k = resolved
+        used_k = seq_k[0]  # real document length (first cumulative end)
+
+        try:
+            from mindiesd import attention_forward
+        except ImportError:
+            raise ImportError(
+                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
+                "Please install MindIE-SD to enable NPU attention support. "
+                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+        # mindiesd always takes BSND input regardless of `layout`; the arg
+        # selects the op-internal layout and laser only supports BNSD, so do
+        # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
+        layout = "BNSD"
+        key = key[:, :used_k]
+        value = value[:, :used_k]
+
+        input_scale = extra.get("laser_input_scale")
+        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
+        if preserve_input_range:
+            query = query / input_scale
+            key = key / input_scale
+            value = value / input_scale
+        scale = self.softmax_scale
+        if preserve_input_range:
+            scale *= input_scale**2
+
+        def _postprocess(out: torch.Tensor) -> torch.Tensor:
+            return out * input_scale if preserve_input_range else out
+
+        return _postprocess(
+            attention_forward(
+                query,
+                key,
+                value,
+                attn_mask=None,
+                opt_mode="manual",
+                op_type="fused_attn_score",  # MINDIE_SD_FA_TYPE env overrides this
+                layout=layout,
+                scale=scale,
+            )
+        )
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
    other case — warmup denoise steps, exempt layers, a layer that does not declare
    ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection or not aligned to the kernel's
-    block size — delegates to FlashAttention, so a model can select this backend
-    unconditionally.
+    segments too short to pay for block selection — delegates to FlashAttention,
+    so a model can select this backend unconditionally. MindIE-SD handles an
+    irregular video tail internally, retaining it outside the sparse blocks.
    """

    def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
                _MIN_VIDEO_BLOCKS,
            )
            return None
-        if video_len % _BLOCK_SIZE != 0:
-            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
-            # builds its block mask from the two segments pooled apart. The counts
-            # only agree when the video segment ends on a block boundary; otherwise
-            # a prefix block is forced on out of range, and the block straddling the
-            # seam holds real prefix rows that selection may drop. Padding the
-            # sequence to realign it is not an option because rf_v2 ignores
-            # attn_mask, so pad keys would take a share of every softmax denominator.
-            logger.warning_once(
-                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
-                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
-                "multiple of %d to enable sparsity.",
-                tuple(latent_shape),
-                video_len,
-                _BLOCK_SIZE,
-                _BLOCK_SIZE,
-            )
-            return None
-
        logger.info_once(
            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
@@ -18,13 +18,26 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
    GridSpec,
    TileTask,
 )
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
+    install_wan_vae_data_movement,
+)
 from vllm_omni.platforms import current_omni_platform

 logger = init_logger(__name__)


+def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None:
+    """Install Wan data-movement methods before offload or compilation setup."""
+    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+    for vae in ModuleDiscovery.discover(pipeline).vaes:
+        if isinstance(vae, OmniAutoencoderKLWan):
+            install_wan_vae_data_movement(vae)
+
+
 class OmniAutoencoderKLWan(AutoencoderKLWan):
    def _execution_context(self):
+        install_wan_vae_data_movement(self)
        try:
            first_param = next(self.parameters())
        except StopIteration:
@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guarded Wan VAE decoder data-movement fast paths.
+
+This module adapts the lossless Wan decoder optimizations from SGLang while
+keeping Diffusers' parameter names and module topology unchanged. Every
+optimized method falls back to the Diffusers 0.38.0 operation order when the
+kernel, device, layout, dtype, autograd state, or module type is unsupported.
+"""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    DupUp3D,
+    WanCausalConv3d,
+    WanDecoder3d,
+    WanResample,
+    WanResidualBlock,
+    WanResidualUpBlock,
+)
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
+
+try:
+    from vllm_omni.diffusion.kernels.wan_vae_data_movement import (
+        cat_pad_5d,
+        dup_up3d_add,
+    )
+except ImportError:  # pragma: no cover - Triton is optional on non-GPU installs.
+    cat_pad_5d = None
+    dup_up3d_add = None
+
+_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed"
+_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv"
+
+
+def _cache_payload(cache: Any) -> torch.Tensor | None:
+    return cache if isinstance(cache, torch.Tensor) else None
+
+
+def _is_dynamic_spatial_conv(conv: nn.Module) -> bool:
+    return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False))
+
+
+def _spatial_shard_is_active() -> bool:
+    return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None
+
+
+def _causal_padding(conv: nn.Module) -> tuple[int, ...]:
+    if _is_dynamic_spatial_conv(conv):
+        return tuple(conv._source_padding)
+    return tuple(conv._padding)
+
+
+def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool:
+    return bool(
+        cat_pad_5d is not None
+        and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv))
+        and x.dim() == 5
+        and x.is_cuda
+        and not torch.is_grad_enabled()
+        and not torch.compiler.is_compiling()
+        and not _spatial_shard_is_active()
+    )
+
+
+def _run_cached_causal_conv(
+    conv: nn.Module,
+    x: torch.Tensor,
+    cache_list: list[Any],
+    index: int,
+) -> torch.Tensor:
+    """Run a causal conv and refresh its feature-cache slot."""
+    cache = cache_list[index]
+    is_repeat_marker = isinstance(cache, str)
+    payload = None if is_repeat_marker else _cache_payload(cache)
+    if _can_use_fused_cache_path(conv, x) and (
+        payload is None or (payload.device == x.device and payload.dtype == x.dtype)
+    ):
+        pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T)
+        if pair is not None:
+            conv_input, cache_list[index] = pair
+            return nn.Conv3d.forward(conv, conv_input)
+
+    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+    if cache_x.shape[2] < CACHE_T and payload is not None:
+        cache_x = torch.cat(
+            [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+            dim=2,
+        )
+    elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
+        cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
+
+    out = conv(x) if payload is None else conv(x, payload)
+    cache_list[index] = cache_x
+    return out
+
+
+def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None):
+    padding = list(self._padding)
+    if (
+        any(padding)
+        and _can_use_fused_cache_path(self, x)
+        and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype))
+    ):
+        conv_input = cat_pad_5d(x, cache_x, padding)
+        if conv_input is not None:
+            return nn.Conv3d.forward(self, conv_input)
+
+    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]
+    x = F.pad(x, padding)
+    return nn.Conv3d.forward(self, x)
+
+
+def _resample_forward(
+    self: WanResample,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    batch, channels, frames, height, width = x.size()
+    if self.mode == "upsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = "Rep"
+            feat_idx[0] += 1
+        else:
+            x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
+            feat_idx[0] += 1
+            x = x.reshape(batch, 2, channels, frames, height, width)
+            x = torch.stack((x[:, 0], x[:, 1]), 3)
+            x = x.reshape(batch, channels, frames * 2, height, width)
+
+    frames = x.shape[2]
+    x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
+    x = self.resample(x)
+    x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+    if self.mode == "downsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = x.clone()
+            feat_idx[0] += 1
+        else:
+            cache_x = x[:, :, -1:, :, :].clone()
+            x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
+            feat_cache[index] = cache_x
+            feat_idx[0] += 1
+    return x
+
+
+def _residual_block_forward(
+    self: WanResidualBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    residual = self.conv_shortcut(x)
+    x = self.nonlinearity(self.norm1(x))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv1(x)
+
+    x = self.dropout(self.nonlinearity(self.norm2(x)))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv2, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv2(x)
+    return x + residual
+
+
+def _residual_up_block_forward(
+    self: WanResidualUpBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    shortcut_source = x.clone()
+    for resnet in self.resnets:
+        if feat_cache is None:
+            x = resnet(x)
+        else:
+            x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    if self.upsampler is not None:
+        if feat_cache is None:
+            x = self.upsampler(x)
+        else:
+            x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    shortcut = self.avg_shortcut
+    if shortcut is None:
+        return x
+    if (
+        dup_up3d_add is not None
+        and type(shortcut) is DupUp3D
+        and x.is_cuda
+        and shortcut_source.is_cuda
+        and x.dtype == shortcut_source.dtype
+        and not torch.is_grad_enabled()
+        and not torch.compiler.is_compiling()
+        and not _spatial_shard_is_active()
+    ):
+        fused = dup_up3d_add(
+            x,
+            shortcut_source,
+            shortcut.factor_t,
+            shortcut.factor_s,
+            shortcut.repeats,
+            first_chunk,
+        )
+        if fused is not None:
+            return fused
+    return x + shortcut(shortcut_source, first_chunk=first_chunk)
+
+
+def _decoder_forward(
+    self: WanDecoder3d,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv_in(x)
+
+    x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
+    for up_block in self.up_blocks:
+        x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+
+    x = self.nonlinearity(self.norm_out(x))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv_out(x)
+    return x
+
+
+def install_wan_vae_data_movement(vae: nn.Module) -> bool:
+    """Install guarded fast paths on a Diffusers Wan decoder once."""
+    if getattr(vae, _INSTALL_ATTR, False):
+        return True
+    if cat_pad_5d is None or dup_up3d_add is None:
+        return False
+
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return False
+
+    decoder.forward = MethodType(_decoder_forward, decoder)
+    patched = 0
+    for module in decoder.modules():
+        module_type = type(module)
+        if module_type is WanCausalConv3d:
+            module.forward = MethodType(_causal_conv_forward, module)
+            patched += 1
+        elif _is_dynamic_spatial_conv(module):
+            # The dynamic wrapper owns halo exchange while a spatial request is
+            # active. Its direct path can still use the lossless fused cache
+            # movement through ``_run_cached_causal_conv``.
+            patched += 1
+        elif module_type is WanResidualBlock:
+            module.forward = MethodType(_residual_block_forward, module)
+        elif module_type is WanResidualUpBlock:
+            module.forward = MethodType(_residual_up_block_forward, module)
+        elif module_type is WanResample:
+            module.forward = MethodType(_resample_forward, module)
+
+    if patched == 0:
+        return False
+    setattr(vae, _INSTALL_ATTR, True)
+    return True
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,405 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bit-exact data-movement kernels for the Wan causal VAE decoder.
+
+Adapted from SGLang PR #34125:
+https://github.com/sgl-project/sglang/pull/34125
+
+The kernels in this module only move values, fill zeros, or perform the same
+two-operand addition as the PyTorch reference path. They preserve the input
+layout instead of selecting a new Conv3d layout, so enabling them does not
+change the decoder's cuDNN accumulation order.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_MAX_INT32 = 2**31 - 1
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@triton.jit
+def _cat_pad_5d_kernel(
+    x_ptr,
+    cache_ptr,
+    out_ptr,
+    keep_ptr,
+    total,
+    channels,
+    frames,
+    height,
+    width,
+    cache_frames,
+    out_frames,
+    out_height,
+    out_width,
+    zero_front_frames,
+    pad_height,
+    pad_width,
+    stride_x_batch,
+    stride_x_channel,
+    stride_x_frame,
+    stride_x_height,
+    stride_x_width,
+    stride_cache_batch,
+    stride_cache_channel,
+    stride_cache_frame,
+    stride_cache_height,
+    stride_cache_width,
+    has_cache: tl.constexpr,
+    keep_frames_const: tl.constexpr,
+    channels_inner: tl.constexpr,
+    index_64bit: tl.constexpr,
+    block_size_const: tl.constexpr,
+):
+    if index_64bit:
+        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
+    else:
+        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
+    mask = offsets < total
+
+    if channels_inner:
+        out_channel = offsets % channels
+        remaining = offsets // channels
+        out_width_idx = remaining % out_width
+        remaining = remaining // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        out_batch = remaining // out_frames
+    else:
+        out_width_idx = offsets % out_width
+        remaining = offsets // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        remaining = remaining // out_frames
+        out_channel = remaining % channels
+        out_batch = remaining // channels
+
+    input_width = out_width_idx - pad_width
+    input_height = out_height_idx - pad_height
+    input_frame = out_frame - zero_front_frames
+
+    spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height)
+    from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames)
+    from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames)
+
+    x_frame = input_frame - cache_frames
+    x_offset = (
+        out_batch * stride_x_batch
+        + out_channel * stride_x_channel
+        + x_frame * stride_x_frame
+        + input_height * stride_x_height
+        + input_width * stride_x_width
+    )
+    values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0)
+    if has_cache:
+        cache_offset = (
+            out_batch * stride_cache_batch
+            + out_channel * stride_cache_channel
+            + input_frame * stride_cache_frame
+            + input_height * stride_cache_height
+            + input_width * stride_cache_width
+        )
+        cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0)
+        values = tl.where(from_cache, cache_values, values)
+    tl.store(out_ptr + offsets, values, mask=mask)
+
+    if keep_frames_const > 0:
+        keep_frame = out_frame - (out_frames - keep_frames_const)
+        keep_mask = mask & spatial_valid & (keep_frame >= 0)
+        if channels_inner:
+            keep_offset = (
+                ((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width
+            ) * channels + out_channel
+        else:
+            keep_offset = (
+                ((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height
+            ) * width + input_width
+        tl.store(keep_ptr + keep_offset, values, mask=keep_mask)
+
+
+def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
+    return tensor.shape[1] > 1 and tensor.stride(1) == 1
+
+
+def cat_pad_5d(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    padding: list[int] | tuple[int, ...],
+    keep_cache_frames: int = 0,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
+    """Fuse temporal concat, constant padding, and optional cache refresh.
+
+    ``padding`` follows the Wan causal Conv3d convention
+    ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
+    the same dense layout that the reference ``cat``/``pad`` path selects.
+    Unsupported inputs return ``None`` so callers can execute the reference
+    implementation.
+    """
+    if len(padding) != 6:
+        return None
+    pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
+    if any(value < 0 for value in padding):
+        return None
+    if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
+        return None
+    if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if keep_cache_frames < 0:
+        return None
+
+    channels_inner = _uses_channels_last_3d(x)
+    cache_frames = 0
+    if cache_x is not None:
+        if (
+            cache_x.dim() != 5
+            or cache_x.dtype != x.dtype
+            or cache_x.device != x.device
+            or cache_x.shape[0] != x.shape[0]
+            or cache_x.shape[1] != x.shape[1]
+            or cache_x.shape[3:] != x.shape[3:]
+            or _uses_channels_last_3d(cache_x) != channels_inner
+        ):
+            return None
+        cache_frames = cache_x.shape[2]
+
+    zero_front_frames = pad_front - cache_frames
+    if zero_front_frames < 0:
+        return None
+
+    batch, channels, frames, height, width = x.shape
+    out_frames = pad_front + frames
+    out_height = height + 2 * pad_height_top
+    out_width = width + 2 * pad_width_left
+    keep_frames = min(keep_cache_frames, out_frames)
+    memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format
+    out = torch.empty(
+        (batch, channels, out_frames, out_height, out_width),
+        device=x.device,
+        dtype=x.dtype,
+        memory_format=memory_format,
+    )
+    total = out.numel()
+    if total == 0 or total > _MAX_INT32 * 4:
+        return None
+
+    if keep_frames:
+        keep_arg = torch.empty(
+            (batch, channels, keep_frames, height, width),
+            device=x.device,
+            dtype=x.dtype,
+            memory_format=memory_format,
+        )
+    else:
+        keep_arg = out
+
+    if cache_x is None:
+        cache_arg = x
+        cache_strides = (0, 0, 0, 0, 0)
+    else:
+        cache_arg = cache_x
+        cache_strides = cache_x.stride()
+
+    block_size = 512
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(x.device):
+        _cat_pad_5d_kernel[grid](
+            x,
+            cache_arg,
+            out,
+            keep_arg,
+            total,
+            channels,
+            frames,
+            height,
+            width,
+            cache_frames,
+            out_frames,
+            out_height,
+            out_width,
+            zero_front_frames,
+            pad_height_top,
+            pad_width_left,
+            *x.stride(),
+            *cache_strides,
+            has_cache=cache_x is not None,
+            keep_frames_const=keep_frames,
+            channels_inner=channels_inner,
+            index_64bit=total >= _MAX_INT32,
+            block_size_const=block_size,
+        )
+
+    if keep_cache_frames:
+        return out, keep_arg
+    return out
+
+
+@triton.jit
+def _dup_up3d_add_kernel(
+    main_ptr,
+    source_ptr,
+    out_ptr,
+    total,
+    out_channels,
+    out_frames,
+    out_height,
+    out_width,
+    frame_offset,
+    stride_main_batch,
+    stride_main_channel,
+    stride_main_frame,
+    stride_main_height,
+    stride_main_width,
+    stride_source_batch,
+    stride_source_channel,
+    stride_source_frame,
+    stride_source_height,
+    stride_source_width,
+    stride_out_batch,
+    stride_out_channel,
+    stride_out_frame,
+    stride_out_height,
+    stride_out_width,
+    factor_temporal_const: tl.constexpr,
+    factor_spatial_const: tl.constexpr,
+    repeats_const: tl.constexpr,
+    channels_inner: tl.constexpr,
+    index_64bit: tl.constexpr,
+    block_size_const: tl.constexpr,
+):
+    if index_64bit:
+        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
+    else:
+        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
+    mask = offsets < total
+
+    if channels_inner:
+        out_channel = offsets % out_channels
+        remaining = offsets // out_channels
+        out_width_idx = remaining % out_width
+        remaining = remaining // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        out_batch = remaining // out_frames
+    else:
+        out_width_idx = offsets % out_width
+        remaining = offsets // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        remaining = remaining // out_frames
+        out_channel = remaining % out_channels
+        out_batch = remaining // out_channels
+
+    unsliced_frame = out_frame + frame_offset
+    source_frame = unsliced_frame // factor_temporal_const
+    temporal_remainder = unsliced_frame % factor_temporal_const
+    source_height = out_height_idx // factor_spatial_const
+    height_remainder = out_height_idx % factor_spatial_const
+    source_width = out_width_idx // factor_spatial_const
+    width_remainder = out_width_idx % factor_spatial_const
+    repeated_channel = (
+        (out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder
+    ) * factor_spatial_const + width_remainder
+    source_channel = repeated_channel // repeats_const
+
+    main_offset = (
+        out_batch * stride_main_batch
+        + out_channel * stride_main_channel
+        + out_frame * stride_main_frame
+        + out_height_idx * stride_main_height
+        + out_width_idx * stride_main_width
+    )
+    source_offset = (
+        out_batch * stride_source_batch
+        + source_channel * stride_source_channel
+        + source_frame * stride_source_frame
+        + source_height * stride_source_height
+        + source_width * stride_source_width
+    )
+    out_offset = (
+        out_batch * stride_out_batch
+        + out_channel * stride_out_channel
+        + out_frame * stride_out_frame
+        + out_height_idx * stride_out_height
+        + out_width_idx * stride_out_width
+    )
+    main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0)
+    source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0)
+    value = main_value.to(tl.float32) + source_value.to(tl.float32)
+    tl.store(out_ptr + out_offset, value, mask=mask)
+
+
+def dup_up3d_add(
+    main: torch.Tensor,
+    source: torch.Tensor,
+    factor_temporal: int,
+    factor_spatial: int,
+    repeats: int,
+    drop_first_frames: bool,
+) -> torch.Tensor | None:
+    """Evaluate ``main + DupUp3D(source)`` without materializing the shortcut."""
+    if main.dim() != 5 or source.dim() != 5:
+        return None
+    if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0:
+        return None
+    if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1):
+        return None
+    if repeats & (repeats - 1):
+        return None
+    if not main.is_cuda or not source.is_cuda:
+        return None
+    if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
+        return None
+
+    batch, source_channels, frames, height, width = source.shape
+    frame_offset = factor_temporal - 1 if drop_first_frames else 0
+    numerator = source_channels * repeats
+    denominator = factor_temporal * factor_spatial * factor_spatial
+    if numerator % denominator:
+        return None
+    expected_shape = (
+        batch,
+        numerator // denominator,
+        frames * factor_temporal - frame_offset,
+        height * factor_spatial,
+        width * factor_spatial,
+    )
+    if tuple(main.shape) != expected_shape:
+        return None
+
+    out = torch.empty_like(main)
+    total = out.numel()
+    if total == 0 or total > _MAX_INT32 * 4:
+        return None
+
+    block_size = 512
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(main.device):
+        _dup_up3d_add_kernel[grid](
+            main,
+            source,
+            out,
+            total,
+            expected_shape[1],
+            expected_shape[2],
+            expected_shape[3],
+            expected_shape[4],
+            frame_offset,
+            *main.stride(),
+            *source.stride(),
+            *out.stride(),
+            factor_temporal_const=factor_temporal,
+            factor_spatial_const=factor_spatial,
+            repeats_const=repeats,
+            channels_inner=out.stride(1) == 1 and expected_shape[1] > 1,
+            index_64bit=total >= _MAX_INT32,
+            block_size_const=block_size,
+        )
+    return out
@@ -21,6 +21,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader

 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -215,6 +216,7 @@ class HeliosPipeline(
        self.vae = AutoencoderKLWan.from_pretrained(
            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
        ).to(self.device)
+        install_wan_vae_data_movement(self.vae)

        transformer_config = load_transformer_config(model, "transformer", local_files_only)
        self.transformer = create_transformer_from_config(
@@ -19,6 +19,7 @@ from torch import nn
 from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor

 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import (
    SupportImageInput,
@@ -374,6 +375,7 @@ class LingBotVideoPipeline(
            torch_dtype=vae_dtype,
            local_files_only=local_files_only,
        ).to(self.device)
+        install_wan_vae_data_movement(self.vae)
        self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
            model,
            subfolder=scheduler_subfolder,
@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
 # lookup and masked out afterwards).
 MINIMAX_H3_ADALN_MODALITY_NUM = 3

+# Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
+# (consumed only via the "laser_input_scale" extra key; other backends and
+# platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
+# workspace, and H3's outlier activations (per-element amax in the hundreds)
+# push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
+# 256 is a power of two, so pre-dividing q/k/v and the compensating
+# kernel-scale/output multiplies are exact in floating point.
+MINIMAX_H3_LASER_INPUT_SCALE = 256.0
+

 def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
    if key not in kwargs or kwargs[key] is None:
@@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
        used = min(max_seqlen, packed_total)
        attn_mask = None
        # Ring attention can dispatch to a different implementation from the
-        # configured backend, so this no-mask fast path is local-only.
-        prefix_slice = (
-            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
+        # configured backend, so the no-mask fast paths are local-only.
+        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
+        # supports_packed_mask_free: backend consumes the packed metadata
+        # without ever reading attn_mask (CUDA packed varlen, NPU
+        # npu_attn_varlen opt-in with its own fallback rebuild).
+        no_mask = not getattr(self.attention, "use_ring", False) and (
+            self.attention.attn_backend.supports_prefix_kv_slicing
+            or self.attention.attn_backend.supports_packed_mask_free()
        )
-        if used < packed_total and not prefix_slice:
+        if used < packed_total and not no_mask:
            attn_mask = torch.arange(packed_total, device=q.device)[None] < used
        metadata = AttentionMetadata(
            attn_mask=attn_mask,
@@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
                "max_seqlen_q": max_seqlen,
                "max_seqlen_k": max_seqlen,
                "valid_kv_length": used,
+                # Opt the NPU flash backend into the packed varlen path so the
+                # quadratic full_qk mask is never materialized. Ring attention
+                # is excluded: it keeps the aligned padding rows for its
+                # fixed-size P2P buffers and still needs the mask.
+                "npu_attn_varlen": not getattr(self.attention, "use_ring", False),
+                # fp16-range protection for the ascend_laser_attention kernel
+                # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
+                # backend/path.
+                "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
            },
            video_layout=video_layout,
        )
@@ -30,6 +30,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    prepare_pipeline_wan_vae_data_movement,
+)
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                    custom_pipeline_name=custom_pipeline_name,
                    device=self.device,
                )
+                prepare_pipeline_wan_vae_data_movement(self.pipeline)
        time_after_load = time.perf_counter()

        logger.info(
@@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
                )
                backend_upper = "FLASH_ATTN"

+            if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
+                # The NPU FLASH_ATTN backend imports mindiesd lazily at first
+                # forward, but CANN snapshots the custom-op registry at the
+                # first custom-op regInfo lookup in the process (e.g. a
+                # vllm-ascend custom op during model load/warmup). Import
+                # mindiesd here so its env.py prepends the mindiesd vendor
+                # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
+                # that snapshot; otherwise aclnnLaserAttention /
+                # FusedAttentionScore fail with EZ1001 "does not support
+                # opType" for the rest of the process.
+                import mindiesd  # noqa: F401
+
            backend = DiffusionAttentionBackendEnum[backend_upper]
            logger.debug("Using diffusion attention backend '%s'", backend_upper)
            return backend.get_path()
Author	SHA1	Message	Date
Hongsheng Liu	5308222abb	[Perf][Diffusion] Fuse lossless Wan VAE data movement Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 09:43:11 +08:00
fan2956andFan	c27623c29c	[Bugfix] rainfusion support all shape video (#6000 ) Signed-off-by: Fan <fan@FandeMacBook-Air.local> Co-authored-by: Fan <fan@FandeMacBook-Air.local>	2026-08-11 05:10:50 +08:00
Yueqian Lin	f8fa86fb1f	[BugFix] Don't fail the TTS ratchet when the branch count goes down (#6008 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-08-11 05:08:31 +08:00
ooooooyeandHongsheng Liu	422d8fd8f9	[Bugfix][Diffusion] Restore supports_packed_mask_free on teacache FakeBackend (#5997 ) Signed-off-by: brandneway <gyuan4892@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 22:20:37 +08:00
Hongsheng Liu	62e239bf1e	[Docs] Fix broken attention backend link (#5998 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 21:55:01 +08:00
ooooooye	740f45abcd	[Perf][Diffusion] MiniMax-H3: opt into packed varlen attention on NPU to eliminate quadratic mask materialization (#5891 ) Signed-off-by: brandneway <gyuan4892@gmail.com>	2026-08-10 19:47:03 +08:00
Hongsheng Liu hsliu_ustc Copilot Autofix powered by AI	5e4b9f8c6e	[Misc] Align CODEOWNERS with module/feature design docs (#5958 ) Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>	2026-08-10 19:25:03 +08:00
Hongsheng Liu	ae811043e9	docs: keep shared task examples in navigation (#5987 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 19:13:52 +08:00