feat(diffusion): add request-gated VAE fast paths

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Bugfix] rainfusion support all shape video (#6000 )
2026-08-11 09:31:27 +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
52 changed files with 2487 additions and 260 deletions
@@ -1,6 +1,11 @@
 # CODEOWNERS is used to automatically assign PR reviewers
 # Format: path pattern @username or @team
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 #
 # The vllm_omni/** source-tree rules below are derived from the `owners` and
 # `primary_code_paths` front matter in docs/design/module/*.md. When updating
 # ownership, update the design doc first, then regenerate the matching rule here.
 # Last-match-wins: more specific rules must come AFTER broader ones.
 # Default owners (fallback reviewers for all files)
 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao
-# Engine and model_executor
+# ComfyUI / external apps
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/apps/ @fhfuih @alex-jw-brooks
 /vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
 /vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/worker/ @tzhouam @gcanlin
-# Diffusion
+# Benchmarks (out-of-tree)
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
 # Diffusion CPU offloading
 /vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
 # Omni and TTS support
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
 /vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
 # Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
 /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
-# Usability and Observability
+# =====================================================================
-/vllm_omni/deploy/ @tzhouam @lishunyang12
+# Module design ownership (docs/design/module/*.md)
-/vllm_omni/metrics/ @lishunyang12
+# Each section header cites the source design doc.
-/vllm_omni/profiler/ @gcanlin
+# =====================================================================
-# ComfyUI
+# Autoregressive runtime — docs/design/module/ar_runtime.md
-/apps/ @fhfuih @alex-jw-brooks
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
 /vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
-# Large Scale Serving
+# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
 /vllm_omni/engine/ @tzhouam @fake0fan
 # File-level overrides: input/output modality contracts own these engine files
 # — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Entrypoints and serving boundaries — docs/design/module/entrypoints.md
 /vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
 # File-level overrides within entrypoints/
 /vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
 /vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
 # Error contracts — docs/design/module/error_contracts.md
 /vllm_omni/errors.py @alex-jw-brooks @NickCao
 # OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
-# Quantization and configuration
+# Model integration — docs/design/module/model_integration.md
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
 /vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
 /vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
 # Config override within model_executor — docs/design/module/vllm_omni_config.md
 /vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
 # Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
 /vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
 /vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
 # Configuration & deploy — docs/design/module/vllm_omni_config.md
 /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
 # Quantization — docs/design/module/quantization.md
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
-# Hardware plugins and accelerators
+# Observability — docs/design/module/observability.md
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+/vllm_omni/metrics/ @lishunyang12 @vraiti
 # Profiling — docs/design/module/profiling.md
 /vllm_omni/profiler/ @gcanlin
 # Benchmarking (in-tree) — docs/design/module/benchmarking.md
 /vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
 # LoRA (no dedicated design doc; ownership intentionally narrowed)
 /vllm_omni/lora/ @SamitHuang @wtomin
 # Transformers utils (related path only in design docs; preserves prior ownership)
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
 # =====================================================================
 # Diffusion module family — docs/design/module/diffusion/
 # Umbrella rule first; sub-module rules after it win via last-match.
 # =====================================================================
 # Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
 # — docs/design/module/diffusion/index.md
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 # Cache management — docs/design/module/cache_management.md
 /vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
 # Diffusion quantization — docs/design/module/quantization.md
 /vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
 # Diffusion profiling — docs/design/module/profiling.md
 /vllm_omni/diffusion/profiler/ @gcanlin
 # Diffusion runtime (scheduler, executor, worker)
 # — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
 /vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
 # Diffusion model integration (pipelines, registration, checkpoint loading)
 # — docs/design/module/diffusion/diffusion_model_integration.md
 /vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
 /vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 /vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 # Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
 # — docs/design/module/diffusion/parallelism.md
 /vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 /vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 # Diffusion offloader (CPU offloading, DLO)
 # — docs/design/module/diffusion/offloader.md
 /vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
@@ -66,7 +66,10 @@ repos:
      entry: python tools/pre_commit/check_tts_adapter.py
      language: python
      pass_filenames: false
-      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
      # The budget-is-now-slack reminder is printed on a PASSING run, and
      # pre-commit hides output from passing hooks unless the hook is verbose.
      verbose: true
    # Keep `suggestion` last
    - id: suggestion
@@ -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
+      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
      - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
    - Online Serving:
      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
+      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
      - Text-To-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
+# Keep the Examples navigation focused on reusable modality/task entry points.
-# navigation category changes. Quantization helpers are feature documentation,
+# Model-specific scripts remain available under ``examples/`` and are surfaced
-# not user-serving examples, so keep them under the existing Features /
+# from these shared pages instead of getting one page per model.
-# Quantization section instead of creating a second top-level Quantization
+GENERAL_EXAMPLE_SLUGS = frozenset(
-# category under Examples.
+    {
-EXAMPLE_NAV_TARGETS = {
+        "image_to_image",
-    "quantization": "quantization",
+        "image_to_video",
-}
+        "speech_to_video",
-GENERATED_QUANTIZATION_EXAMPLES_TITLE = "Generated Examples"
+        "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
+    # Build the new Examples section - start with preserved items.
    # not explicitly routed elsewhere retain the historical Examples behavior.
    examples_section = preserved_items.copy()
    for category in sorted(examples_by_category.keys()):
        if EXAMPLE_NAV_TARGETS.get(category, "examples") != "examples":
            continue
        items = category_items(category)
        if items:
            category_title = fix_case(category.replace("_", " ").title())
@@ -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
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
-is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
-tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
-that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+promoted rows remain visible to every query and key.
 attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
 speedup, so pick an aligned geometry when you want one.
-Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
-spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
-patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
-are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+and kernel tiling aligned, at the cost of retaining more blocks.
 distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
 to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
-"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
-resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
-which you can buy back with `start_step`.
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
 that implements protected video tails for rf_v2.
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends
 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 ### Multi-NPU: 768P combined-service configuration
 Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
-degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
+degree 8, native tiled VAE patch parallelism degree 8, and distributed
 layerwise offload:
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
  --usp 8 \
  --ring 1 \
  --text-encoder-tp-size 8 \
-  --enable-layerwise-offload \
+  --enable-distributed-layerwise-offload \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --vae-patch-parallel-size 8 \
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
 partition restart is required. Layerwise offload applies to both DiTs.
 On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
 768P without offloading or sharding: use distributed layerwise offload (as
 above) or HSDP — see
 [§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
 ### Optional optimizations
 Two independent optimizations may be enabled on top of the configuration
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
 over the whole sequence, so ring parallelism would split away the keys it
 needs. Scale with `--usp` instead.
 ## Memory and attention optimizations (A3)
 ### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
 Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
 combined MiniMax-H3 service at 768P. Enable **one** of the following at
 server startup:
 **Distributed layerwise offload** — DiT layers are offloaded with parameters
 gathered across the parallel group instead of replicated per rank:
 ```bash
  --enable-distributed-layerwise-offload
 ```
 Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
 input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
 usage is high with this option.
 **HSDP** — hybrid sharded data parallelism for the DiT parameters. The
 multi-stream memory-reuse knob is mandatory with this flag:
 ```bash
 export MULTI_STREAM_MEMORY_REUSE=2
 ```
 ```bash
  --use-hsdp
 ```
 Host memory usage is small, but large shapes may still OOM; HBM usage
 optimizations for this path are ongoing.
 ### FLASH_ATTN backend with MindIE-SD
 Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
 [§ Environment](#environment)). On NPU this backend carries most of the
 memory-reduction and performance work: a mask-free packed varlen path that
 never materializes the quadratic `full_qk` padding mask, and K/V prefix
 slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
 transformer.
 ### Optional: LaserAttention fused kernel
 For an additional attention speedup, select the Ascend LaserAttention fused
 kernel before starting the server:
 ```bash
 export MINDIE_SD_FA_TYPE="ascend_laser_attention"
 ```
 This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
 exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
 kernel's fp16 workspace cannot overflow on outlier activations. Measured
 speedup numbers will be added here.
 ## HTTP API examples
 The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:
 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
 These measurements describe the validated shapes rather than a general
 throughput guarantee.
@@ -9,6 +9,10 @@ This script tests two main scenarios:
 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
 """
 import sys
 from types import SimpleNamespace
 from unittest.mock import Mock
 import pytest
 import torch
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
        impl.forward_cuda(query, query, query, metadata)
 # ---------------------------------------------------------------------------
 # NPU mask-free packed attention (PR #5891).
 #
 # These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
 #   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
 #     sync, no MindIE-SD import);
 #   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
 #   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
 #
 # They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
 # into ``sys.modules`` (same pattern as ``test_paged_attention`` /
 # ``benchmarks/conftest``). No production code is exercised for real kernels.
 def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
    return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
 def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
    """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
    Restored automatically by monkeypatch. Defaults are no-op callables so the
    top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
    instance methods that would otherwise call them.
    """
    fake = SimpleNamespace(
        attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
        attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
    )
    monkeypatch.setitem(sys.modules, "mindiesd", fake)
    return fake
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 def test_resolve_packed_seq_accepts_real_plus_pad_contract():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
 def test_resolve_packed_seq_single_document_no_padding():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
@pytest.mark.parametrize(
    "case",
    [
        "causal",
        "missing_key",
        "batch_ne_1",
        "cu_too_long",
        "cu_qk_mismatch",
        "max_seqlen_not_int",
        "used_gt_total",
        "used_zero",
        "pad_longer_than_real",
    ],
 )
 def test_resolve_packed_seq_rejects_malformed_metadata(case):
    """Any metadata outside the exact [real, pad] contract returns None so the
    caller falls back to the masked path."""
    cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
    cu2 = torch.tensor([0, 8], dtype=torch.int32)
    cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
    cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
    base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
    q = torch.randn(1, 8, 2, 4)
    impl = _npu_impl(causal=(case == "causal"))
    if case == "missing_key":
        extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5}  # 3 of 4 keys absent
    elif case == "batch_ne_1":
        q = torch.randn(2, 8, 2, 4)
        extra = base
    elif case == "cu_too_long":
        extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
    elif case == "cu_qk_mismatch":
        extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
    elif case == "max_seqlen_not_int":
        extra = {**base, "max_seqlen_q": 5.0}
    elif case == "used_gt_total":
        extra = {**base, "max_seqlen_q": 9}  # 9 > total length 8
    elif case == "used_zero":
        extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
    elif case == "pad_longer_than_real":
        # real=3, pad=5: the real document must be the longer one.
        extra = {
            "cu_seqlens_q": cu_short,
            "cu_seqlens_k": cu_short,
            "max_seqlen_q": 3,
            "max_seqlen_k": 3,
        }
    else:
        extra = base
    assert impl._resolve_packed_seq_npu(q, q, extra) is None
 # --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
 def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
    """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
    Lets a test assert which branch ``forward_fa_npu`` dispatched to without
    running any real kernel.
    """
    impl = _npu_impl()
    sentinel = torch.tensor([1.0])
    impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
    impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
    return impl, sentinel
 def test_npu_env_unset_uses_varlen(monkeypatch):
    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_env_laser_uses_slice(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_prefix_kv_slice_npu.assert_called_once()
    impl._forward_varlen_packed_npu.assert_not_called()
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
 def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
    """Locks in the CURRENT dispatch behavior (PR #5891).
    Any value other than ``ascend_laser_attention`` — including the valid
    ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
    the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
    comment #1 on the PR and is intentionally captured here as a regression
    net: update this test when the dispatch starts honoring/rejecting these.
    """
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
    """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
    the existing masked attention_forward path; the new branches are skipped."""
    fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
    _fake_mindiesd(monkeypatch, attention_forward=fake_forward)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(
        attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
        extra={},  # no npu_attn_varlen opt-in
    )
    q = torch.randn(1, 8, 2, 4)
    out = impl.forward_fa_npu(q, q, q, metadata)
    impl._forward_varlen_packed_npu.assert_not_called()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
    fake_forward.assert_called_once()
    assert out is fake_forward.return_value
 # --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
 def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured["key"] = key
        captured["value"] = value
        captured.update(kwargs)
        return torch.ones_like(query)  # deterministic stand-in kernel output
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()  # softmax_scale == 0.5
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        "laser_input_scale": 256.0,
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    # K/V sliced to the valid prefix (used_k == 5); q keeps full length.
    assert captured["key"].shape == (1, 5, 2, 4)
    assert captured["value"].shape == (1, 5, 2, 4)
    assert captured["query"].shape == (1, 8, 2, 4)
    assert captured["attn_mask"] is None
    # Pre-divide q/k/v by the factor and compensate the kernel scale by its
    # square (exact for power-of-two); output scaled back by the factor.
    torch.testing.assert_close(captured["query"], q / 256.0)
    torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
    assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
 def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured.update(kwargs)
        return torch.ones_like(query)
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        # no laser_input_scale → no pre-scaling
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    torch.testing.assert_close(captured["query"], q)
    assert captured["scale"] == pytest.approx(0.5)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
 if __name__ == "__main__":
    print("Running FlashAttention Padding Tests...")
    print("=" * 60)
@@ -5,10 +5,9 @@
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
-mask from the video and prefix segments separately, so the two only agree on the
+to the always-kept prefix segment before generating the block mask.
 block count when the video segment ends on a boundary.
 """
 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
    assert plan.latent_shape == list(ALIGNED_GRID)
-def test_misaligned_video_segment_stays_dense():
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
-    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
-    # pad keys would take a share of every real query's softmax denominator.
+
-    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+    assert plan is not None
    assert plan.prefix_len == PREFIX_ROWS
    assert plan.used_len == PREFIX_ROWS + 62496
    assert plan.latent_shape == list(MISALIGNED_GRID)
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
 def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
    torch.manual_seed(1)
    prefix_len, grid = 384, (4, 16, 10)
    video_len = grid[0] * grid[1] * grid[2]
    used = prefix_len + video_len
    heads, head_dim = 4, 128
    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
    impl = make_impl()
    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
    out = impl._forward_sparse_npu(q, k, v, plan)
    reference = torch.nn.functional.scaled_dot_product_attention(
        q.transpose(1, 2),
        k.transpose(1, 2),
        v.transpose(1, 2),
        scale=impl.softmax_scale,
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
    class FakeBackend:
        supports_prefix_kv_slicing = False
        # Mirror AttentionBackend.supports_packed_mask_free: H3's packed
        # attention path calls attn_backend.supports_packed_mask_free() to decide
        # whether to build the padding mask. Return False so the fake exercises
        # the masked path (the abstract backend default).
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    def __init__(self, **kwargs):
        del kwargs
        super().__init__()
@@ -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
@@ -113,6 +113,14 @@ class TestCreateDefaultDiffusion:
        stages = StageConfigFactory.create_default_diffusion({"model": "x"})
        assert stages[0]["engine_args"]["cache_backend"] == "none"
    def test_vae_decode_precision_roundtrip(self):
        od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16")
        assert od.vae_decode_precision == "bf16"
    def test_invalid_vae_decode_precision_is_rejected(self):
        with pytest.raises(ValueError, match="vae_decode_precision"):
            OmniDiffusionConfig(model="x", vae_decode_precision="int8")
    def test_single_gpu_default_devices(self):
        stages = StageConfigFactory.create_default_diffusion({"model": "x"})
        assert stages[0]["runtime"]["devices"] == "0"
@@ -512,6 +512,35 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
    assert runner.pipeline.last_req.num_reqs == 1
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
    ("quality", "expected_enabled"),
    [(None, False), ("lossless", False), ("high", True)],
 )
 def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled):
    runner = _make_runner(cache_backend=None, cache_backend_name="none")
    req = _make_request()
    req.sampling_params.quality = quality
    gate_calls = []
    @contextmanager
    def _record_gate(pipeline, enabled):
        gate_calls.append((pipeline, enabled, "enter"))
        yield
        gate_calls.append((pipeline, enabled, "exit"))
    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
    monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate)
    DiffusionModelRunner.execute_model(runner, req)
    assert gate_calls == [
        (runner.pipeline, expected_enabled, "enter"),
        (runner.pipeline, expected_enabled, "exit"),
    ]
 class _BatchPipeline:
    """Pipeline returning a configurable list of outputs from forward()."""
@@ -579,6 +608,18 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
    sched = _make_scheduler_output(num_reqs=2)
    sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
    sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
    with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -728,6 +769,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.core_model
@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    memory_pool_active = False
    optimization_pool_states = []
    @contextmanager
    def _memory_pool_context(tag):
        nonlocal memory_pool_active
        assert tag == "weights"
        memory_pool_active = True
        try:
            yield
        finally:
            memory_pool_active = False
    class _DummyLoader:
        def __init__(self, load_config, od_config=None):
            del load_config, od_config
@@ -775,16 +829,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
    monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
    monkeypatch.setattr(
        model_runner_module,
        "optimize_pipeline_vaes",
        lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
    )
    monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
    monkeypatch.setattr(
        model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
    )
-    DiffusionModelRunner.load_model(runner)
+    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
    assert runner.cache_backend is None
    assert runner.od_config.cache_backend is None
    assert dummy_cache_backend.enabled is False
    assert optimization_pool_states == [True]
@pytest.mark.core_model
@@ -102,6 +102,7 @@ class TestDiffusionWorkerSleep:
    def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 1 (offload weights only)."""
        clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches")
        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
@@ -128,6 +129,7 @@ class TestDiffusionWorkerSleep:
        # Verify sleep was called with correct tags
        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
        clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
        assert bool(result) is True
        # Verify buffers were NOT saved (level 1 doesn't save buffers)
        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
@@ -0,0 +1,300 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Correctness contracts for request-gated VAE decoder fast paths."""
 from __future__ import annotations
 import pytest
 import torch
 import torch.nn.functional as F
 from torch import nn
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.vae_optimizations.gate import (
    VaeFastPathGate,
    register_vae_fast_path_gate,
    use_pipeline_vae_fast_path,
    use_vae_fast_path,
 )
 pytestmark = [pytest.mark.diffusion]
 class _PipelineWithVae(nn.Module):
    def __init__(self, vae: nn.Module) -> None:
        super().__init__()
        self.vae = vae
 class _DistributedExecutorConfig:
    def __init__(self, parallel_mode: str) -> None:
        self.parallel_mode = parallel_mode
@pytest.mark.core_model
@pytest.mark.cpu
 def test_vae_fast_path_gate_restores_nested_state() -> None:
    vae = nn.Identity()
    gate = VaeFastPathGate()
    register_vae_fast_path_gate(vae, gate)
    with use_vae_fast_path(vae, True):
        assert gate.enabled
        with use_vae_fast_path(vae, False):
            assert not gate.enabled
        assert gate.enabled
    assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
 def test_pipeline_vae_fast_path_controls_discovered_vae() -> None:
    vae = nn.Identity()
    pipeline = _PipelineWithVae(vae)
    gate = VaeFastPathGate()
    register_vae_fast_path_gate(vae, gate)
    with use_pipeline_vae_fast_path(pipeline, True):
        assert gate.enabled
    assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
 def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None:
    from vllm_omni.diffusion.vae_optimizations import image
    gate = VaeFastPathGate()
    gate.enabled = True
    norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate)
    inputs = torch.randn(1, 4, 3, 3)
    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
    monkeypatch.setattr(
        image,
        "group_norm_silu_4d",
        lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"),
    )
    actual = norm(inputs)
    expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps))
    torch.testing.assert_close(actual, expected)
 def _small_image_vae() -> nn.Module:
    from diffusers.models.autoencoders import AutoencoderKL
    return AutoencoderKL(
        in_channels=3,
        out_channels=3,
        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"),
        block_out_channels=(32, 32),
        layers_per_block=1,
        latent_channels=4,
        norm_num_groups=8,
        sample_size=16,
    )
 def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module:
    if vae_cls is None:
        from diffusers.models.autoencoders import AutoencoderKLWan
        vae_cls = AutoencoderKLWan
    return vae_cls(
        base_dim=8,
        decoder_base_dim=8,
        z_dim=4,
        dim_mult=[1, 2],
        num_res_blocks=1,
        attn_scales=[],
        temperal_downsample=[False],
        latents_mean=[0.0] * 4,
        latents_std=[1.0] * 4,
        scale_factor_temporal=1,
        scale_factor_spatial=2,
    )
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
 def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None:
    from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
    from vllm_omni.diffusion.vae_optimizations.image import (
        FusedGroupNormSiLU,
        FusedUpsample2xConv2d,
        maybe_optimize_image_vae,
    )
    torch.manual_seed(0)
    vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval()
    latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16)
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents).sample
    maybe_optimize_image_vae(vae)
    maybe_optimize_image_vae(vae)
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules())
    assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules())
    assert torch.equal(vae.decode(latents).sample, reference)
    with use_vae_fast_path(vae, True):
        fast = vae.decode(latents).sample
    torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0)
    assert any(
        buffer is not None
        for module in vae.modules()
        for name, buffer in module._buffers.items()
        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
    )
    clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae))
    assert all(
        buffer is None
        for module in vae.modules()
        for name, buffer in module._buffers.items()
        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
    )
    with use_vae_fast_path(vae, True):
        torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0)
    assert torch.equal(vae.decode(latents).sample, reference)
    assert vae.decoder.conv_in.weight.is_contiguous()
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None:
    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d
    torch.manual_seed(0)
    x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last)
    weight = torch.randn(64, device="cuda", dtype=dtype)
    bias = torch.randn(64, device="cuda", dtype=dtype)
    actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5)
    expected = F.silu(F.group_norm(x, 8, weight, bias))
    assert actual is not None
    atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2)
    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
    "activation_dtype,affine_dtype,atol,rtol",
    [
        (torch.float32, torch.float32, 1e-5, 1e-5),
        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),
    ],
 )
 def test_wan_rmsnorm_silu_matches_reference(
    activation_dtype: torch.dtype,
    affine_dtype: torch.dtype,
    atol: float,
    rtol: float,
 ) -> None:
    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu
    torch.manual_seed(0)
    x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous(
        memory_format=torch.channels_last_3d
    )
    gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype)
    bias = torch.randn_like(gamma)
    normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype)
    expected = F.silu(normalized * 96**0.5 * gamma + bias)
    actual = wan_rmsnorm_silu(x, gamma, bias)
    assert actual is not None
    assert actual.stride() == x.stride()
    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
 def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None:
    from vllm_omni.diffusion.vae_optimizations.wan import (
        FusedWanRMSNormSiLU,
        maybe_optimize_wan_vae,
    )
    torch.manual_seed(0)
    vae = _small_wan_vae().cuda().eval()
    latents = torch.randn(1, 4, 1, 8, 8, device="cuda")
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents).sample
    maybe_optimize_wan_vae(vae)
    maybe_optimize_wan_vae(vae)
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules())
    assert torch.equal(vae.decode(latents).sample, reference)
    with use_vae_fast_path(vae, True):
        fast = vae.decode(latents).sample
    torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4)
    assert torch.equal(vae.decode(latents).sample, reference)
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"])
 def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None:
    import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations
    monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True)
    vae = _small_wan_vae()
    vae.distributed_executor = _DistributedExecutorConfig(parallel_mode)
    wan_optimizations.maybe_optimize_wan_vae(vae)
    assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules())
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
    ("precision", "decode_dtype"),
    [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
 )
 def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None:
    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
    from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
    torch.manual_seed(0)
    vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval()
    configure_wan_decode_precision(vae, precision)
    assert next(vae.encoder.parameters()).dtype == torch.float32
    assert next(vae.quant_conv.parameters()).dtype == torch.float32
    assert next(vae.decoder.parameters()).dtype == decode_dtype
    assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype
    _, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype))
    assert grid_spec.output_dtype == decode_dtype
    encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16))
    decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32))
    assert encoded.latent_dist.parameters.dtype == torch.float32
    assert decoded.sample.dtype == decode_dtype
@@ -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):
+def test_only_general_serving_examples_get_generated_pages():
-    nav_file = tmp_path / "docs" / ".nav.yml"
+    general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
-    nav_file.parent.mkdir(parents=True)
+    model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))
    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",
    )
-    monkeypatch.setattr(generate_examples, "ROOT_DIR", tmp_path)
+    assert generate_examples.is_general_example(general_example)
-    monkeypatch.setattr(generate_examples, "NAV_FILE", nav_file)
+    assert not generate_examples.is_general_example(model_example)
    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
@@ -76,6 +76,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
    assert args.cfg_parallel_size == 4
 def test_serve_parser_accepts_wan_decode_precision() -> None:
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    OmniServeCommand().subparser_init(subparsers)
    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
    assert args.vae_decode_precision == "bf16"
    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
 def _make_headless_args(**kwargs) -> TrackingNamespace:
    defaults = {
        "model": "fake-model",
@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
    return found
-def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
+def _check(
    label: str,
    path: Path,
    actual: int,
    budget: int,
    sample: list[str],
    budget_name: str,
    notices: list[str],
 ) -> list[str]:
    if actual > budget:
        listing = "\n".join(f"      {line}" for line in sample)
        return [
            f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
        ]
    if actual < budget:
-        return [
+        # Progress is never a build failure. Removing branches is the goal, so a
-            f"{path}: {actual} {label}, down from a budget of {budget}. "
+        # PR that does it must not hand-edit this file to stay green -- nor break
-            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
+        # `main` for everyone else when it merges from a stale base.
-            f"so the ratchet holds the new ground.",
+        notices.append(
-        ]
+            f"{path}: {actual} {label}, below the budget of {budget}. "
            f"Lower {budget_name} to {actual} in "
            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
        )
    return []
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
            print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
            return 1
    notices: list[str] = []
    branches = _model_type_branches(ast.parse(serving_path.read_text()))
    detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
        MAX_MODEL_TYPE_BRANCHES,
        [f"line {lineno}: {src}" for lineno, src in branches],
        "MAX_MODEL_TYPE_BRANCHES",
        notices,
    )
    errors += _check(
        "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
        MAX_LEGACY_DETECTORS,
        [f"line {lineno}: {name}" for lineno, name in detectors],
        "MAX_LEGACY_DETECTORS",
        notices,
    )
    for notice in notices:
        print(f"check_tts_adapter: {notice}", file=sys.stderr)
    for error in errors:
        print(f"check_tts_adapter: {error}", file=sys.stderr)
    return 1 if errors else 0
@@ -573,6 +573,7 @@ class _DiffusionConfigProjection:
    fa_deterministic: bool = False
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None
    mask_strategy_file_path: str | None = None
    skip_time_steps: int = 15
    VSA_sparsity: float = 0.0
@@ -401,6 +401,7 @@ class StageDeployConfig:
    step_execution: bool | None = None
    vae_use_slicing: bool | None = None
    vae_use_tiling: bool | None = None
    vae_decode_precision: str | None = None
    boundary_ratio: float | None = None
    flow_shift: float | None = None
    diffusion_kv_cache_dtype: str | None = None
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
    # avoid a slower masked-attention plan when tail padding is not semantic.
    supports_prefix_kv_slicing: bool = False
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        """Whether packed attention never reads attn_mask on this platform.
        When True, models that pack a [real, pad] two-document layout and
        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
        constructing the padding mask entirely. Backends whose mask-free
        behavior is platform-dependent must check current_omni_platform.
        """
        return False
    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
    # Platform resolution rejects an explicit selection outside this set, so a
    # hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
    #     the packed cu_seqlens tensors.
    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
    #     backend that advertises supports_prefix_kv_slicing.
    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
    #     Requires the [real, pad] two-document packing contract; see
    #     FlashAttentionImpl._forward_varlen_packed_npu.
    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
    #     with this set (>1), q/k/v are divided by the factor before the op,
    #     the kernel scale_value is multiplied by its square, and the output
    #     is scaled back (exact for power-of-two factors). Absent means no
    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
    # Piecewise attention metadata (mixed causal/full masks).
    # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from functools import partial
 import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
    piecewise_attn,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supports_piecewise_spans: bool = True
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
        # before ever reading attn_mask. NPU: forward_fa_npu honors the
        # npu_attn_varlen opt-in and rebuilds the padding mask itself if the
        # packed contract fails. XPU reads attn_mask, so models must keep
        # constructing it there.
        return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
    @classmethod
    def supports_attention_mask(cls) -> bool:
        return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
        # model marks extra["npu_attn_varlen"] and carries packed metadata, so
        # the padding document is excluded without reading or materializing
        # the quadratic full_qk mask.
        #   - default: packed varlen via mindiesd attention_forward_varlen;
        #   - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
        #     slicing via mindiesd attention_forward (so the env-selected op
        #     type is honored; the TND varlen op cannot be).
        extra = attn_metadata.extra if attn_metadata else {}
        if extra.get("npu_attn_varlen", False):
            if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
            else:
                out = self._forward_varlen_packed_npu(query, key, value, extra)
            if out is not None:
                return out
        attention_mask = attn_metadata.attn_mask if attn_metadata else None
        if attention_mask is None and extra.get("npu_attn_varlen", False):
            # Models skip mask construction when this opt-in is set
            # (FlashAttentionBackend.supports_packed_mask_free). The packed
            # paths above declined (contract mismatch), so rebuild the padding
            # mask here to keep the masked fallback correct.
            used = extra.get("valid_kv_length")
            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
                raise ValueError(
                    "npu_attn_varlen packed metadata is unusable and no attn_mask "
                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
                    "refusing to run unmasked attention over padding rows."
                )
            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
        # NPU aclnnFlashAttentionScore requires mask shape to be one of:
        # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
            op_type="fused_attn_score",
            layout=layout,
        )
    def _resolve_packed_seq_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        extra: dict,
    ) -> tuple[list[int], list[int]] | None:
        """Resolve packed document boundaries (cumulative end offsets per doc).
        Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
        or None unless the metadata matches the exact contract MiniMax-H3
        produces:
          - q/k/v are [1, T, N, D] (BSND, single packed batch);
          - cu_seqlens describe a [real, pad] two-document packing, with the
            padding document as a strict suffix;
          - max_seqlen_q/k are Python ints equal to the real document length
            (so boundaries are derived without any device sync).
        """
        if self.causal:
            return None
        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
        if any(k not in extra for k in packed_keys):
            return None
        cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
        used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
        if query.shape[0] != 1 or key.shape[0] != 1:
            return None
        if not isinstance(used_q, int) or not isinstance(used_k, int):
            return None
        total_q, total_k = query.shape[1], key.shape[1]
        # .shape is host-side metadata: counting documents never syncs.
        if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
            return None
        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
            return None
        if used_q == total_q and used_k == total_k:
            # No padding document: one full-length document.
            return [total_q], [total_k]
        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
            # [real, pad] packing: the real document must be the longer one
            # (consistent with the max_seqlen naming).
            return [used_q, total_q], [used_k, total_k]
        return None
    def _forward_varlen_packed_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
        Returns None (caller falls back to the mask path) when the packed
        contract does not hold; see _resolve_packed_seq_npu.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        seq_q, seq_k = resolved
        try:
            from mindiesd import attention_forward_varlen
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        q = query.squeeze(0)  # [T, N, D] == TND
        k = key.squeeze(0)
        v = value.squeeze(0)
        # attention_forward_varlen wants cu_seqlens as host lists of cumulative
        # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
        out = attention_forward_varlen(
            q,
            k,
            v,
            [0, *seq_q],
            [0, *seq_k],
            softmax_scale=self.softmax_scale,
        )
        return out.unsqueeze(0)
    def _forward_prefix_kv_slice_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Mask-free attention by slicing K/V to the valid prefix (zero-copy
        views), then mindiesd attention_forward without a mask.
        Same numerical contract as the varlen path: the padding document is a
        strict suffix, so dropping it from K/V is identical to masking it out.
        Query keeps full length; outputs on padding rows are never consumed
        downstream. Returns None when the packed contract does not hold.
        Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
        cannot honor the env-selected op type, while attention_forward can.
        The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
        bf16 activations with large outliers overflow 65504 into ±inf and the
        affected rows turn NaN. Models may opt into exact power-of-two input
        pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
        are divided by the factor, the kernel scale_value is multiplied by its
        square, and the output is scaled back. Absent or invalid factor means
        no pre-scaling.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        _, seq_k = resolved
        used_k = seq_k[0]  # real document length (first cumulative end)
        try:
            from mindiesd import attention_forward
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # mindiesd always takes BSND input regardless of `layout`; the arg
        # selects the op-internal layout and laser only supports BNSD, so do
        # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
        layout = "BNSD"
        key = key[:, :used_k]
        value = value[:, :used_k]
        input_scale = extra.get("laser_input_scale")
        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
        if preserve_input_range:
            query = query / input_scale
            key = key / input_scale
            value = value / input_scale
        scale = self.softmax_scale
        if preserve_input_range:
            scale *= input_scale**2
        def _postprocess(out: torch.Tensor) -> torch.Tensor:
            return out * input_scale if preserve_input_range else out
        return _postprocess(
            attention_forward(
                query,
                key,
                value,
                attn_mask=None,
                opt_mode="manual",
                op_type="fused_attn_score",  # MINDIE_SD_FA_TYPE env overrides this
                layout=layout,
                scale=scale,
            )
        )
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
    other case — warmup denoise steps, exempt layers, a layer that does not declare
    ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection or not aligned to the kernel's
+    segments too short to pay for block selection — delegates to FlashAttention,
-    block size — delegates to FlashAttention, so a model can select this backend
+    so a model can select this backend unconditionally. MindIE-SD handles an
-    unconditionally.
+    irregular video tail internally, retaining it outside the sparse blocks.
    """
    def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
                _MIN_VIDEO_BLOCKS,
            )
            return None
        if video_len % _BLOCK_SIZE != 0:
            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
            # builds its block mask from the two segments pooled apart. The counts
            # only agree when the video segment ends on a block boundary; otherwise
            # a prefix block is forced on out of range, and the block straddling the
            # seam holds real prefix rows that selection may drop. Padding the
            # sequence to realign it is not an option because rf_v2 ignores
            # attn_mask, so pad keys would take a share of every softmax denominator.
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
                "multiple of %d to enable sparsity.",
                tuple(latent_shape),
                video_len,
                _BLOCK_SIZE,
                _BLOCK_SIZE,
            )
            return None
        logger.info_once(
            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
@@ -711,6 +711,9 @@ class OmniDiffusionConfig:
    # VAE memory optimization parameters
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    # Optional Wan decode-only precision. When set, the VAE encoder remains
    # FP32 while decoder-side modules use the selected precision.
    vae_decode_precision: str | None = None
    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
@@ -923,6 +926,11 @@ class OmniDiffusionConfig:
        )
    def __post_init__(self):
        if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"):
            raise ValueError(
                "vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, "
                f"got {self.vae_decode_precision!r}"
            )
        if self.diffusion_compile_granularity not in {"regional", "full"}:
            raise ValueError(
                "diffusion_compile_granularity must be 'regional' or 'full', "
@@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
            split_dims=(3, 4),
            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
            tile_spec=tile_spec,
-            output_dtype=self.dtype,
+            output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype),
        )
        return tiletask_list, grid_spec
@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
 # lookup and masked out afterwards).
 MINIMAX_H3_ADALN_MODALITY_NUM = 3
 # Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
 # (consumed only via the "laser_input_scale" extra key; other backends and
 # platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
 # workspace, and H3's outlier activations (per-element amax in the hundreds)
 # push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
 # 256 is a power of two, so pre-dividing q/k/v and the compensating
 # kernel-scale/output multiplies are exact in floating point.
 MINIMAX_H3_LASER_INPUT_SCALE = 256.0
 def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
    if key not in kwargs or kwargs[key] is None:
@@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
        used = min(max_seqlen, packed_total)
        attn_mask = None
        # Ring attention can dispatch to a different implementation from the
-        # configured backend, so this no-mask fast path is local-only.
+        # configured backend, so the no-mask fast paths are local-only.
-        prefix_slice = (
+        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
-            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
+        # supports_packed_mask_free: backend consumes the packed metadata
        # without ever reading attn_mask (CUDA packed varlen, NPU
        # npu_attn_varlen opt-in with its own fallback rebuild).
        no_mask = not getattr(self.attention, "use_ring", False) and (
            self.attention.attn_backend.supports_prefix_kv_slicing
            or self.attention.attn_backend.supports_packed_mask_free()
        )
-        if used < packed_total and not prefix_slice:
+        if used < packed_total and not no_mask:
            attn_mask = torch.arange(packed_total, device=q.device)[None] < used
        metadata = AttentionMetadata(
            attn_mask=attn_mask,
@@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
                "max_seqlen_q": max_seqlen,
                "max_seqlen_k": max_seqlen,
                "valid_kv_length": used,
                # Opt the NPU flash backend into the packed varlen path so the
                # quadratic full_qk mask is never materialized. Ring attention
                # is excluded: it keeps the aligned padding rows for its
                # fixed-size P2P buffers and still needs the mask.
                "npu_attn_varlen": not getattr(self.attention, "use_ring", False),
                # fp16-range protection for the ascend_laser_attention kernel
                # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
                # backend/path.
                "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
            },
            video_layout=video_layout,
        )
@@ -0,0 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Optional VAE decoder optimizations."""
 from vllm_omni.diffusion.vae_optimizations.gate import (
    use_pipeline_vae_fast_path,
    use_vae_fast_path,
 )
 from vllm_omni.diffusion.vae_optimizations.optimize import (
    clear_pipeline_vae_fast_path_caches,
    optimize_pipeline_vaes,
 )
 __all__ = [
    "clear_pipeline_vae_fast_path_caches",
    "optimize_pipeline_vaes",
    "use_pipeline_vae_fast_path",
    "use_vae_fast_path",
 ]
@@ -0,0 +1,67 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Decode-scoped gates for optional, non-bit-exact VAE fast paths."""
 from __future__ import annotations
 from collections.abc import Iterator
 from contextlib import ExitStack, contextmanager
 from weakref import WeakKeyDictionary
 from torch import nn
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 class VaeFastPathGate:
    """Mutable flag shared by wrappers installed on one VAE instance."""
    __slots__ = ("enabled",)
    def __init__(self) -> None:
        self.enabled = False
 _VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = WeakKeyDictionary()
 def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None:
    """Associate an installed fast path with its owning VAE."""
    _VAE_FAST_PATH_GATES[vae] = gate
@contextmanager
 def use_vae_fast_path(vae: nn.Module, enabled: bool) -> Iterator[None]:
    """Enable an installed fast path for one decode and always restore it."""
    gate = _VAE_FAST_PATH_GATES.get(vae)
    if gate is None:
        yield
        return
    previous_enabled = gate.enabled
    gate.enabled = enabled
    try:
        yield
    finally:
        gate.enabled = previous_enabled
@contextmanager
 def use_pipeline_vae_fast_path(pipeline: nn.Module, enabled: bool) -> Iterator[None]:
    """Apply one request's quality choice to every declared VAE component."""
    vaes = ModuleDiscovery.discover(pipeline).vaes
    with ExitStack() as stack:
        for vae in vaes:
            stack.enter_context(use_vae_fast_path(vae, enabled))
        yield
 __all__ = [
    "VaeFastPathGate",
    "register_vae_fast_path_gate",
    "use_pipeline_vae_fast_path",
    "use_vae_fast_path",
 ]
@@ -0,0 +1,354 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders."""
 from __future__ import annotations
 from types import MethodType
 from typing import Any
 import torch
 import torch.nn.functional as F
 from torch import nn
 from vllm.logger import init_logger
 from vllm_omni.diffusion.vae_optimizations.gate import (
    VaeFastPathGate,
    register_vae_fast_path_gate,
 )
 logger = init_logger(__name__)
 try:
    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import (
        group_norm_silu_4d,
        group_norm_silu_rows,
    )
    _HAS_TRITON = True
 except ImportError:  # pragma: no cover - exercised on builds without Triton
    _HAS_TRITON = False
 class FusedGroupNormSiLU(nn.Module):
    """Quality-gated GroupNorm + SiLU with an exact eager fallback."""
    def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
        super().__init__()
        self.num_groups = norm.num_groups
        self.num_channels = norm.num_channels
        self.eps = norm.eps
        self.affine = norm.affine
        # Register parameters directly so state-dict FQNs do not change.
        self.weight = norm.weight
        self.bias = norm.bias
        self._vllm_gate = gate
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling():
            output = group_norm_silu_4d(
                x,
                self.weight,
                self.bias,
                self.num_groups,
                self.eps,
                apply_silu=True,
            )
            if output is not None:
                return output
        return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps))
 def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
    def is_fusable(norm: object) -> bool:
        return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None
    count = 0
    for module in decoder.modules():
        if (
            type(module) is resnet_cls
            and module.time_emb_proj is None
            and module.time_embedding_norm in ("default", "group")
            and module.upsample is None
            and module.downsample is None
            and isinstance(module.nonlinearity, nn.SiLU)
            and is_fusable(module.norm1)
            and is_fusable(module.norm2)
        ):
            module.norm1 = FusedGroupNormSiLU(module.norm1, gate)
            module.norm2 = FusedGroupNormSiLU(module.norm2, gate)
            module.nonlinearity = nn.Identity()
            count += 2
    if (
        type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm
        and isinstance(getattr(decoder, "conv_act", None), nn.SiLU)
        and is_fusable(decoder.conv_norm_out)
    ):
        decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate)
        decoder.conv_act = nn.Identity()
        count += 1
    return count
 _UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)}
 def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
    weight = conv.weight.detach().float()
    output_channels, input_channels = weight.shape[:2]
    folded = weight.new_zeros(input_channels, output_channels, 4, 4)
    for row in range(4):
        for column in range(4):
            value = weight.new_zeros(output_channels, input_channels)
            for source_row in _UPSAMPLE_TAP_MAP[row]:
                for source_column in _UPSAMPLE_TAP_MAP[column]:
                    value += weight[:, :, source_row, source_column]
            folded[:, :, row, column] = value.t()
    folded = folded.to(conv.weight.dtype)
    if conv.weight.is_contiguous(memory_format=torch.channels_last):
        folded = folded.contiguous(memory_format=torch.channels_last)
    return folded.to(conv.weight.device)
 class FusedUpsample2xConv2d(nn.Module):
    """Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d."""
    def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:
        super().__init__()
        # Keep only ``conv`` registered so checkpoint parameter names stay
        # ``...upsamplers.N.conv.*``. The original module is an eager fallback.
        object.__setattr__(self, "_original", upsample)
        self.conv = upsample.conv
        self.channels = upsample.channels
        self._vllm_gate = gate
        self.register_buffer("_fused_weight", None, persistent=False)
    def forward(
        self,
        hidden_states: torch.Tensor,
        output_size: int | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor:
        if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels:
            return self._original(hidden_states, output_size=output_size, *args, **kwargs)
        weight = self._fused_weight
        if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype:
            weight = _fold_upsample2x_conv2d_weight(self.conv)
            self._fused_weight = weight
        return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1)
 def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
    count = 0
    for block in decoder.up_blocks:
        upsamplers = getattr(block, "upsamplers", None)
        if not upsamplers:
            continue
        for index, upsample in enumerate(upsamplers):
            if type(upsample) is not upsample_cls:
                continue
            conv = getattr(upsample, "conv", None)
            if (
                upsample.use_conv
                and not upsample.use_conv_transpose
                and upsample.interpolate
                and upsample.norm is None
                and upsample.name == "conv"
                and type(conv) is nn.Conv2d
                and conv.kernel_size == (3, 3)
                and conv.stride == (1, 1)
                and conv.padding == (1, 1)
                and conv.dilation == (1, 1)
                and conv.groups == 1
                and conv.padding_mode == "zeros"
            ):
                upsamplers[index] = FusedUpsample2xConv2d(upsample, gate)
                count += 1
    return count
 def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
    value_weight = module.to_v.weight.detach().float()
    value_bias = module.to_v.bias.detach().float()
    output_weight = module.to_out[0].weight.detach().float()
    output_bias = module.to_out[0].bias.detach().float()
    dtype = module.to_v.weight.dtype
    return (
        (output_weight @ value_weight).to(dtype).contiguous(),
        (output_weight @ value_bias + output_bias).to(dtype),
    )
 def _attention_fast_forward(
    self: nn.Module,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    temb: torch.Tensor | None = None,
    **cross_attention_kwargs: Any,
 ) -> torch.Tensor:
    if (
        not self._vllm_gate.enabled
        or encoder_hidden_states is not None
        or attention_mask is not None
        or temb is not None
        or hidden_states.ndim != 4
    ):
        return type(self).forward(
            self,
            hidden_states,
            encoder_hidden_states=encoder_hidden_states,
            attention_mask=attention_mask,
            temb=temb,
            **cross_attention_kwargs,
        )
    residual = hidden_states
    batch_size, channels, height, width = hidden_states.shape
    hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
    if self.group_norm is not None:
        norm = self.group_norm
        normalized = None
        if not torch.compiler.is_compiling():
            normalized = group_norm_silu_rows(
                hidden_states,
                norm.weight,
                norm.bias,
                norm.num_groups,
                norm.eps,
                apply_silu=False,
            )
        hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2)
    query = self.to_q(hidden_states)
    key = self.to_k(hidden_states)
    folded_weight = self._vllm_folded_value_weight
    folded_bias = self._vllm_folded_value_bias
    if (
        folded_weight is None
        or folded_bias is None
        or folded_weight.device != hidden_states.device
        or folded_weight.dtype != self.to_v.weight.dtype
    ):
        folded_weight, folded_bias = _fold_attention_value_projection(self)
        self._vllm_folded_value_weight = folded_weight
        self._vllm_folded_value_bias = folded_bias
    value = F.linear(hidden_states, folded_weight, folded_bias)
    output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1))
    output = output.squeeze(1).to(query.dtype)
    output = self.to_out[1](output)
    output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
    if self.residual_connection:
        output = output + residual
    return output / self.rescale_output_factor
 def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool:
    return (
        type(module) is attention_cls
        and isinstance(module.processor, processor_cls)
        and module.heads == 1
        and module.scale_qk
        and module.spatial_norm is None
        and not module.norm_cross
        and getattr(module, "norm_q", None) is None
        and getattr(module, "norm_k", None) is None
        and getattr(module, "add_k_proj", None) is None
        and type(module.to_q) is nn.Linear
        and type(module.to_k) is nn.Linear
        and type(module.to_v) is nn.Linear
        and type(module.to_out[0]) is nn.Linear
        and isinstance(module.to_out[1], nn.Dropout)
        and module.to_v.bias is not None
        and module.to_out[0].bias is not None
    )
 def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
    use_channels_last = self._vllm_gate.enabled
    if use_channels_last != self._vllm_channels_last:
        memory_format = torch.channels_last if use_channels_last else torch.contiguous_format
        self.to(memory_format=memory_format)
        self._vllm_channels_last = use_channels_last
        logger.debug(
            "%s decoder switched to %s layout",
            self._vllm_label,
            "channels_last" if use_channels_last else "contiguous",
        )
    return type(self).forward(self, *args, **kwargs)
 def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module:
    """Install image VAE decoder fast paths when the structure is fully supported."""
    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
        return vae
    from diffusers.models.attention_processor import Attention, AttnProcessor2_0
    from diffusers.models.autoencoders import AutoencoderKL
    from diffusers.models.autoencoders.vae import Decoder
    from diffusers.models.resnet import ResnetBlock2D
    from diffusers.models.upsampling import Upsample2D
    supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,)
    try:
        from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
        supported_types += (AutoencoderKLFlux2,)
    except ImportError:  # pragma: no cover - older supported Diffusers builds
        pass
    if not isinstance(vae, supported_types):
        return vae
    decoder = getattr(vae, "decoder", None)
    if type(decoder) is not Decoder:
        return vae
    attention_modules = [
        module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0)
    ]
    attention_count = sum(isinstance(module, Attention) for module in decoder.modules())
    if len(attention_modules) != attention_count:
        logger.info(
            "%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths",
            type(vae).__name__,
            attention_count - len(attention_modules),
            attention_count,
        )
        return vae
    gate = VaeFastPathGate()
    decoder._vllm_gate = gate
    decoder._vllm_label = type(vae).__name__
    decoder._vllm_channels_last = False
    decoder.forward = MethodType(_decoder_layout_forward, decoder)
    upsample_count = _install_fused_upsample(decoder, Upsample2D, gate)
    for module in attention_modules:
        module._vllm_gate = gate
        module.register_buffer("_vllm_folded_value_weight", None, persistent=False)
        module.register_buffer("_vllm_folded_value_bias", None, persistent=False)
        module.forward = MethodType(_attention_fast_forward, module)
    norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate)
    register_vae_fast_path_gate(vae, gate)
    vae._vllm_vae_fast_path_installed = True
    logger.info(
        "%s installed quality-gated VAE decoder fast paths "
        "(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)",
        type(vae).__name__,
        upsample_count,
        len(attention_modules),
        norm_count,
    )
    return vae
 __all__ = [
    "FusedGroupNormSiLU",
    "FusedUpsample2xConv2d",
    "maybe_optimize_image_vae",
 ]
@@ -0,0 +1,70 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Pipeline-level installation for optional VAE optimizations."""
 from __future__ import annotations
 from torch import nn
 from vllm.logger import init_logger
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.vae_optimizations.image import maybe_optimize_image_vae
 from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
 from vllm_omni.diffusion.vae_optimizations.wan import maybe_optimize_wan_vae
 from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 _FOLDED_CACHE_NAMES = (
    "_fused_weight",
    "_vllm_folded_value_weight",
    "_vllm_folded_value_bias",
 )
 def optimize_pipeline_vaes(pipeline: nn.Module, od_config: OmniDiffusionConfig) -> None:
    """Configure declared VAE components once after checkpoint loading."""
    precision = getattr(od_config, "vae_decode_precision", None)
    vaes = ModuleDiscovery.discover(pipeline).vaes
    for vae in vaes:
        # An explicit precision override is a serving contract, so failures on
        # a recognized Wan VAE must abort startup instead of being hidden.
        configure_wan_decode_precision(vae, precision)
    if not current_omni_platform.is_cuda():
        return
    for vae in vaes:
        for optimizer in (maybe_optimize_image_vae, maybe_optimize_wan_vae):
            try:
                optimizer(vae)
            except Exception:
                # A wrapper may already have installed its exact fallback
                # before a later optional rewrite failed. With no registered
                # enabled gate, those wrappers remain on the reference path.
                logger.warning(
                    "Failed to install optional VAE fast paths on %s; using the reference path",
                    type(vae).__name__,
                    exc_info=True,
                )
 def clear_pipeline_vae_fast_path_caches(pipeline: nn.Module) -> None:
    """Release lazily folded buffers before the worker enters sleep mode."""
    cleared = 0
    for vae in ModuleDiscovery.discover(pipeline).vaes:
        if not getattr(vae, "_vllm_vae_fast_path_installed", False):
            continue
        for module in vae.modules():
            for name in _FOLDED_CACHE_NAMES:
                if name in module._buffers and module._buffers[name] is not None:
                    setattr(module, name, None)
                    cleared += 1
    if cleared:
        logger.debug("Released %d lazy VAE fast-path buffers before sleep", cleared)
 __all__ = ["clear_pipeline_vae_fast_path_caches", "optimize_pipeline_vaes"]
@@ -0,0 +1,91 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Decode-only precision configuration for structurally split Wan VAEs."""
 from __future__ import annotations
 from types import MethodType
 from typing import Any
 import torch
 from torch import nn
 from vllm.logger import init_logger
 logger = init_logger(__name__)
 PRECISION_TO_DTYPE = {
    "fp32": torch.float32,
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
 }
 def _decode_with_precision(
    self: nn.Module,
    latents: torch.Tensor,
    *args: Any,
    **kwargs: Any,
 ) -> Any:
    return type(self).decode(
        self,
        latents.to(dtype=self._vllm_decode_dtype),
        *args,
        **kwargs,
    )
 def _encode_in_fp32(
    self: nn.Module,
    samples: torch.Tensor,
    *args: Any,
    **kwargs: Any,
 ) -> Any:
    return type(self).encode(
        self,
        samples.to(dtype=torch.float32),
        *args,
        **kwargs,
    )
 def configure_wan_decode_precision(vae: nn.Module, precision: str | None) -> nn.Module:
    """Keep Wan encode in FP32 while materializing only decode modules lower."""
    if precision is None:
        return vae
    from diffusers.models.autoencoders import AutoencoderKLWan
    if not isinstance(vae, AutoencoderKLWan):
        return vae
    try:
        decode_dtype = PRECISION_TO_DTYPE[precision]
    except KeyError as exc:
        raise ValueError(
            f"Unsupported vae_decode_precision={precision!r}; expected one of {sorted(PRECISION_TO_DTYPE)}"
        ) from exc
    encoder = getattr(vae, "encoder", None)
    quant_conv = getattr(vae, "quant_conv", None)
    decoder = getattr(vae, "decoder", None)
    post_quant_conv = getattr(vae, "post_quant_conv", None)
    if not all(isinstance(module, nn.Module) for module in (encoder, decoder, post_quant_conv)):
        logger.info("Wan VAE lacks separate encoder/decoder modules; skipping decode-only precision")
        return vae
    encoder.to(dtype=torch.float32)
    if isinstance(quant_conv, nn.Module):
        quant_conv.to(dtype=torch.float32)
    decoder.to(dtype=decode_dtype)
    post_quant_conv.to(dtype=decode_dtype)
    vae.encode = MethodType(_encode_in_fp32, vae)
    vae._vllm_decode_dtype = decode_dtype
    vae.decode = MethodType(_decode_with_precision, vae)
    logger.info(
        "Wan VAE configured with FP32 encode and %s decode",
        precision,
    )
    return vae
 __all__ = ["PRECISION_TO_DTYPE", "configure_wan_decode_precision"]
@@ -0,0 +1,237 @@
 # ruff: noqa: N803
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Channels-last two-pass GroupNorm(+SiLU) Triton kernels.
 The kernels keep the channel dimension innermost throughout a 2D VAE
 decoder. Statistics are accumulated in FP32, then a separate apply pass
 performs the affine transform and optional SiLU epilogue. Unsupported inputs
 return ``None`` so callers can fail closed to the reference implementation.
 """
 from __future__ import annotations
 import torch
 import triton
 import triton.language as tl
 _SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
 _MAX_CHANNELS = 2048
@triton.jit
 def _gn_partial_rows_kernel(
    x_ptr,
    psum_ptr,
    psq_ptr,
    rows,
    rows_per_prog,
    C: tl.constexpr,
    BLOCK_R: tl.constexpr,
 ):
    chunk = tl.program_id(0).to(tl.int64)
    batch = tl.program_id(1).to(tl.int64)
    num_chunks = tl.num_programs(0)
    row0 = chunk * rows_per_prog
    cols = tl.arange(0, C)
    acc_sum = tl.zeros((C,), tl.float32)
    acc_sq = tl.zeros((C,), tl.float32)
    x_base = x_ptr + batch * rows * C
    for row_offset in range(0, rows_per_prog, BLOCK_R):
        row_indices = row0 + row_offset + tl.arange(0, BLOCK_R)
        mask = row_indices < rows
        x = tl.load(
            x_base + row_indices[:, None] * C + cols[None, :],
            mask=mask[:, None],
            other=0.0,
        ).to(tl.float32)
        acc_sum += tl.sum(x, 0)
        acc_sq += tl.sum(x * x, 0)
    output_offsets = (batch * num_chunks + chunk) * C + cols
    tl.store(psum_ptr + output_offsets, acc_sum)
    tl.store(psq_ptr + output_offsets, acc_sq)
@triton.jit
 def _gn_finalize_kernel(
    psum_ptr,
    psq_ptr,
    weight_ptr,
    bias_ptr,
    scale_shift_ptr,
    num_chunks,
    group_numel,
    eps,
    channels,
    CPG: tl.constexpr,
    BLOCK_K: tl.constexpr,
 ):
    group = tl.program_id(0).to(tl.int64)
    batch = tl.program_id(1).to(tl.int64)
    cols = group * CPG + tl.arange(0, CPG)
    acc_sum = tl.zeros((), tl.float32)
    acc_sq = tl.zeros((), tl.float32)
    for chunk0 in range(0, num_chunks, BLOCK_K):
        chunks = chunk0 + tl.arange(0, BLOCK_K)
        mask = chunks < num_chunks
        offsets = (batch * num_chunks + chunks)[:, None] * channels + cols[None, :]
        acc_sum += tl.sum(tl.load(psum_ptr + offsets, mask=mask[:, None], other=0.0))
        acc_sq += tl.sum(tl.load(psq_ptr + offsets, mask=mask[:, None], other=0.0))
    mean = acc_sum / group_numel
    variance = acc_sq / group_numel - mean * mean
    variance = tl.maximum(variance, 0.0)
    reciprocal_std = tl.rsqrt(variance + eps)
    weight = tl.load(weight_ptr + cols).to(tl.float32)
    bias = tl.load(bias_ptr + cols).to(tl.float32)
    scale = weight * reciprocal_std
    shift = bias - mean * scale
    tl.store(scale_shift_ptr + batch * 2 * channels + cols, scale)
    tl.store(scale_shift_ptr + (batch * 2 + 1) * channels + cols, shift)
@triton.jit
 def _gn_apply_rows_kernel(
    x_ptr,
    scale_shift_ptr,
    y_ptr,
    rows,
    C: tl.constexpr,
    BLOCK_R: tl.constexpr,
    SILU: tl.constexpr,
 ):
    program = tl.program_id(0).to(tl.int64)
    batch = tl.program_id(1).to(tl.int64)
    cols = tl.arange(0, C)
    scale = tl.load(scale_shift_ptr + batch * 2 * C + cols)
    shift = tl.load(scale_shift_ptr + (batch * 2 + 1) * C + cols)
    row_indices = program * BLOCK_R + tl.arange(0, BLOCK_R)
    mask = row_indices < rows
    offsets = batch * rows * C + row_indices[:, None] * C + cols[None, :]
    x = tl.load(x_ptr + offsets, mask=mask[:, None], other=0.0).to(tl.float32)
    y = x * scale[None, :] + shift[None, :]
    if SILU:
        y = y * tl.sigmoid(y)
    tl.store(y_ptr + offsets, y.to(y_ptr.dtype.element_ty), mask=mask[:, None])
 def _group_norm_silu_rows(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    num_groups: int,
    eps: float,
    apply_silu: bool,
 ) -> torch.Tensor:
    batch_size, rows, channels = x.shape
    channels_per_group = channels // num_groups
    block_rows = max(1, 8192 // channels)
    rows_per_program = block_rows * 32
    num_chunks = triton.cdiv(rows, rows_per_program)
    partial_sum = torch.empty((batch_size, num_chunks, channels), device=x.device, dtype=torch.float32)
    partial_sq = torch.empty_like(partial_sum)
    with torch.cuda.device(x.device):
        _gn_partial_rows_kernel[(num_chunks, batch_size)](
            x,
            partial_sum,
            partial_sq,
            rows,
            rows_per_program,
            C=channels,
            BLOCK_R=block_rows,
            num_warps=4,
        )
        scale_shift = torch.empty((batch_size, 2, channels), device=x.device, dtype=torch.float32)
        block_chunks = max(1, min(4096 // channels_per_group, triton.next_power_of_2(num_chunks)))
        _gn_finalize_kernel[(num_groups, batch_size)](
            partial_sum,
            partial_sq,
            weight,
            bias,
            scale_shift,
            num_chunks,
            rows * channels_per_group,
            eps,
            channels,
            CPG=channels_per_group,
            BLOCK_K=block_chunks,
            num_warps=4,
        )
        output = torch.empty_like(x)
        _gn_apply_rows_kernel[(triton.cdiv(rows, block_rows), batch_size)](
            x,
            scale_shift,
            output,
            rows,
            C=channels,
            BLOCK_R=block_rows,
            SILU=apply_silu,
            num_warps=4,
        )
    return output
 def _is_supported(
    x: torch.Tensor,
    weight: torch.Tensor | None,
    bias: torch.Tensor | None,
    num_groups: int,
 ) -> bool:
    if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()):
        return False
    if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES:
        return False
    if weight is None or bias is None:
        return False
    channels = x.shape[1] if x.dim() == 4 else x.shape[-1]
    if weight.shape != (channels,) or bias.shape != (channels,):
        return False
    if not (
        weight.device == bias.device == x.device
        and weight.dtype == bias.dtype == x.dtype
        and weight.is_contiguous()
        and bias.is_contiguous()
    ):
        return False
    if num_groups < 1 or channels % num_groups != 0:
        return False
    return triton.next_power_of_2(channels) == channels and channels <= _MAX_CHANNELS
 def group_norm_silu_4d(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    num_groups: int,
    eps: float,
    apply_silu: bool = True,
 ) -> torch.Tensor | None:
    """Run GroupNorm(+SiLU) on a channels-last ``(N, C, H, W)`` tensor."""
    if x.dim() != 4 or not _is_supported(x, weight, bias, num_groups):
        return None
    batch_size, channels, height, width = x.shape
    if not (channels > 1 and (height > 1 or width > 1) and x.is_contiguous(memory_format=torch.channels_last)):
        return None
    rows = x.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
    output = _group_norm_silu_rows(rows, weight, bias, num_groups, eps, apply_silu)
    return output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
 def group_norm_silu_rows(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    num_groups: int,
    eps: float,
    apply_silu: bool = True,
 ) -> torch.Tensor | None:
    """Run GroupNorm(+SiLU) over contiguous ``(N, rows, C)`` input."""
    if x.dim() != 3 or not x.is_contiguous():
        return None
    if not _is_supported(x, weight, bias, num_groups):
        return None
    return _group_norm_silu_rows(x, weight, bias, num_groups, eps, apply_silu)
 __all__ = ["group_norm_silu_4d", "group_norm_silu_rows"]
@@ -0,0 +1,160 @@
 # ruff: noqa: N803
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Channels-last-3D Wan VAE RMSNorm(+SiLU) Triton kernel."""
 from __future__ import annotations
 import torch
 import triton
 import triton.language as tl
 _SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
 _MAX_CHANNELS = 1024
@triton.jit
 def _wan_rmsnorm_silu_kernel(
    x_ptr,
    gamma_ptr,
    bias_ptr,
    output_ptr,
    channels: tl.constexpr,
    time_size,
    height_size,
    width_size,
    x_stride_batch,
    x_stride_channel,
    x_stride_time,
    x_stride_height,
    x_stride_width,
    output_stride_batch,
    output_stride_channel,
    output_stride_time,
    output_stride_height,
    output_stride_width,
    rms_scale,
    eps,
    has_bias: tl.constexpr,
    BLOCK_C: tl.constexpr,
 ):
    row = tl.program_id(0).to(tl.int64)
    channel_offsets = tl.arange(0, BLOCK_C)
    channel_mask = channel_offsets < channels
    width = row % width_size
    remainder = row // width_size
    height = remainder % height_size
    remainder = remainder // height_size
    time = remainder % time_size
    batch = remainder // time_size
    x_base = batch * x_stride_batch + time * x_stride_time + height * x_stride_height + width * x_stride_width
    output_base = (
        batch * output_stride_batch
        + time * output_stride_time
        + height * output_stride_height
        + width * output_stride_width
    )
    x = tl.load(
        x_ptr + x_base + channel_offsets * x_stride_channel,
        mask=channel_mask,
        other=0.0,
    ).to(tl.float32)
    norm = tl.sqrt(tl.sum(x * x, axis=0))
    inverse_norm = 1.0 / tl.maximum(norm, eps)
    # Match the eager dtype boundaries: normalize and scale in activation
    # dtype, affine in the promoted dtype, then SiLU in FP32.
    y = (x * inverse_norm).to(x_ptr.dtype.element_ty)
    gamma = tl.load(gamma_ptr + channel_offsets, mask=channel_mask, other=1.0)
    y = (y * rms_scale).to(x_ptr.dtype.element_ty)
    y = (y.to(tl.float32) * gamma.to(tl.float32)).to(output_ptr.dtype.element_ty)
    if has_bias:
        bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0)
        y = (y.to(tl.float32) + bias.to(tl.float32)).to(output_ptr.dtype.element_ty)
    y = y.to(tl.float32)
    y = y * tl.sigmoid(y)
    tl.store(
        output_ptr + output_base + channel_offsets * output_stride_channel,
        y,
        mask=channel_mask,
    )
 def _affine_supported(x: torch.Tensor, affine: torch.Tensor) -> bool:
    return (
        affine.is_cuda
        and affine.device == x.device
        and (affine.dtype == x.dtype or affine.dtype == torch.float32)
        and affine.numel() == x.shape[1]
    )
 def can_use_wan_rmsnorm_silu(
    x: torch.Tensor,
    gamma: torch.Tensor,
    bias: torch.Tensor | None,
 ) -> bool:
    """Return whether ``wan_rmsnorm_silu`` supports this exact input."""
    return (
        x.is_cuda
        and not torch.is_grad_enabled()
        and not x.requires_grad
        and x.dtype in _SUPPORTED_DTYPES
        and x.ndim == 5
        and x.numel() > 0
        and 0 < x.shape[1] <= _MAX_CHANNELS
        and x.is_contiguous(memory_format=torch.channels_last_3d)
        and _affine_supported(x, gamma)
        and (bias is None or _affine_supported(x, bias))
    )
 def wan_rmsnorm_silu(
    x: torch.Tensor,
    gamma: torch.Tensor,
    bias: torch.Tensor | None = None,
    rms_scale: float | None = None,
    eps: float = 1e-12,
 ) -> torch.Tensor | None:
    """Fuse ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)``."""
    if not can_use_wan_rmsnorm_silu(x, gamma, bias):
        return None
    batch_size, channels, time_size, height_size, width_size = x.shape
    output_dtype = torch.promote_types(x.dtype, gamma.dtype)
    output = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=output_dtype)
    block_channels = triton.next_power_of_2(channels)
    num_warps = 1 if block_channels <= 64 else 4 if block_channels <= 512 else 8
    gamma = gamma.reshape(channels).contiguous()
    has_bias = bias is not None
    bias = gamma if bias is None else bias.reshape(channels).contiguous()
    if rms_scale is None:
        rms_scale = channels**0.5
    with torch.cuda.device(x.device):
        _wan_rmsnorm_silu_kernel[(batch_size * time_size * height_size * width_size,)](
            x,
            gamma,
            bias,
            output,
            channels,
            time_size,
            height_size,
            width_size,
            *x.stride(),
            *output.stride(),
            float(rms_scale),
            eps,
            has_bias,
            BLOCK_C=block_channels,
            num_warps=num_warps,
        )
    return output
 __all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"]
@@ -0,0 +1,162 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quality-gated CUDA fast path for Diffusers Wan VAE decoders."""
 from __future__ import annotations
 from types import MethodType
 from typing import Any
 import torch
 import torch.nn.functional as F
 from torch import nn
 from vllm.logger import init_logger
 from vllm_omni.diffusion.vae_optimizations.gate import (
    VaeFastPathGate,
    register_vae_fast_path_gate,
 )
 logger = init_logger(__name__)
 try:
    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import (
        wan_rmsnorm_silu,
    )
    _HAS_TRITON = True
 except ImportError:  # pragma: no cover - exercised on builds without Triton
    _HAS_TRITON = False
 class FusedWanRMSNormSiLU(nn.Module):
    """Quality-gated Wan RMSNorm + SiLU with an exact eager fallback."""
    def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None:
        super().__init__()
        self.gamma = norm.gamma
        self.bias = norm.bias
        self.scale = float(norm.scale)
        self._vllm_gate = gate
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._vllm_gate.enabled and not torch.compiler.is_compiling():
            bias = self.bias if isinstance(self.bias, torch.Tensor) else None
            output = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale)
            if output is not None:
                return output
        needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any(
            token in str(x.dtype) for token in ("float4_", "float8_")
        )
        normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype)
        return F.silu(normalized * self.scale * self.gamma + self.bias)
 def _plain_silu(activation: object) -> bool:
    return isinstance(activation, nn.SiLU) and not activation.inplace
 def _fusable_norm(norm: object, norm_cls: type) -> bool:
    return (
        type(norm) is norm_cls
        and getattr(norm, "channel_first", False)
        and isinstance(getattr(norm, "gamma", None), torch.Tensor)
    )
 def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResidualBlock, WanRMS_norm
    residual_blocks = [module for module in decoder.modules() if isinstance(module, WanResidualBlock)]
    eligible = [
        module
        for module in residual_blocks
        if type(module) is WanResidualBlock
        and _plain_silu(module.nonlinearity)
        and _fusable_norm(module.norm1, WanRMS_norm)
        and _fusable_norm(module.norm2, WanRMS_norm)
    ]
    if len(eligible) != len(residual_blocks):
        logger.info(
            "Wan VAE has %d/%d non-standard residual blocks; skipping fast path",
            len(residual_blocks) - len(eligible),
            len(residual_blocks),
        )
        return None
    if not (
        _fusable_norm(getattr(decoder, "norm_out", None), WanRMS_norm)
        and _plain_silu(getattr(decoder, "nonlinearity", None))
    ):
        logger.info("Wan VAE has a non-standard output head; skipping fast path")
        return None
    count = 0
    for module in eligible:
        module.norm1 = FusedWanRMSNormSiLU(module.norm1, gate)
        module.norm2 = FusedWanRMSNormSiLU(module.norm2, gate)
        module.nonlinearity = nn.Identity()
        count += 2
    decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate)
    decoder.nonlinearity = nn.Identity()
    return count + 1
 def _set_wan_decoder_memory_format(decoder: nn.Module, *, channels_last: bool) -> None:
    for module in decoder.modules():
        if isinstance(module, nn.Conv3d):
            module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format)
        elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
            module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format)
 def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
    use_channels_last = self._vllm_gate.enabled
    if use_channels_last != self._vllm_channels_last:
        _set_wan_decoder_memory_format(self, channels_last=use_channels_last)
        self._vllm_channels_last = use_channels_last
        logger.debug(
            "Wan VAE decoder switched to %s layout",
            "channels_last_3d" if use_channels_last else "contiguous",
        )
    return type(self).forward(self, *args, **kwargs)
 def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
    """Install the Wan decoder fast path when its full structure is supported."""
    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
        return vae
    from diffusers.models.autoencoders import AutoencoderKLWan
    from diffusers.models.autoencoders.autoencoder_kl_wan import WanDecoder3d
    if not isinstance(vae, AutoencoderKLWan):
        return vae
    decoder = getattr(vae, "decoder", None)
    if type(decoder) is not WanDecoder3d:
        return vae
    distributed_executor = getattr(vae, "distributed_executor", None)
    parallel_mode = getattr(distributed_executor, "parallel_mode", "tile")
    if parallel_mode in ("auto", "spatial_shard_height", "spatial_shard_width"):
        logger.info("Wan spatial-shard decode is configured; skipping quality-gated VAE fast path")
        return vae
    gate = VaeFastPathGate()
    norm_count = _install_norm_silu(decoder, gate)
    if norm_count is None:
        return vae
    decoder._vllm_gate = gate
    decoder._vllm_channels_last = False
    decoder.forward = MethodType(_decoder_layout_forward, decoder)
    register_vae_fast_path_gate(vae, gate)
    vae._vllm_vae_fast_path_installed = True
    logger.info(
        "Wan VAE installed quality-gated decoder fast path (%d RMSNorm+SiLU fusions)",
        norm_count,
    )
    return vae
 __all__ = ["FusedWanRMSNormSiLU", "maybe_optimize_wan_vae"]
@@ -43,6 +43,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
 from vllm_omni.diffusion.vae_optimizations import (
    optimize_pipeline_vaes,
    use_pipeline_vae_fast_path,
 )
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import (
@@ -243,6 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                    custom_pipeline_name=custom_pipeline_name,
                    device=self.device,
                )
                # Keep precision rematerialization in the tagged weights pool
                # so sleep level 1 can offload every resulting parameter.
                optimize_pipeline_vaes(self.pipeline, self.od_config)
        time_after_load = time.perf_counter()
        logger.info(
@@ -499,6 +506,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
        if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False):
            raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.")
        quality = getattr(reqs[0].sampling_params, "quality", None)
        if any(getattr(req.sampling_params, "quality", None) != quality for req in reqs[1:]):
            raise RuntimeError("A diffusion worker batch cannot mix request quality levels.")
        use_fast_vae = quality == "high"
        # Use no_grad() for HSDP compatibility, inference_mode() otherwise for
        # better perf. HSDP2's fully_shard pre-forward hooks need tensor version
        # counters, which inference tensors do not track.
@@ -523,6 +535,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
            with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
                with record_function(record_name):
                    with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae):
                        raw_outputs = self.pipeline.forward(batch)
                    outputs = _normalize_pipeline_outputs(
                        raw_outputs,
@@ -774,6 +787,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                            if should_decode:
                                clear_pipeline_stage_durations(self.pipeline)
                                with use_pipeline_vae_fast_path(
                                    self.pipeline,
                                    getattr(req.sampling, "quality", None) == "high",
                                ):
                                    result = self.pipeline.post_decode(req)
                                if result is not None:
                                    self._attach_stepwise_metrics(
@@ -576,6 +576,14 @@ class DiffusionWorker:
        usage_before = allocator.get_current_usage()
        if self.model_runner is not None and self.model_runner.pipeline is not None:
            from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
            # Folded weights are allocated lazily on the request path rather
            # than in the tagged weight pool. Drop them before sleep so they
            # neither remain resident nor need a separate CPU backup.
            clear_pipeline_vae_fast_path_caches(self.model_runner.pipeline)
        if level == 2 and self.model_runner is not None:
            if hasattr(self.model_runner, "graph_runners"):
                self.model_runner.graph_runners.clear()
@@ -491,6 +491,7 @@ class OrchestratorArgs:
    step_execution: bool = False
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    vae_decode_precision: str | None = None
    enable_multithread_weight_load: bool = True
    num_weight_load_threads: int = 4
    enable_cpu_offload: bool = False
@@ -1026,6 +1026,7 @@ class AsyncOmniEngine:
            "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
            "vae_use_slicing": kwargs.get("vae_use_slicing", False),
            "vae_use_tiling": kwargs.get("vae_use_tiling", False),
            "vae_decode_precision": kwargs.get("vae_decode_precision", None),
            "cache_backend": cache_backend,
            "cache_config": cache_config,
            "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
@@ -629,6 +629,15 @@ class OmniServeCommand(CLISubcommand):
            action="store_true",
            help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
        )
        omni_config_group.add_argument(
            "--vae-decode-precision",
            choices=("fp32", "fp16", "bf16"),
            default=None,
            help=(
                "Optional Wan decode-only VAE precision. The encoder remains FP32; "
                "unset preserves the pipeline's existing precision behavior."
            ),
        )
        # Parallel weight loading (faster diffusion startup)
        omni_config_group.add_argument(
@@ -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	094abd70c1	feat(diffusion): add request-gated VAE fast paths Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 09:31:27 +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