[Perf][Diffusion] Compile and prewarm VAE decode

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Bugfix] rainfusion support all shape video (#6000 )
2026-08-11 07:14:35 +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
49 changed files with 1368 additions and 297 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.
@@ -4,6 +4,17 @@ This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
 resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
 recipes it uses **no offload of any kind** — see the capacity note below.
 Validated on:
 - Host: DGX Spark (GB10), aarch64
 - GPUs: 1 (unified memory)
 - Driver: <FILL_DRIVER_VERSION>
 - vLLM: 0.26.0
 - vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
 - Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
  `seed=1101`
 ## Capacity requirements
 | Resource | DGX Spark (GB10) |
@@ -11,7 +22,8 @@ recipes it uses **no offload of any kind** — see the capacity note below.
 | Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
 | GPUs | 1 |
 | Checkpoint storage | 135 GiB per partition |
-| Measured peak allocator high-water (FP8) | 97.7 GiB |
+| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
 | Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
@@ -29,9 +41,21 @@ math relative to the RTX and datacenter recipes:
  fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
  Qwen3-VL text encoder and both VAEs stay BF16.
-With ~97.7 GiB of the pool held by the allocator at peak, there is roughly 23 GiB
+With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
-of headroom for the OS, page cache, and any other process on the box. Treat this
+23 GiB of headroom for the OS, page cache, and any other process on the box.
-platform as single-tenant while a request is in flight.
+Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
 and the allocator figure understates true pool usage.
 Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
 `MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
 a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
 roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
 which is the CUDA context, non-PyTorch allocations, and page cache the header
 does not count.
 Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
 second process, and adding reference images, reference videos, or a larger output
 shape will run into the OOM killer.
 ## One GB10: 960x576, 8 seconds
@@ -94,6 +118,52 @@ ffprobe -v error -show_entries \
 Both requests return `200 OK` with an MP4 body: H.264 video at the requested
 geometry plus a 32 kHz stereo AAC track.
 ## One GB10: Ref2VA, 960x576, 8 seconds
 `FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
 Stop the FL2VA server before starting this one.
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
 FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
 vllm serve /path/to/MiniMax-H3/Ref2VA \
  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
  --init-timeout 3600 \
  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
  --vae-parallel-mode tile --vae-use-tiling \
  --quantization fp8 --enforce-eager \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
 takes over an hour on this platform. Use 14400 s.
 Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
 omitted when `width` and `height` are supplied.
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -D ref2va-headers.txt \
  -w 'client_e2e_s=%{time_total}\n' \
  --max-time 14400 \
  -F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
  -F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-ref2va.mp4
 ```
 The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
 32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
 ## Measured latency
 Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
@@ -102,32 +172,69 @@ Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
 | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
 | T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
 | T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
 | Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
 | Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
 How these were measured:
 - Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
-  `X-Stage-Durations` response header of `/v1/videos/sync`.
+  `X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
  `MiniMaxH3Pipeline.encode_prompt` (text encode),
  `MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
  (VAE decode).
 - MP4 mux is not a header field. It is the
  `Video response encoding (MP4 bytes)` line in the server log.
 - Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
  high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
  the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
  higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
-  memory has no discrete VRAM to report and the field returns `[N/A]`.
+  memory has no discrete VRAM to report and the field returns `[N/A]`. See the
- Text encode and VAE decode were measured on a separate 10-step run at the same
+  capacity section for a `/proc/meminfo` cross-check.
-  960x576 / 8 s geometry. Both are independent of step count, so they transfer to
+- For T2VA, text encode and VAE decode were measured on a separate 10-step run at
-  the 50-step row unchanged.
+  the same 960x576 / 8 s geometry. Both are independent of step count, so they
  transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
  same way, from the 10-step run at identical output geometry and duration.
 - Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
  (`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
  geometry and reference input rather than step count.
 - Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
  engine bookkeeping sit outside the profiled stages.
 Denoise dominates at 96% of end-to-end, so step count and geometry are the only
 knobs that materially change the wall clock on this platform.
-Repeated identical 10-step requests varied by up to 23% in denoise time
+Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
 4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
 versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
 The cost is in denoise. A single reference image expands the Qwen3-VL
 presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
 6883 tokens, 1 reference images`), and those tokens sit in the attention context
 of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
 1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
 expected for an identical output shape.
 Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
 (397 s to 490 s) on a machine that had already been under sustained load. Treat
 these as single-run measurements on one machine, not a throughput guarantee, and
 expect the slower end of that range for back-to-back jobs.
 Ref2VA per-step time is load-dependent in the same way. On an idle machine the
 first 10-step request ran at 66.55 s per step. The 50-step request that followed
 occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
 further 10-step requests taken immediately afterwards settled at 74.88 s and
 74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
 is a stable thermal state rather than measurement noise.
 Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
 only reachable on a cold machine and does not hold for the length of one request.
 Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
 8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
 VAE decode was stable at 67.9-69.8 s throughout.
 Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
-default 1800 s is shorter than a 50-step run on this platform.
+default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
 shorter than a 50-step Ref2VA run.
 ## Known limitations
@@ -139,11 +246,17 @@ default 1800 s is shorter than a 50-step run on this platform.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
  `height` are both supplied they take precedence and `aspect_ratio` only has to
  name one of the supported ratios.
- Sustained load degrades throughput. See the variance note above.
+- Sustained load degrades throughput. See the variance notes above.
 - Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
-  the same command with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video
+  with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
-  count and prompt length increase activation memory; begin with one request at
+  above; reference videos and additional images push the token count higher and
-  a time and keep the short edge at 576.
+  were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
  single-image request, there is very little room to grow — re-measure before
  adding references or raising the output shape.
 - Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
  under **Validated on**: FP8 weight loading and image-only Ref2VA are both
  broken on `release/v0.26.0`.suggestion is  `v0.26.1`.
 - Online FP8 is incompatible with layerwise offload — the offload path produces a
  weight stride the Cutlass FP8 kernel rejects. This is not a practical
  restriction here since offload is unusable on GB10 anyway.
 ```
@@ -751,6 +751,7 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
                "  - stage_id: 0",
                "    diffusion_compile_granularity: regional",
                "    diffusion_compile_dynamic: false",
                "    diffusion_compile_vae: false",
            ]
        )
    )
@@ -762,13 +763,16 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
        cli_overrides={
            "diffusion_compile_granularity": "full",
            "diffusion_compile_dynamic": True,
            "diffusion_compile_vae": True,
        },
    ).stage_by_id(0)
    assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional"
    assert configured_stage.diffusion_config.diffusion_compile_dynamic is False
    assert configured_stage.diffusion_config.diffusion_compile_vae is False
    assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
    assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
    assert overridden_stage.diffusion_config.diffusion_compile_vae is True
 def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
@@ -776,6 +780,11 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
        omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 def test_structured_diffusion_config_rejects_non_boolean_compile_vae():
    with pytest.raises(ValidationError, match="diffusion_compile_vae"):
        omni_config_module._DiffusionConfigProjection(diffusion_compile_vae="false")
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():
    with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
        omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
@@ -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
@@ -69,6 +69,7 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    attention = object.__new__(MiniMaxH3Attention)
    nn.Module.__init__(attention)
    attention.rot_dim = 96
    attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
    attention.rope._forward_method = attention.rope.forward_native
@@ -5,7 +5,7 @@ import pytest
 import torch.nn as nn
 import vllm_omni.diffusion.compile as compile_module
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -74,3 +74,50 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
        regionally_compile(model, dynamic=True)
    assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
 def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
    compiled_calls = 0
    eager_calls = 0
    def eager(value):
        nonlocal eager_calls
        eager_calls += 1
        return value + 1
    def compile_fn(fn, **kwargs):
        assert fn is eager
        assert kwargs == {"mode": "default"}
        def compiled(value):
            nonlocal compiled_calls
            compiled_calls += 1
            raise RuntimeError("lazy failure")
        return compiled
    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")
    assert activated
    assert callable_fn(1) == 2
    assert callable_fn(2) == 3
    assert compiled_calls == 1
    assert eager_calls == 2
    assert "disabling compilation and retrying eagerly" in caplog.text
 def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
    def eager(value):
        return value + 1
    def compile_fn(*args, **kwargs):
        raise RuntimeError("setup failure")
    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")
    assert not activated
    assert callable_fn is eager
    assert callable_fn(1) == 2
    assert "continuing eagerly" in caplog.text
@@ -515,16 +515,19 @@ class TestDiffusionCompileConfig:
        assert config.diffusion_compile_granularity == "regional"
        assert config.diffusion_compile_dynamic is True
        assert config.diffusion_compile_vae is True
    def test_from_kwargs_preserves_compile_controls(self) -> None:
        config = OmniDiffusionConfig.from_kwargs(
            model="test",
            diffusion_compile_granularity="full",
            diffusion_compile_dynamic=False,
            diffusion_compile_vae=False,
        )
        assert config.diffusion_compile_granularity == "full"
        assert config.diffusion_compile_dynamic is False
        assert config.diffusion_compile_vae is False
    def test_config_rejects_invalid_compile_granularity(self) -> None:
        with pytest.raises(ValueError, match="diffusion_compile_granularity"):
@@ -534,6 +537,10 @@ class TestDiffusionCompileConfig:
        with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
            OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
    def test_config_rejects_non_boolean_compile_vae(self) -> None:
        with pytest.raises(TypeError, match="diffusion_compile_vae"):
            OmniDiffusionConfig(model="test", diffusion_compile_vae="false")
    @pytest.mark.parametrize(
        "kwargs, feature",
        [
@@ -130,6 +130,23 @@ class _CompileTrackingModel:
        return self
 class _DecodeVAE(torch.nn.Module):
    def decode(self, value):
        return value + 1
 class _PipelineWithVAE(torch.nn.Module):
    _dit_modules = []
    _encoder_modules = []
    _vae_modules = ["vae", "vae_alias"]
    _resident_modules = []
    def __init__(self):
        super().__init__()
        self.vae = _DecodeVAE()
        self.vae_alias = self.vae
 def _make_request():
    sampling_params = SimpleNamespace(
        generator=None,
@@ -320,6 +337,31 @@ def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypa
    assert "lazy compilation errors" in caplog.text
@pytest.mark.core_model
@pytest.mark.cpu
 def test_compile_vae_decodes_uses_discovery_and_deduplicates_aliases(monkeypatch):
    runner = object.__new__(DiffusionModelRunner)
    runner.pipeline = _PipelineWithVAE()
    compile_calls = []
    def compile_fn(fn, **kwargs):
        compile_calls.append((fn, kwargs))
        return lambda *args, **call_kwargs: fn(*args, **call_kwargs) + 10
    monkeypatch.setattr(model_runner_module.torch, "compile", compile_fn)
    DiffusionModelRunner._compile_vae_decodes(runner)
    assert len(compile_calls) == 1
    assert compile_calls[0][1] == {
        "mode": "default",
        "fullgraph": False,
        "dynamic": None,
    }
    assert runner.pipeline.vae.decode(1) == 12
    assert runner.pipeline.vae_alias.decode is runner.pipeline.vae.decode
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
@@ -971,6 +971,44 @@ class TestDiffusionEngine:
        admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
        assert admitted_request.request_id == "dummy_req_id"
    def test_compile_warmup_repeats_through_normal_request_path(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine.od_config = SimpleNamespace(
            model_class_name="mock_model",
            diffusion_load_format="default",
            diffusion_compile_vae=True,
            enforce_eager=False,
        )
        engine.add_req_and_wait_for_response = mocker.Mock(
            side_effect=[DiffusionOutput(output=None), DiffusionOutput(output=None)]
        )
        engine._dummy_run()
        requests = [call.args[0] for call in engine.add_req_and_wait_for_response.call_args_list]
        assert len(requests) == 2
        assert requests[0].is_dummy_run()
        assert not requests[1].is_dummy_run()
        assert requests[1].prompt == {"prompt": "A detailed image."}
        assert requests[1].sampling_params.need_kv_receive is False
        assert requests[1].sampling_params.num_inference_steps == 1
    def test_compile_real_path_warmup_failure_is_non_fatal(self, mocker: MockerFixture, caplog) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine.od_config = SimpleNamespace(
            model_class_name="mock_model",
            diffusion_load_format="default",
            diffusion_compile_vae=True,
            enforce_eager=False,
        )
        engine.add_req_and_wait_for_response = mocker.Mock(
            side_effect=[DiffusionOutput(output=None), RuntimeError("prewarm failure")]
        )
        engine._dummy_run()
        assert "serving will continue" in caplog.text
 class TestStepScheduler:
    def setup_method(self) -> None:
@@ -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
@@ -87,6 +87,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
            "enforce_eager": None,
            "diffusion_compile_granularity": None,
            "diffusion_compile_dynamic": None,
            "diffusion_compile_vae": None,
        }
    )[0]
@@ -98,6 +99,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
    assert stage_cfg["engine_args"]["enforce_eager"] is False
    assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
    assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is True
 def test_default_stage_config_propagates_ulysses_mode():
@@ -315,7 +317,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
 def test_serve_cli_accepts_diffusion_compile_controls():
-    """Ensure both compile controls reach the diffusion stage."""
+    """Ensure all compile controls reach the diffusion stage."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)
@@ -328,6 +330,7 @@ def test_serve_cli_accepts_diffusion_compile_controls():
            "--diffusion-compile-granularity",
            "full",
            "--no-diffusion-compile-dynamic",
            "--no-diffusion-compile-vae",
        ]
    )
@@ -336,8 +339,10 @@ def test_serve_cli_accepts_diffusion_compile_controls():
    assert args.diffusion_compile_granularity == "full"
    assert args.diffusion_compile_dynamic is False
    assert args.diffusion_compile_vae is False
    assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full"
    assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is False
 def test_serve_cli_accepts_diffusion_attention_backend():
@@ -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
@@ -570,6 +570,7 @@ class _DiffusionConfigProjection:
    pin_cpu_memory: bool = True
    diffusion_compile_granularity: Literal["regional", "full"] = "regional"
    diffusion_compile_dynamic: bool = Field(default=True, strict=True)
    diffusion_compile_vae: bool = Field(default=True, strict=True)
    fa_deterministic: bool = False
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
@@ -394,6 +394,7 @@ class StageDeployConfig:
    # Diffusion execution, cache, and VAE behavior.
    diffusion_compile_granularity: str | None = None
    diffusion_compile_dynamic: bool | None = None
    diffusion_compile_vae: bool | None = None
    fa_deterministic: bool | None = None
    cache_backend: str | None = None
    cache_config: dict[str, Any] | 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 "
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Callable
 from typing import Any
 import torch
@@ -9,6 +10,59 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 class _FallbackCompiledCallable:
    """Run a compiled callable until it fails, then permanently use eager."""
    def __init__(
        self,
        eager: Callable[..., Any],
        compiled: Callable[..., Any],
        *,
        label: str,
    ) -> None:
        self._eager = eager
        self._compiled: Callable[..., Any] | None = compiled
        self._label = label
    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        compiled = self._compiled
        if compiled is None:
            return self._eager(*args, **kwargs)
        try:
            return compiled(*args, **kwargs)
        except Exception as exc:
            self._compiled = None
            logger.warning(
                "%s failed during lazy torch.compile execution (%s); disabling compilation and retrying eagerly.",
                self._label,
                exc,
            )
            return self._eager(*args, **kwargs)
 def compile_callable_with_fallback(
    fn: Callable[..., Any],
    *,
    label: str,
    **compile_kwargs: Any,
 ) -> tuple[Callable[..., Any], bool]:
    """Compile ``fn`` with a fail-closed eager fallback.
    The boolean return value reports whether lazy compilation was activated.
    Synchronous setup failures return the original callable unchanged.
    """
    try:
        compiled = torch.compile(fn, **compile_kwargs)
    except Exception as exc:
        logger.warning(
            "%s torch.compile setup failed (%s); continuing eagerly.",
            label,
            exc,
        )
        return fn, False
    return _FallbackCompiledCallable(fn, compiled, label=label), True
 def _matches_repeated_block(
    name: str,
    module: nn.Module,
@@ -726,6 +726,9 @@ class OmniDiffusionConfig:
    # provide its own setup_compile() implementation.
    diffusion_compile_granularity: str = "regional"
    diffusion_compile_dynamic: bool = True
    # Compile discovered VAE ``decode`` callables whenever generic diffusion
    # torch.compile is enabled. Runtime failures permanently fall back to eager.
    diffusion_compile_vae: bool = True
    # Parallel weight loading (for faster diffusion model startup)
    enable_multithread_weight_load: bool = True
@@ -930,6 +933,8 @@ class OmniDiffusionConfig:
            )
        if not isinstance(self.diffusion_compile_dynamic, bool):
            raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
        if not isinstance(self.diffusion_compile_vae, bool):
            raise TypeError(f"diffusion_compile_vae must be a bool, got {type(self.diffusion_compile_vae)!r}")
        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
        if self.omni_kv_config is None:
@@ -56,6 +56,7 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 _ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
 _VAE_COMPILE_WARMUP_REQUEST_ID = "vae_compile_warmup_req_id"
 __all__ = [
    "DiffusionEngine",
@@ -851,28 +852,26 @@ class DiffusionEngine:
        num_inference_steps = 1
        height = 512
        width = 512
        prompt: OmniTextPrompt = {"prompt": "dummy run"}
        supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
        if supports_image_input:
            # Provide a dummy image input if the model supports it
            color_format = image_color_format(self.od_config.model_class_name)
            dummy_image = PIL.Image.new(color_format, (width, height))
            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
        if supports_audio_input:
            audio_sr = 16000
            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
        num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
        if num_frames <= 0:
            logger.info("Skipping dummy warmup run (num_frames=0)")
            return
-        req = OmniDiffusionRequest(
+
-            prompt=prompt,
+        def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest:
-            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            prompt: OmniTextPrompt = {"prompt": prompt_text}
-            sampling_params=OmniDiffusionSamplingParams(
+            if supports_image_input:
                color_format = image_color_format(self.od_config.model_class_name)
                prompt.setdefault("multi_modal_data", {})["image"] = PIL.Image.new(
                    color_format,
                    (width, height),
                )
            if supports_audio_input:
                audio_sr = 16000
                prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
            sampling_params = OmniDiffusionSamplingParams(
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
@@ -885,13 +884,47 @@ class DiffusionEngine:
                # Disable CFG for warmup to avoid triggering CFG parallel
                # validation when cfg_parallel_size > 1.
                extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
-            ),
+            )
-        )
+            if real_path:
                # Internal pretraffic must exercise ordinary pipeline branches
                # without waiting for KV produced by an upstream stage.
                sampling_params.need_kv_receive = False
            return OmniDiffusionRequest(
                prompt=prompt,
                request_id=request_id,
                sampling_params=sampling_params,
            )
        req = build_request(DUMMY_DIFFUSION_REQUEST_ID, "dummy run", real_path=False)
        logger.info("dummy run to warm up the model")
        output = self.add_req_and_wait_for_response(req)
        if output.error:
            raise RuntimeError(f"Dummy run failed: {output.error}")
        compile_real_path = bool(
            getattr(self.od_config, "diffusion_compile_vae", False)
            and not getattr(self.od_config, "enforce_eager", True)
        )
        if not compile_real_path:
            return
        logger.info("real request-path warmup for compiled VAE decode")
        real_path_req = build_request(
            _VAE_COMPILE_WARMUP_REQUEST_ID,
            "A detailed image.",
            real_path=True,
        )
        try:
            real_path_output = self.add_req_and_wait_for_response(real_path_req)
        except Exception as exc:
            logger.warning("Real request-path compile warmup failed (%s); serving will continue.", exc)
            return
        if real_path_output.error:
            logger.warning(
                "Real request-path compile warmup returned an error (%s); serving will continue.",
                real_path_output.error,
            )
    def _submit_rpc(
        self,
        method: str,
@@ -68,6 +68,7 @@ class RMSNorm(CustomOp):
    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
        self.hidden_size = hidden_size
        self.variance_epsilon = eps
    def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
@@ -107,6 +108,19 @@ class RMSNorm(CustomOp):
        except Exception:
            return self.forward_native(x)
    def forward_musa(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
        # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
        # H3 Q/K norm with its inline RoPE path on MUSA.
        return F.rms_norm(
            x,
            (self.hidden_size,),
            self.weight,
            self.variance_epsilon,
        )
    def forward_npu(
        self,
        x: torch.Tensor,
@@ -193,7 +193,16 @@ class RotaryEmbedding(CustomOp):
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(x, cos, sin)
+        # Keep the full-dimension NeoX form inline for MUSA compilation.
        # Other layouts use the shared native implementation.
        if self.half_head_dim or self.interleaved:
            return self.forward_native(x, cos, sin)
        if cos.dim() == 3:
            cos, sin = cos[0], sin[0]
        cos = cos.unsqueeze(-2)
        sin = sin.unsqueeze(-2)
        x1, x2 = x.chunk(2, dim=-1)
        return x * cos + torch.cat((-x2, x1), dim=-1) * sin
    def forward_native(
        self,
@@ -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:
@@ -333,6 +342,7 @@ class MiniMaxH3Attention(nn.Module):
        )
        self.num_heads = self.qkv_proj.num_heads
        self.num_kv_heads = self.qkv_proj.num_kv_heads
        self.rot_dim = 6 * arch.rope_inv_freq_len
        self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
        self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
        self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
@@ -365,7 +375,7 @@ class MiniMaxH3Attention(nn.Module):
        x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
        are cast to the activation dtype before the elementwise math.
        """
-        rot_dim = freqs.shape[-1]
+        rot_dim = self.rot_dim
        x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
        cos = torch.cos(freqs).to(x.dtype)  # [T, rot_dim]
        sin = torch.sin(freqs).to(x.dtype)
@@ -401,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,
@@ -415,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,
        )
@@ -28,7 +28,7 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
    resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.interface import (
    supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 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
@@ -188,6 +189,39 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
            compile_dynamic,
        )
    def _compile_vae_decodes(self) -> None:
        """Compile declared VAE decode callables with eager fallback."""
        assert self.pipeline is not None
        modules = ModuleDiscovery.discover(self.pipeline)
        compiled_count = 0
        for name, vae in zip(modules.vae_names, modules.vaes, strict=True):
            decode = getattr(vae, "decode", None)
            if not callable(decode):
                logger.debug("Model runner: VAE component %s has no decode callable; skipping compile.", name)
                continue
            label = f"VAE component {name}.decode"
            compiled_decode, activated = compile_callable_with_fallback(
                decode,
                label=label,
                mode="default",
                fullgraph=False,
                dynamic=None,
            )
            if not activated:
                continue
            setattr(vae, "decode", compiled_decode)
            compiled_count += 1
        if compiled_count:
            logger.info(
                "Model runner: configured lazy torch.compile for %d VAE decode callable(s); "
                "startup warmup will compile the real request path.",
                compiled_count,
            )
        else:
            logger.info("Model runner: no compile-compatible VAE decode callable was discovered.")
    def load_model(
        self,
        memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None,
@@ -277,7 +311,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
        # Apply torch.compile if not in eager mode
        if not self.od_config.enforce_eager:
            if current_omni_platform.supports_torch_inductor():
-                if hasattr(self.pipeline, "setup_compile"):
+                pipeline_owns_compile = hasattr(self.pipeline, "setup_compile")
                if pipeline_owns_compile:
                    try:
                        self.pipeline.setup_compile()
                    except Exception as exc:
@@ -291,6 +326,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                        transformer_attrs = ("transformer", "transformer_2")
                    for attr_name in transformer_attrs:
                        self._compile_transformer(attr_name)
                if self.od_config.diffusion_compile_vae and not pipeline_owns_compile:
                    self._compile_vae_decodes()
            else:
                logger.warning(
                    "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",
@@ -485,6 +485,7 @@ class OrchestratorArgs:
    diffusion_attention_config: str | None = None
    diffusion_compile_granularity: str | None = None
    diffusion_compile_dynamic: bool | None = None
    diffusion_compile_vae: bool | None = None
    cache_backend: str = "none"
    cache_config: str | None = None
    enable_cache_dit_summary: bool = False
@@ -1043,6 +1043,9 @@ class AsyncOmniEngine:
            "diffusion_compile_dynamic": (
                True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
            ),
            "diffusion_compile_vae": (
                True if kwargs.get("diffusion_compile_vae") is None else kwargs["diffusion_compile_vae"]
            ),
            "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
            "boundary_ratio": kwargs.get("boundary_ratio", None),
            "flow_shift": kwargs.get("flow_shift", None),
@@ -446,6 +446,15 @@ class OmniServeCommand(CLISubcommand):
                "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
            ),
        )
        omni_config_group.add_argument(
            "--diffusion-compile-vae",
            action=argparse.BooleanOptionalAction,
            default=None,
            help=(
                "Compile discovered VAE decode callables when diffusion torch.compile is enabled. "
                "Use --no-diffusion-compile-vae to keep VAE decode eager."
            ),
        )
        omni_config_group.add_argument(
            "--fa-deterministic",
            dest="fa_deterministic",
@@ -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	1cdaddd79b	[Perf][Diffusion] Compile and prewarm VAE decode Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 07:14:35 +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
R0CKSTAR	636d44e272	[Perf][MUSA] Restore MiniMax-H3 dynamic RoPE fusion (#5881 ) Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>	2026-08-10 17:24:46 +08:00
yiminghub2024	8e44e78038	Add Ref2VA measurements to the MiniMax-H3 DGX Spark (GB10) recipe (#5972 ) Signed-off-by: yiminghub2024 <482890@qq.com>	2026-08-10 15:51:15 +08:00
Hongsheng Liu	b98b7c8529	[Docs] Preserve generated quantization URLs and link recipes in supported models (#5969 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-10 15:45:50 +08:00