docs(review): enforce configuration provenance

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-09 22:39:57 +08:00
136 changed files with 1954 additions and 8475 deletions
@@ -404,7 +404,7 @@ steps:
      - label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
        timeout_in_minutes: 90
        commands:
-          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -m "full_model and cuda" --run-level "full_model"
        mirror_hardwares: h100_1
      - label: ":full_moon: Diffusion X2V · Doc Test"
@@ -94,6 +94,7 @@ If docs and live code disagree, verify the code/tests and report the drift.
 | --- | --- |
 | [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, or stage config is added. |
 | [perf-verification.md](references/checks/perf-verification.md) | The PR makes a latency, throughput, memory, or quality claim. |
 | [config-provenance.md](references/checks/config-provenance.md) | Changed or newly load-bearing behavior reads config, CLI, environment, legacy, or default values. |
 | [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
 | [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
 | [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
@@ -152,13 +153,19 @@ contract they protect or use only the applicable evidence checks.
 Apply every category in [general-checks.md](references/process/general-checks.md) before
 lower-priority comments.
-For each changed value or behavior, trace:
+For each changed or newly load-bearing value or behavior, trace:
 ```text
 public ingress -> validation/defaulting -> producer -> transformations
  -> stage/worker/connector boundary -> final consumer -> terminal cleanup
 ```
 Treat unchanged context as in scope when the diff makes it authoritative for a
 new consumer. When config, CLI, environment, legacy, or default sources can
 control that value, apply
 [config-provenance.md](references/checks/config-provenance.md) and prove the
 documented source at the final sink with distinct sentinel values.
 Cover every applicable offline/online, streaming/non-streaming, sync/async,
 feature-on/off, topology, and compatibility path. Search bounded callers and
 sibling implementations rather than assuming the changed hunk is the only path.
@@ -0,0 +1,49 @@
 # Configuration Provenance
 Read this reference when changed or newly load-bearing behavior reads config,
 CLI, environment, legacy, alias, or default values. Always use it for endpoint,
 host, port, address, connector metadata, and `*_extra_config` consumers, even
 when the schema or documentation is unchanged.
 ## Trace the source of truth
 For each semantic value, record one compact ledger:
 ```text
 documented input -> parsed field -> competing sources and precedence
  -> runtime projection/transport -> final consumer
 ```
 - Search the exact field, environment variable, default literal, and aliases in
  code, tests, docs, examples, and deployment files.
 - Treat unchanged producers and context lines as review scope when the diff
  starts relying on them or makes their output authoritative.
 - Resolve the intended precedence from active branch-local docs, enforced
  code/tests, and compatibility policy. Report a documented input that is
  silently ignored or reinterpreted at a live sink.
 - Check every supported construction path and feature-off behavior without
  inventing parity for unsupported paths.
 ## Require discriminating evidence
 Give every competing source a distinct value so the test reveals which source
 wins. Prefer non-default sentinels; for example:
 ```text
 documented config port = 25301
 environment/default port = 8998
 expected final endpoint = http://10.0.0.8:25301
 ```
 Assert the final URL, connection target, runtime argument, or other live sink,
 not only an intermediate object. A test that changes the environment value to
 the expected configured value proves formatting, not configuration propagation.
 For a bug fix, show the discriminating test fails on the frozen base and passes
 on the reviewed head when the environment permits execution.
 ## Finding bar
 A supported non-default value reaching the wrong endpoint, device, topology, or
 runtime behavior is a correctness or compatibility finding. A missing
 discriminating test alone is a validation gap unless repository policy requires
 that regression coverage for the change.
@@ -16,6 +16,9 @@ were reverted or broken. Check:
  instead of the changed behavior being mocked away;
 - fakes preserve relevant types, MRO, shapes, devices, async behavior, and
  lifecycle transitions;
 - configuration tests give competing config, CLI, environment, legacy, and
  default sources distinct non-default values and assert the final consumer;
  setting only the source read by the implementation does not prove precedence;
 - seeds, ordering, timing, synchronization, external services, and numeric
  tolerances are deterministic or explicitly controlled;
 - normal, invalid, boundary, feature-off, failure/cancellation, and regression
@@ -15,6 +15,9 @@ and [design documents](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/
 - Trace each changed value through public ingress, validation/defaulting,
  producer, transformations, process or device boundaries, final consumer, and
  terminal cleanup.
 - Treat unchanged code as in scope when the diff makes its value newly
  load-bearing for a changed consumer; validate that value's origin and contract
  rather than reviewing only added lines.
 - Check every applicable sync/async, offline/online, streaming/non-streaming,
  feature-off, topology, and compatibility path without demanding unsupported
  modes merely for symmetry.
@@ -123,7 +126,9 @@ distinct live caller, invariant, or compatibility need.
 ## Finding bar
-Anchor each finding to a changed `path:line`; name the trigger or call path,
+Anchor each finding to the exact `path:line` that causes the behavior; it may be
-current behavior, impact, and smallest fix direction. Treat pending CI, missing
+unchanged when the diff makes it newly load-bearing, in which case name the
-hardware, or unsupported measurements as validation gaps unless the repository
+changed call path that brings it into scope. Include the trigger, impact, and
-contract makes that evidence a merge requirement.
+smallest fix direction. Treat pending CI, missing hardware, or unsupported
 measurements as validation gaps unless the repository contract makes that
 evidence a merge requirement.
@@ -166,6 +166,10 @@ index patch, worktree patch, and untracked-content manifest. Any mismatch makes
 the review stale; discard the affected evidence and restart. Prefer one
 root-cause comment to several symptom comments.
 When the root cause is an unchanged line made newly load-bearing by the diff,
 cite it in the consolidated review body instead of moving the comment to an
 unrelated changed line merely to make it inline.
 Local presentation is the default. Only explicit authorization permits GitHub
 posting. When authorized, publish one consolidated final review after
 validation; do not post preliminary or incremental comments. Submit a review
@@ -50,6 +50,7 @@ compatibility changes.
 | --- | --- | --- |
 | New or expanded model, loader, processor, registry, or stage config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) |
 | Latency, throughput, memory, scaling, precision, or quality claim | [perf-verification.md](../checks/perf-verification.md) | [`diffusion-perf-opt`](../../../diffusion-perf-opt/SKILL.md) for diffusion |
 | Config/CLI/environment/default precedence; endpoint, host, port, address, connector metadata, or `*_extra_config` | [config-provenance.md](../checks/config-provenance.md) | None |
 | Tests changed, absent for risky behavior, or test-only | [test-quality-evaluation.md](../checks/test-quality-evaluation.md) | [`vllm-omni-test`](../../../vllm-omni-test/SKILL.md) |
 | CI, examples, docs, public behavior, or contributor evidence | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | None |
 | Suitable hardware/server and runnable affected path | [verification.md](../checks/verification.md) | None |
@@ -1,11 +1,6 @@
 # CODEOWNERS is used to automatically assign PR reviewers
 # Format: path pattern @username or @team
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 #
 # The vllm_omni/** source-tree rules below are derived from the `owners` and
 # `primary_code_paths` front matter in docs/design/module/*.md. When updating
 # ownership, update the design doc first, then regenerate the matching rule here.
 # Last-match-wins: more specific rules must come AFTER broader ones.
 # Default owners (fallback reviewers for all files)
 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -26,10 +21,25 @@
 # CI
 /.buildkite/ @yenuo26 @congw729 @NickCao
-# ComfyUI / external apps
+# Engine and model_executor
-/apps/ @fhfuih @alex-jw-brooks
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
 /vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
 /vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/worker/ @tzhouam @gcanlin
-# Benchmarks (out-of-tree)
+# Diffusion
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
 # Diffusion CPU offloading
 /vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
 # Omni and TTS support
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
 /vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
 # Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
 /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -37,112 +47,20 @@
 /benchmarks/glm_image/ @princepride @JaredforReal
 /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
-# =====================================================================
+# Usability and Observability
-# Module design ownership (docs/design/module/*.md)
+/vllm_omni/deploy/ @tzhouam @lishunyang12
-# Each section header cites the source design doc.
+/vllm_omni/metrics/ @lishunyang12
 # =====================================================================
 # Autoregressive runtime — docs/design/module/ar_runtime.md
 /vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
 /vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
 # Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
 /vllm_omni/engine/ @tzhouam @fake0fan
 # File-level overrides: input/output modality contracts own these engine files
 # — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
 /vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 /vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
 # Entrypoints and serving boundaries — docs/design/module/entrypoints.md
 /vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
 # File-level overrides within entrypoints/
 /vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
 /vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
 # Error contracts — docs/design/module/error_contracts.md
 /vllm_omni/errors.py @alex-jw-brooks @NickCao
 # OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
 # Model integration — docs/design/module/model_integration.md
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
 /vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
 /vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
 # Config override within model_executor — docs/design/module/vllm_omni_config.md
 /vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
 # Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
 /vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
 /vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
 # Configuration & deploy — docs/design/module/vllm_omni_config.md
 /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
 # Quantization — docs/design/module/quantization.md
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
 # Observability — docs/design/module/observability.md
 /vllm_omni/metrics/ @lishunyang12 @vraiti
 # Profiling — docs/design/module/profiling.md
 /vllm_omni/profiler/ @gcanlin
-# Benchmarking (in-tree) — docs/design/module/benchmarking.md
+# ComfyUI
-/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
+/apps/ @fhfuih @alex-jw-brooks
-# LoRA (no dedicated design doc; ownership intentionally narrowed)
+# Large Scale Serving
-/vllm_omni/lora/ @SamitHuang @wtomin
+/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
-# Transformers utils (related path only in design docs; preserves prior ownership)
+# Quantization and configuration
-/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
+/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
-# =====================================================================
+# Hardware plugins and accelerators
-# Diffusion module family — docs/design/module/diffusion/
+/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
 # Umbrella rule first; sub-module rules after it win via last-match.
 # =====================================================================
 # Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
 # — docs/design/module/diffusion/index.md
 /vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 # Cache management — docs/design/module/cache_management.md
 /vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
 # Diffusion quantization — docs/design/module/quantization.md
 /vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
 # Diffusion profiling — docs/design/module/profiling.md
 /vllm_omni/diffusion/profiler/ @gcanlin
 # Diffusion runtime (scheduler, executor, worker)
 # — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
 /vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
 /vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
 # Diffusion model integration (pipelines, registration, checkpoint loading)
 # — docs/design/module/diffusion/diffusion_model_integration.md
 /vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
 /vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 /vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
 # Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
 # — docs/design/module/diffusion/parallelism.md
 /vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 /vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
 # Diffusion offloader (CPU offloading, DLO)
 # — docs/design/module/diffusion/offloader.md
 /vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
@@ -66,10 +66,7 @@ repos:
      entry: python tools/pre_commit/check_tts_adapter.py
      language: python
      pass_filenames: false
-      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
+      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
      # The budget-is-now-slack reminder is printed on a PASSING run, and
      # pre-commit hides output from passing hooks unless the hook is verbose.
      verbose: true
    # Keep `suggestion` last
    - id: suggestion
@@ -16,21 +16,31 @@ nav:
  - Examples:
    - examples/README.md
    - Offline Inference:
      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
      - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
    - Online Serving:
      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
-      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
  - General:
    - usage/*
@@ -7,7 +7,6 @@ owners:
  - "@yinpeiqi"
  - "@fake0fan"
  - "@Sy0307"
  - "@Gaohan123"
 primary_code_paths:
  - vllm_omni/core/**
  - vllm_omni/worker/**
@@ -4,7 +4,6 @@ kind: module
 status: draft
 owners:
  - "@alex-jw-brooks"
  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/benchmarks/**
  - vllm_omni/entrypoints/cli/benchmark/**
@@ -5,7 +5,6 @@ status: draft
 owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
 primary_code_paths:
  - vllm_omni/diffusion/cache/**
 related_code_paths:
@@ -3,10 +3,9 @@ title: Diffusion Model Integration
 kind: module
 status: draft
 owners:
-  - "@fhfuih"
+  - "@Isotr0py"
  - "@Bounty-hunter"
  - "@wtomin"
-  - "@RuixiangMa"
+  - "@ZJY0516"
 primary_code_paths:
  - vllm_omni/diffusion/models/**
  - vllm_omni/diffusion/model_loader/**
@@ -6,7 +6,6 @@ owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
  - "@fhfuih"
 primary_code_paths:
  - vllm_omni/diffusion/sched/**
  - vllm_omni/diffusion/executor/**
@@ -4,8 +4,7 @@ kind: module
 status: draft
 owners:
  - "@yuanheng-zhao"
-  - "@david6666666"
+  - "@Isotr0py"
  - "@lishunyang12"
 primary_code_paths:
  - vllm_omni/diffusion/offloader/**
 related_code_paths:
@@ -6,7 +6,6 @@ owners:
  - "@princepride"
  - "@yuanheng-zhao"
  - "@xuechendi"
  - "@Bounty-hunter"
 primary_code_paths:
  - vllm_omni/diffusion/distributed/**
  - vllm_omni/diffusion/attention/parallel/**
@@ -6,7 +6,6 @@ architecture_state: current-plus-in-flight
 owners:
  - "@alex-jw-brooks"
  - "@linyueqian"
  - "@NickCao"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -8,15 +8,8 @@ owners:
  - "@xuechendi"
 primary_code_paths:
  - vllm_omni/platforms/**
 related_code_paths:
  - vllm_omni/attention/**
 primary_code_path_owners:
  # Additional owners scoped to attention kernels (quantized/custom attention)
  - paths:
      - vllm_omni/attention/**
    owners:
      - "@david6666666"
      - "@lishunyang12"
 related_code_paths: []
 depends_on:
  - ar_runtime.md
  - diffusion/index.md
@@ -6,8 +6,6 @@ architecture_state: current
 owners:
  - "@Sy0307"
  - "@amy-why-3459"
  - "@Gaohan123"
  - "@alex-jw-brooks"
 document_stewards:
  - "@hsliuustc0106"
  - "@Gaohan123"
@@ -9,14 +9,6 @@ primary_code_paths:
  - vllm_omni/model_executor/**
  - vllm_omni/model_extras/**
  - vllm_omni/plugins/**
 primary_code_path_owners:
  # Additional owners scoped to a sub-path of model_executor/**
  - paths:
      - vllm_omni/model_executor/models/**
    owners:
      - "@Sy0307"
      - "@amy-why-3459"
      - "@linyueqian"
 related_code_paths:
  - vllm_omni/transformers_utils/**
  - vllm_omni/tokenizers/**
@@ -4,7 +4,6 @@ kind: module
 status: draft
 owners:
  - "@lishunyang12"
  - "@vraiti"
 primary_code_paths:
  - vllm_omni/metrics/**
 related_code_paths:
@@ -123,7 +123,6 @@ def extract_omni_serve_subparser_init():
                        "TrackingNamespace": None,  # type annotation only
                        "make_arg_parser": lambda parser: parser,  # no-op for doc
                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                        "_nonneg_finite_float": float,
                        "VLLM_SUBCMD_PARSER_EPILOG": "",
                        "logger": logger,
                        "DummySubparsers": DummySubparsers,
@@ -24,23 +24,6 @@ SERVING_MODE_TITLES = {
    "online_serving": "Online serving",
 }
 # Keep the Examples navigation focused on reusable modality/task entry points.
 # Model-specific scripts remain available under ``examples/`` and are surfaced
 # from these shared pages instead of getting one page per model.
 GENERAL_EXAMPLE_SLUGS = frozenset(
    {
        "image_to_image",
        "image_to_video",
        "speech_to_video",
        "text_to_audio",
        "text_to_image",
        "text_to_speech",
        "text_to_video",
        "x_to_text",
        "x_to_video_audio",
    }
 )
 def load_model_display_names() -> dict[str, str]:
    try:
@@ -63,13 +46,6 @@ def load_model_display_names() -> dict[str, str]:
 MODEL_DISPLAY_NAMES = load_model_display_names()
 def is_general_example(example: "Example") -> bool:
    """Return whether an example should get a generated serving page."""
    if example.category not in SERVING_MODE_TITLES:
        return True
    return example.path.stem in GENERAL_EXAMPLE_SLUGS
 def fix_case(text: str) -> str:
    subs = {
        "api": "API",
@@ -364,21 +340,22 @@ def update_nav_file(examples: list[Example]):
            examples_by_category[category] = []
        examples_by_category[category].append(example)
-    def category_items(category: str) -> list[dict[str, str]]:
+    # Build the new Examples section - start with preserved items
-        items = []
+    examples_section = preserved_items.copy()
-        for example in sorted(examples_by_category[category], key=lambda e: e.path.stem):
+
    # Add examples grouped by category, sorted by category name
    for category in sorted(examples_by_category.keys()):
        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
        category_items = []
        for example in category_examples:
            doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
            rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            items.append({example.nav_title: rel_path.as_posix()})
+            category_items.append({example.nav_title: rel_path.as_posix()})
        return items
-    # Build the new Examples section - start with preserved items.
+        if category_items:
-    examples_section = preserved_items.copy()
+            # Format category name (e.g., "offline_inference" -> "Offline Inference")
    for category in sorted(examples_by_category.keys()):
        items = category_items(category)
        if items:
            category_title = fix_case(category.replace("_", " ").title())
-            examples_section.append({category_title: items})
+            examples_section.append({category_title: category_items})
    # Update the nav structure
    nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
@@ -413,8 +390,6 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
        for path in category.glob("*/*.md"):
            examples.append(Example(path.parent, category.stem))
    examples = [example for example in examples if is_general_example(example)]
    # Generate the example documentation
    for example in sorted(examples, key=lambda e: e.path.stem):
        example_name = f"{example.path.stem}.md"
@@ -17,93 +17,84 @@ th {
 }
 </style>
-| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
+| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
-|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
+|--------------|--------|-------------------|------------|---------|-----|-----------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   | — |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
-| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
+| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
-| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
+| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
-| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
-| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
+| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   |
-| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
+| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
-| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
+| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
-| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
+| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
-| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
-| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
+| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | |
-| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
-| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
-| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | |
-| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
+| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
-| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
-| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
-| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
-| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
-| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
+| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
-| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
+| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
-| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
+| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
+| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
-| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
+| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
-| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
-| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
+| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
-| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
+| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
-| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
+| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
-| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
+| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
-| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
+| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
-| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
-| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
+| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
-| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
+| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
-| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
+| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
-| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
+| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
-| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
+| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
-| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
-| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
-| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
+| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
-| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
+| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
-|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
-| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
+| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
-| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
 link points to the repository recipe. `—` means that no recipe link has been
 added for that model row yet.
 For rows with an audited recipe link, the hardware checkmarks reflect hardware
 explicitly documented in that recipe. Other rows retain the implementation
 support metadata until their recipe is audited.
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
@@ -283,20 +283,23 @@ blocks — use Ulysses SP (`ring_degree=1`).
 ### Which geometries run sparse
-MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
+Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
-necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
+is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
-does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
+tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
-promoted rows remain visible to every query and key.
+that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
 attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
 speedup, so pick an aligned geometry when you want one.
-For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
+Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
-128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
+spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
-MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
+patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
-and kernel tiling aligned, at the cost of retaining more blocks.
+are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
 distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
 to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
-8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
+"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
-creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
+resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
-grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+which you can buy back with `start_step`.
 that implements protected video tails for rf_v2.
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
@@ -159,7 +159,7 @@ The following tables show which models support each feature:
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
-| **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
+| **MiniMax-H3**               |     ❌     |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
 **Frame Interpolation Support**
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends
 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -33,25 +33,6 @@ This folder provides a unified CLI script for image-to-video generation using vL
 Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
 ### LingBot-Video TI2V
 LingBot-Video accepts exactly one first-frame image and uses the same dense or
 MoE checkpoint as its T2I and T2V modes:
 ```bash
 python examples/offline_inference/image_to_video/image_to_video.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --model-class-name LingBotVideoPipeline \
  --image /path/to/input.png \
  --prompt "the fox looks toward the camera" \
  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
  --output lingbot_ti2v.mp4
 ```
 The input image keeps its original geometry until the LingBot pipeline applies
 its model-specific resize and center crop.
 ## Prerequisites
 Download the example image used in the snippets below:
@@ -43,23 +43,6 @@ This folder provides several entrypoints for experimenting with text-to-image di
 Default model: `Qwen/Qwen-Image`
 ### LingBot-Video
 LingBot-Video uses the same checkpoint for image and video generation. Run the
 shared T2I entry point with:
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --prompt "a red fox standing in fresh snow" \
  --height 192 --width 320 --num-inference-steps 2 \
  --guidance-scale 3.0 --extra-body '{"flow_shift":3.0}' \
  --output lingbot_t2i.png
 ```
 The LingBot prompt builder selects the `image` output modality and requests one
 frame. Passing an input image is not supported in LingBot T2I mode.
 ## Quick Start
 ### Python API
@@ -55,21 +55,6 @@ python text_to_video.py \
  --output vace_t2v_output.mp4
 ```
 ### LingBot-Video
 The dedicated LingBot example explicitly selects the video output modality:
 ```bash
 python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
  --output lingbot_t2v.mp4
 ```
 Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
 checkpoint. Requested frame counts are rounded upward to the causal
 VAE's `4n+1` grid.
 ### LTX-2
 ```bash
@@ -20,9 +20,9 @@ from vllm_omni.platforms import current_omni_platform
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
-    parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
-    parser.add_argument("--negative-prompt", default=None)
+    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--output", default="lingbot_video_output.mp4")
    parser.add_argument("--height", type=int, default=192)
    parser.add_argument("--width", type=int, default=320)
@@ -58,12 +58,9 @@ def _to_frame_array(video: Any) -> np.ndarray:
    if isinstance(video, torch.Tensor):
        tensor = video.detach().cpu()
        if tensor.dim() == 5:
-            tensor = tensor[0]
+            tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
-        if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
+        elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
-            if tensor.shape[0] in (3, 4):
+            tensor = tensor.permute(1, 2, 3, 0)
                tensor = tensor.permute(1, 2, 3, 0)
            elif tensor.shape[1] in (3, 4):
                tensor = tensor.permute(0, 2, 3, 1)
        if tensor.is_floating_point():
            if float(tensor.min()) < 0.0:
                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
@@ -90,12 +87,8 @@ def main() -> None:
        parallel_config=DiffusionParallelConfig(),
    )
-    prompt: dict[str, Any] = {
+    prompt: dict[str, str] = {"prompt": args.prompt}
-        "prompt": args.prompt,
+    if args.negative_prompt:
        "modalities": ["video"],
        "num_frames": args.num_frames,
    }
    if args.negative_prompt is not None:
        prompt["negative_prompt"] = args.negative_prompt
    sampling_params = OmniDiffusionSamplingParams(
@@ -1,74 +0,0 @@
 # LingBot-Video
 LingBot-Video uses one `LingBotVideoPipeline` for text-to-image (T2I),
 text-to-video (T2V), and text-image-to-video (TI2V) generation. Both the dense
 and MoE checkpoints use the same request format.
 ## Start the server
 ```bash
 MODEL=robbyant/lingbot-video-dense-1.3b bash run_server.sh
 ```
 The MoE checkpoint uses the same server and request scripts, but requires
 substantially more GPU memory:
 ```bash
 MODEL=robbyant/lingbot-video-moe-30b-a3b bash run_server.sh
 ```
 ## Text to image
 The image endpoint selects T2I mode and always generates one frame:
 ```bash
 bash run_curl_text_to_image.sh
 ```
 The script sends a `320x192`, two-step smoke request and writes
 `lingbot_t2i.png`.
 ## Text or text-image to video
 Run the video script without an image to select T2V mode:
 ```bash
 bash run_curl_text_image_to_video.sh
 ```
 Pass a first-frame image to the same script to select TI2V mode:
 ```bash
 INPUT_IMAGE=/path/to/input.png bash run_curl_text_image_to_video.sh
 ```
 The client scripts omit the optional `model` request field, so they target
 whichever dense or MoE checkpoint the server loaded. The video example uses the
 lightweight `320x192`, 9-frame, two-step configuration.
 Until the shared `/v1/videos` reference-image resizing is removed, TI2V target
 dimensions must be sent through `extra_params`, for example
 `{"size":"320x192"}`. Do not use the top-level `size`, `width`, or `height`
 fields for TI2V because the serving layer currently applies those dimensions
 to the reference image before the model receives it. T2V requests can continue
 to use the top-level dimension fields.
 LingBot video frame counts use the causal VAE `4n+1` grid. The pipeline rounds
 any requested frame count upward to the next valid value. An explicit
 `num_frames` takes precedence over `seconds`; otherwise, the server first
 resolves `seconds * fps` and the pipeline applies the same alignment.
 Official `resolution`/`ratio` presets can be sent through `extra_params`, for
 example `{"resolution":"720p","ratio":"16:9"}`. The `2k` and `4k` entries
 only define output dimensions; whether they run successfully depends on the
 checkpoint, GPU memory, and memory optimizations available in the deployment.
 For `/v1/images/generations`, the server resolves these aliases to their final
 pixel dimensions before applying `--max-generated-image-size`. Requests above
 the configured limit return HTTP 400 before engine dispatch. LingBot produces
 one output per prompt; image requests with `n>1` are also rejected with HTTP
 400.
 LingBot TI2V accepts exactly one image reference. Image editing, video
 references, audio references, batching, and Refiner execution are not supported
 by this pipeline mode.
@@ -1,67 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 if [ -n "${INPUT_IMAGE:-}" ] && [ ! -f "${INPUT_IMAGE}" ]; then
    echo "INPUT_IMAGE does not exist: ${INPUT_IMAGE}"
    exit 1
 fi
 BASE_URL="${BASE_URL:-http://localhost:8099}"
 POLL_INTERVAL="${POLL_INTERVAL:-2}"
 if [ -n "${INPUT_IMAGE:-}" ]; then
    default_output="lingbot_ti2v.mp4"
 else
    default_output="lingbot_t2v.mp4"
 fi
 OUTPUT_PATH="${OUTPUT_PATH:-${default_output}}"
 curl_args=(
    -sS -X POST "${BASE_URL}/v1/videos"
    -H "Accept: application/json"
    -F "prompt=the subject turns toward the camera with smooth natural motion"
    -F "num_frames=9"
    -F "fps=24"
    -F "num_inference_steps=2"
    -F "guidance_scale=3.0"
    -F "flow_shift=3.0"
    -F "seed=42"
 )
 if [ -n "${INPUT_IMAGE:-}" ]; then
    curl_args+=(
        -F "input_reference=@${INPUT_IMAGE}"
        -F 'extra_params={"size":"320x192"}'
    )
 else
    curl_args+=(-F "width=320" -F "height=192")
 fi
 create_response="$(curl "${curl_args[@]}")"
 video_id="$(echo "${create_response}" | jq -er '.id')"
 echo "Created video job ${video_id}"
 while true; do
    status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
    status="$(echo "${status_response}" | jq -er '.status')"
    case "${status}" in
        queued|in_progress)
            sleep "${POLL_INTERVAL}"
            ;;
        completed)
            break
            ;;
        failed)
            echo "${status_response}" | jq .
            exit 1
            ;;
        *)
            echo "Unexpected video job status: ${status}"
            exit 1
            ;;
    esac
 done
 curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
 echo "Saved video to ${OUTPUT_PATH}"
@@ -1,20 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 BASE_URL="${BASE_URL:-http://localhost:8099}"
 OUTPUT_PATH="${OUTPUT_PATH:-lingbot_t2i.png}"
 curl -sS -X POST "${BASE_URL}/v1/images/generations" \
    -H "Content-Type: application/json" \
    -d "{
        \"prompt\": \"a red fox standing in fresh snow\",
        \"size\": \"320x192\",
        \"num_inference_steps\": 2,
        \"guidance_scale\": 3.0,
        \"flow_shift\": 3.0,
        \"seed\": 42,
        \"response_format\": \"b64_json\"
    }" | jq -er '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
 echo "Saved image to ${OUTPUT_PATH}"
@@ -1,10 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 MODEL="${MODEL:-robbyant/lingbot-video-dense-1.3b}"
 PORT="${PORT:-8099}"
 vllm serve "${MODEL}" --omni \
    --model-class-name LingBotVideoPipeline \
    --port "${PORT}"
@@ -0,0 +1,247 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Verify a ModelOpt FP8 diffusers checkpoint exported by
 quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
 Three checks:
  A. transformer/config.json has a sane quantization_config block.
  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
  C. transformer disk size is materially smaller than a BF16 baseline.
 Example:
    python examples/quantization/check_modelopt_fp8_export.py \\
        --output ./hv15-480p-modelopt-fp8
    # Optional: compare disk size against a local or HF BF16 baseline.
    python examples/quantization/check_modelopt_fp8_export.py \\
        --output ./hv15-480p-modelopt-fp8 \\
        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
 """
 from __future__ import annotations
 import argparse
 import json
 import sys
 from collections import Counter
 from pathlib import Path
 SUPPORTED_ALGOS = ("FP8",)
 def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
    cfg_path = transformer_dir / "config.json"
    if not cfg_path.exists():
        print(f"[FAIL] {cfg_path} missing.")
        return 1, None
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
        return 1, None
    print(f"[A] quantization_config from {cfg_path}:")
    print(json.dumps(qc, indent=2))
    quant_algo = qc.get("quant_algo")
    issues = []
    if qc.get("quant_method") != "modelopt":
        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
    if quant_algo not in SUPPORTED_ALGOS:
        issues.append(
            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
        )
    if issues:
        print("[A] WARN — config looks incomplete:")
        for issue in issues:
            print(f"    - {issue}")
        return 2, quant_algo
    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
    return 0, quant_algo
 def _read_safetensors_header(path: Path) -> dict:
    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
    """
    import struct
    with open(path, "rb") as f:
        header_len = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(header_len))
    header.pop("__metadata__", None)
    return header
 def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
    """
    if not weight_scale_shapes:
        return "no weight_scale tensors found"
    def meaningful_dims(shape: list[int]) -> int:
        return sum(1 for d in shape if d > 1)
    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
    total = len(weight_scale_shapes)
    if per_tensor == total:
        return "per-tensor (all scalar scales)"
    if per_channel == total:
        return "per-channel (1 meaningful dim)"
    if per_block == total:
        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
 def _check_safetensors(transformer_dir: Path) -> int:
    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
    files = sorted(transformer_dir.glob("*.safetensors"))
    if not files:
        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
        return 1
    header_dtype_counts: Counter[str] = Counter()
    sample_quant_weight_keys: list[str] = []
    sample_scale_keys: list[str] = []
    weight_scale_shapes: list[list[int]] = []
    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
    for f in files:
        try:
            header = _read_safetensors_header(f)
        except Exception as exc:
            print(f"[B] WARN — could not parse header of {f}: {exc}")
            continue
        for k, info in header.items():
            dtype = info.get("dtype", "?")
            header_dtype_counts[dtype] += 1
            # FP8 stores weights as F8_E4M3 directly.
            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
                sample_quant_weight_keys.append(k)
            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
                sample_scale_keys.append(k)
            if k.endswith(".weight_scale"):
                weight_scale_shapes.append(info.get("shape", []))
                if len(sample_weight_scale_entries) < 5:
                    sample_weight_scale_entries.append((k, info.get("shape", [])))
    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
        marker = "  <-- FP8" if dtype.startswith("F8") else ""
        print(f"    {dtype:10s} {count:>6d}{marker}")
    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
    if quant_count == 0:
        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
        return 1
    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
    if sample_quant_weight_keys:
        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
    if sample_scale_keys:
        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
    for key, shape in sample_weight_scale_entries[:3]:
        print(f"      {key}: shape {shape}")
    return 0
 def _disk_size_gib(p: Path) -> float:
    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
 def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
    """Returns 0 always (informational only)."""
    quant_size = _disk_size_gib(transformer_dir)
    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
    if baseline is None:
        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
        return 0
    baseline_path = Path(baseline)
    if not baseline_path.exists():
        # Try HF download.
        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
            return 0
        print(f"    Downloading baseline transformer from HF: {baseline}")
        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
    bf16_size = _disk_size_gib(bf16_dir)
    if bf16_size == 0:
        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
        return 0
    # Expected reduction for FP8: ~50%.
    min_reduction = 30
    reduction = (1 - quant_size / bf16_size) * 100
    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
    if reduction < min_reduction:
        print(
            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
        )
    return 0
 def main() -> None:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
    p.add_argument(
        "--baseline",
        default=None,
        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
    )
    args = p.parse_args()
    out_root = Path(args.output).expanduser().resolve()
    transformer_dir = out_root / "transformer"
    if not transformer_dir.exists():
        print(f"[FAIL] {transformer_dir} does not exist.")
        sys.exit(1)
    print(f"Checking: {out_root}\n")
    fail = 0
    config_status, _quant_algo = _check_config(transformer_dir)
    fail |= config_status
    fail |= _check_safetensors(transformer_dir)
    _check_size_vs_baseline(transformer_dir, args.baseline)
    print()
    if fail == 0:
        print("=" * 60)
        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
    elif fail == 1:
        print("=" * 60)
        print("FAILURES detected — calibration may need to be re-run.")
        sys.exit(1)
    else:
        print("=" * 60)
        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
 if __name__ == "__main__":
    main()
@@ -0,0 +1,448 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
 Calibrates the DiT transformer using a small video prompt set and exports a
 diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
 The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
 (see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
 Layers kept full precision match the #2728 / #2795 pattern: modulation,
 AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
 and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
 empirically degrades under FP8 (see #2920 ablation).
 Example:
    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
        --output ./hv15-480p-modelopt-fp8 \\
        --overwrite
 """
 from __future__ import annotations
 import argparse
 import copy
 import json
 import re
 import shutil
 import sys
 from pathlib import Path
 from typing import Any
 import torch
 from diffusers import DiffusionPipeline
 DEFAULT_PROMPTS = [
    "A dog running across a field of golden wheat.",
    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
    "A close-up of a blooming rose covered in morning dew, soft natural light.",
    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
 ]
 def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
    p.add_argument("--height", type=int, default=480)
    p.add_argument("--width", type=int, default=832)
    p.add_argument(
        "--num-frames",
        type=int,
        default=33,
        help="Frames per calibration sample. 33 matches the typical short benchmark.",
    )
    p.add_argument("--guidance-scale", type=float, default=6.0)
    p.add_argument(
        "--calib-steps",
        type=int,
        default=10,
        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
    )
    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--prompt",
        action="append",
        default=[],
        help="Custom calibration prompt. Repeat to provide multiple.",
    )
    p.add_argument(
        "--quantize-mha",
        action="store_true",
        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
    )
    p.add_argument(
        "--weight-block-size",
        type=str,
        default=None,
        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
    )
    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
    return p
 def _parse_block_size(spec: str | None) -> list[int] | None:
    if spec is None:
        return None
    parts = [int(x) for x in spec.split(",") if x.strip()]
    if len(parts) != 2:
        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
    return parts
 def _require_modelopt() -> Any:
    try:
        import modelopt.torch.quantization as mtq
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "NVIDIA ModelOpt is not installed. Install with:\n"
            "  pip install 'nvidia-modelopt[all]'\n"
            f"Original error: {exc}"
        ) from exc
    return mtq
 def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
    model_path = args.model
    output_dir = Path(args.output).expanduser().resolve()
    if output_dir.exists():
        if not args.overwrite:
            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
        shutil.rmtree(output_dir)
    return model_path, output_dir
 def _select_dtype(name: str) -> torch.dtype:
    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
 def _build_prompts(args: argparse.Namespace) -> list[str]:
    prompts = args.prompt or DEFAULT_PROMPTS
    if args.calib_size <= 0:
        raise SystemExit("--calib-size must be positive.")
    if len(prompts) < args.calib_size:
        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
        prompts = (prompts * repeats)[: args.calib_size]
    return prompts[: args.calib_size]
 # Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
 # - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
 # - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
 # - proj_out: final output projection
 # - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
 def _filter_func_hv15(name: str) -> bool:
    pattern = re.compile(
        r"(proj_out.*|"
        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
        r"time_embed|cond_type_embed|"
        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
        r"token_refiner).*)"
    )
    return pattern.match(name) is not None
 def _mha_filter_func(name: str) -> bool:
    pattern = re.compile(
        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
    )
    return pattern.match(name) is not None
 def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
    if not hasattr(mtq, "disable_quantizer"):
        return
    mtq.disable_quantizer(backbone, _filter_func_hv15)
    if not quantize_mha:
        mtq.disable_quantizer(backbone, _mha_filter_func)
 def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
    if hasattr(pipe, "set_progress_bar_config"):
        pipe.set_progress_bar_config(disable=True)
    pipe.to("cuda")
    return pipe
 def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
    generator = torch.Generator(device="cuda")
    # Try to set guidance on the pipeline's guider object up front (modern
    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
    # back silently — calibration uses whatever default the pipeline ships with.
    guider = getattr(pipe, "guider", None)
    if guider is not None and hasattr(guider, "guidance_scale"):
        try:
            guider.guidance_scale = args.guidance_scale
        except Exception:
            pass
    base_kwargs = dict(
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.calib_steps,
        output_type="latent",
    )
    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
        with torch.inference_mode():
            for idx, prompt in enumerate(prompts):
                generator.manual_seed(args.seed + idx)
                # Try with guidance_scale first; fall back without on TypeError
                # for pipelines (like HV-1.5) that take CFG via guider config.
                try:
                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
                except TypeError as exc:
                    if "guidance_scale" not in str(exc):
                        raise
                    pipe(prompt=prompt, generator=generator, **base_kwargs)
    return forward_loop
 def _summarize_export(output_dir: Path) -> None:
    cfg_path = output_dir / "transformer" / "config.json"
    if not cfg_path.exists():
        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
        return
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
        return
    print("Export summary:")
    print(f"  quant_method: {qc.get('quant_method')}")
    print(f"  quant_algo:   {qc.get('quant_algo')}")
    producer = qc.get("producer")
    if isinstance(producer, dict):
        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
    print(f"  config path:  {cfg_path}")
 def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
    """Convert in-memory weights of quantized modules to actual FP8 storage.
    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
    in ModelOpt's recognized-model registry), so we must call the per-weight
    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
    helper.
    """
    from modelopt.torch.export.quant_utils import (
        QUANTIZATION_NONE,
        get_quantization_format,
        quantizer_attr_names,
        weight_attr_names,
    )
    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
    exported = 0
    for name, module in backbone.named_modules():
        try:
            quantization_format = get_quantization_format(module)
        except Exception as exc:
            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
            continue
        if quantization_format == QUANTIZATION_NONE:
            continue
        for weight_name in weight_attr_names(module):
            quantizer_attrs = quantizer_attr_names(weight_name)
            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
                continue
            _export_quantized_weight(module, dtype, weight_name)
            exported += 1
    return exported
 def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
    `weight_block_size` is given, advertise block-wise weight quantization in
    the saved metadata (so consumers know to expect multi-element scale tensors).
    """
    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
    if weight_block_size is not None:
        weights_cfg["strategy"] = "block"
        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
    return {
        "config_groups": {
            "group_0": {
                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
                "weights": weights_cfg,
                "targets": ["Linear"],
            }
        },
        "ignore": [
            "context_embedder*",
            "context_embedder_2*",
            "cond_type_embed*",
            "image_embedder*",
            "norm1.linear*",
            "norm1_context.linear*",
            "norm2*",
            "norm2_context*",
            "norm_out*",
            "proj_out*",
            "time_embed*",
            "token_refiner*",
            "x_embedder*",
        ],
        "producer": {"name": "modelopt"},
        "quant_algo": "FP8",
        "quant_method": "modelopt",
    }
 def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
    cfg_path = output_dir / "transformer" / "config.json"
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
    existing = cfg.get("quantization_config")
    if isinstance(existing, dict):
        producer = existing.get("producer")
        if isinstance(producer, dict):
            new_qc["producer"] = producer
    cfg["quantization_config"] = new_qc
    with cfg_path.open("w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)
 def _save_pipeline_with_fp8_transformer(
    pipe: DiffusionPipeline,
    model_path: str,
    output_dir: Path,
    max_shard_size: str = "5GB",
 ) -> None:
    """Save the pipeline with the (now FP8) transformer.
    Copies the source directory verbatim except for `transformer/`, then
    saves the transformer with quantizers hidden so the state dict contains
    only the FP8 weights + scale tensors.
    """
    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
    src = Path(model_path)
    if not src.exists():
        from huggingface_hub import snapshot_download
        src = Path(snapshot_download(model_path))
    if output_dir.exists():
        shutil.rmtree(output_dir)
    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
    transformer_out = output_dir / "transformer"
    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
    # nn.Module (transformer), not the diffusers Pipeline wrapper.
    with hide_quantizers_from_state_dict(pipe.transformer):
        pipe.transformer.save_pretrained(
            str(transformer_out),
            safe_serialization=True,
            max_shard_size=max_shard_size,
        )
 def main() -> None:
    args = _build_parser().parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
    mtq = _require_modelopt()
    model_path, output_dir = _ensure_paths(args)
    dtype = _select_dtype(args.dtype)
    prompts = _build_prompts(args)
    print("Quantization plan:")
    weight_block_size = _parse_block_size(args.weight_block_size)
    print(f"  input:           {args.model}")
    print(f"  output:          {output_dir}")
    print(f"  dtype:           {dtype}")
    print(f"  height/width:    {args.height}x{args.width}")
    print(f"  num_frames:      {args.num_frames}")
    print(f"  calib_size:      {len(prompts)}")
    print(f"  calib_steps:     {args.calib_steps}")
    print(f"  quantize_mha:    {args.quantize_mha}")
    print(
        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
    )
    pipe = _load_pipeline(model_path, dtype)
    backbone = pipe.transformer
    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
    if weight_block_size is not None:
        # Switch from per-tensor (default) to block-wise weight quantization.
        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
        quant_config["quant_cfg"]["*weight_quantizer"] = {
            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
        }
        print(
            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
        )
    forward_loop = _build_forward_loop(pipe, args, prompts)
    quantized = mtq.quantize(backbone, quant_config, forward_loop)
    if quantized is not None:
        pipe.transformer = quantized
        backbone = quantized
    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
    print("so we have to call the per-weight export helper ourselves)...")
    exported = _force_export_quantized_weights(backbone, dtype)
    print(f"  -> {exported} weights converted to FP8 in memory")
    if exported == 0:
        raise SystemExit(
            "No quantized weights were exported. Calibration may have skipped every layer "
            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
            "weight quantizers."
        )
    print("\nSaving pipeline with FP8 transformer...")
    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
    print(f"Saved to: {output_dir}")
    _summarize_export(output_dir)
    print("\nNext: validate the checkpoint with vllm-omni:")
    print(
        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
        f"    --model {output_dir} \\\n"
        "    --quantization fp8 \\\n"
        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
        "    --enforce-eager"
    )
    print(
        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
        "checkpoint's config.json has modelopt metadata.)"
    )
 if __name__ == "__main__":
    main()
@@ -0,0 +1,471 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
 Calibrates the DiT transformer using a small video prompt set and exports a
 diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
 The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
 (see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
 Layers kept full precision match the #2728 / #2795 pattern: condition embedder
 (time/text/image), patch embedding, modulation (scale_shift_table), final
 norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
 are quantized — static calibration handles the numerics that online FP8
 couldn't (see #2920 ablation).
 Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
 fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
 Example:
    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
        --output ./wan22-ti2v-modelopt-fp8 \\
        --overwrite
 """
 from __future__ import annotations
 import argparse
 import copy
 import json
 import re
 import shutil
 import sys
 from pathlib import Path
 from typing import Any
 import torch
 from diffusers import DiffusionPipeline
 DEFAULT_PROMPTS = [
    "A dog running across a field of golden wheat.",
    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
    "A close-up of a blooming rose covered in morning dew, soft natural light.",
    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
 ]
 def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
    p.add_argument(
        "--num-frames",
        type=int,
        default=49,
        help="Frames per calibration sample. 49 matches the typical short benchmark; "
        "use 17 to reduce memory pressure during calibration.",
    )
    p.add_argument("--guidance-scale", type=float, default=5.0)
    p.add_argument(
        "--calib-steps",
        type=int,
        default=10,
        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
    )
    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--prompt",
        action="append",
        default=[],
        help="Custom calibration prompt. Repeat to provide multiple.",
    )
    p.add_argument(
        "--quantize-mha",
        action="store_true",
        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
        "sequences amplified FP8 drift in the online ablation (see #2920).",
    )
    p.add_argument(
        "--weight-block-size",
        type=str,
        default=None,
        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
        "for the HV-1.5 investigation status before relying on this.",
    )
    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
    return p
 def _parse_block_size(spec: str | None) -> list[int] | None:
    if spec is None:
        return None
    parts = [int(x) for x in spec.split(",") if x.strip()]
    if len(parts) != 2:
        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
    return parts
 def _require_modelopt() -> Any:
    try:
        import modelopt.torch.quantization as mtq
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "NVIDIA ModelOpt is not installed. Install with:\n"
            "  pip install 'nvidia-modelopt[all]'\n"
            f"Original error: {exc}"
        ) from exc
    return mtq
 def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
    model_path = args.model
    output_dir = Path(args.output).expanduser().resolve()
    if output_dir.exists():
        if not args.overwrite:
            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
        shutil.rmtree(output_dir)
    return model_path, output_dir
 def _select_dtype(name: str) -> torch.dtype:
    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
 def _build_prompts(args: argparse.Namespace) -> list[str]:
    prompts = args.prompt or DEFAULT_PROMPTS
    if args.calib_size <= 0:
        raise SystemExit("--calib-size must be positive.")
    if len(prompts) < args.calib_size:
        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
        prompts = (prompts * repeats)[: args.calib_size]
    return prompts[: args.calib_size]
 # Layers to KEEP at full precision. Wan2.2's module naming:
 # - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
 # - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
 # - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
 # - norm_out: AdaLayerNorm final
 # - proj_out: final nn.Linear
 # - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
 def _filter_func_wan22(name: str) -> bool:
    pattern = re.compile(
        r"(proj_out.*|"
        r".*(condition_embedder|patch_embedding|"
        r"norm_out|scale_shift_table|"
        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
    )
    return pattern.match(name) is not None
 def _mha_filter_func(name: str) -> bool:
    pattern = re.compile(
        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
    )
    return pattern.match(name) is not None
 def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
    if not hasattr(mtq, "disable_quantizer"):
        return
    mtq.disable_quantizer(backbone, _filter_func_wan22)
    if not quantize_mha:
        mtq.disable_quantizer(backbone, _mha_filter_func)
 def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
    if hasattr(pipe, "set_progress_bar_config"):
        pipe.set_progress_bar_config(disable=True)
    pipe.to("cuda")
    return pipe
 def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
    generator = torch.Generator(device="cuda")
    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
    guider = getattr(pipe, "guider", None)
    if guider is not None and hasattr(guider, "guidance_scale"):
        try:
            guider.guidance_scale = args.guidance_scale
        except Exception:
            pass
    base_kwargs = dict(
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.calib_steps,
        output_type="latent",
    )
    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
        with torch.inference_mode():
            for idx, prompt in enumerate(prompts):
                generator.manual_seed(args.seed + idx)
                # Try with guidance_scale first; fall back without on TypeError
                # for pipelines that take CFG via guider config only.
                try:
                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
                except TypeError as exc:
                    if "guidance_scale" not in str(exc):
                        raise
                    pipe(prompt=prompt, generator=generator, **base_kwargs)
    return forward_loop
 def _summarize_export(output_dir: Path) -> None:
    cfg_path = output_dir / "transformer" / "config.json"
    if not cfg_path.exists():
        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
        return
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
        return
    print("Export summary:")
    print(f"  quant_method: {qc.get('quant_method')}")
    print(f"  quant_algo:   {qc.get('quant_algo')}")
    producer = qc.get("producer")
    if isinstance(producer, dict):
        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
    print(f"  config path:  {cfg_path}")
 def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
    """Convert in-memory weights of quantized modules to actual FP8 storage.
    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
    in ModelOpt's recognized-model registry), so we must call the per-weight
    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
    calibration helpers.
    """
    from modelopt.torch.export.quant_utils import (
        QUANTIZATION_NONE,
        get_quantization_format,
        quantizer_attr_names,
        weight_attr_names,
    )
    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
    exported = 0
    for name, module in backbone.named_modules():
        try:
            quantization_format = get_quantization_format(module)
        except Exception as exc:
            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
            continue
        if quantization_format == QUANTIZATION_NONE:
            continue
        for weight_name in weight_attr_names(module):
            quantizer_attrs = quantizer_attr_names(weight_name)
            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
                continue
            _export_quantized_weight(module, dtype, weight_name)
            exported += 1
    return exported
 def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
    if weight_block_size is not None:
        weights_cfg["strategy"] = "block"
        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
    return {
        "config_groups": {
            "group_0": {
                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
                "weights": weights_cfg,
                "targets": ["Linear"],
            }
        },
        "ignore": [
            "condition_embedder*",
            "norm_out*",
            "output_scale_shift_prepare*",
            "patch_embedding*",
            "proj_out*",
            "scale_shift_table*",
            "timestep_proj_prepare*",
        ],
        "producer": {"name": "modelopt"},
        "quant_algo": "FP8",
        "quant_method": "modelopt",
    }
 def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
    for sub in ("transformer", "transformer_2"):
        cfg_path = output_dir / sub / "config.json"
        if not cfg_path.exists():
            continue
        with cfg_path.open(encoding="utf-8") as f:
            cfg = json.load(f)
        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
        existing = cfg.get("quantization_config")
        if isinstance(existing, dict):
            producer = existing.get("producer")
            if isinstance(producer, dict):
                new_qc["producer"] = producer
        cfg["quantization_config"] = new_qc
        with cfg_path.open("w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2)
 def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
    """Return [(attr_name, module)] for every DiT on the pipeline.
    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
    (high-noise). TI2V-5B has only ``transformer``.
    """
    out = []
    for attr in ("transformer", "transformer_2"):
        mod = getattr(pipe, attr, None)
        if mod is not None:
            out.append((attr, mod))
    return out
 def _save_pipeline_with_fp8_transformer(
    pipe: DiffusionPipeline,
    model_path: str,
    output_dir: Path,
    max_shard_size: str = "5GB",
 ) -> None:
    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
    as the single-DiT TI2V-5B.
    """
    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
    src = Path(model_path)
    if not src.exists():
        from huggingface_hub import snapshot_download
        src = Path(snapshot_download(model_path))
    if output_dir.exists():
        shutil.rmtree(output_dir)
    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
    for attr, backbone in _list_transformers(pipe):
        transformer_out = output_dir / attr
        # Pass the nn.Module (transformer), not the Pipeline wrapper.
        with hide_quantizers_from_state_dict(backbone):
            backbone.save_pretrained(
                str(transformer_out),
                safe_serialization=True,
                max_shard_size=max_shard_size,
            )
 def main() -> None:
    args = _build_parser().parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
    mtq = _require_modelopt()
    model_path, output_dir = _ensure_paths(args)
    dtype = _select_dtype(args.dtype)
    prompts = _build_prompts(args)
    weight_block_size = _parse_block_size(args.weight_block_size)
    print("Quantization plan:")
    print(f"  input:           {args.model}")
    print(f"  output:          {output_dir}")
    print(f"  dtype:           {dtype}")
    print(f"  height/width:    {args.height}x{args.width}")
    print(f"  num_frames:      {args.num_frames}")
    print(f"  calib_size:      {len(prompts)}")
    print(f"  calib_steps:     {args.calib_steps}")
    print(f"  quantize_mha:    {args.quantize_mha}")
    print(
        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
    )
    pipe = _load_pipeline(model_path, dtype)
    transformers = _list_transformers(pipe)
    if not transformers:
        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
    if weight_block_size is not None:
        quant_config["quant_cfg"]["*weight_quantizer"] = {
            "num_bits": (4, 3),
            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
        }
        print(
            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
        )
    forward_loop = _build_forward_loop(pipe, args, prompts)
    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
    # on every attached transformer (A14B: transformer + transformer_2) get
    # calibrated during the same passes.
    for attr, backbone in transformers:
        print(f"\nQuantizing {attr}...")
        quantized = mtq.quantize(backbone, quant_config, forward_loop)
        if quantized is not None:
            setattr(pipe, attr, quantized)
        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
    # inside the loop would convert transformer's weights to real FP8, and the
    # next transformer's forward_loop (full pipeline) would then fake-quantize an
    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
    print("so we have to call the per-weight export helper ourselves)...")
    total_exported = 0
    for attr, backbone in _list_transformers(pipe):
        exported = _force_export_quantized_weights(backbone, dtype)
        print(f"  -> {exported} weights converted to FP8 in {attr}")
        total_exported += exported
    if total_exported == 0:
        raise SystemExit(
            "No quantized weights were exported. Calibration may have skipped every layer "
            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
            "weight quantizers."
        )
    print("\nSaving pipeline with FP8 transformer(s)...")
    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
    print(f"Saved to: {output_dir}")
    _summarize_export(output_dir)
    print("\nNext: validate the checkpoint with vllm-omni:")
    print(
        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
        f"    --model {output_dir} \\\n"
        "    --quantization fp8 \\\n"
        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
        "    --output outputs/wan22_modelopt_fp8.mp4"
    )
    print(
        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
        "checkpoint's config.json has modelopt metadata.)"
    )
 if __name__ == "__main__":
    main()
@@ -79,8 +79,7 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 ### Multi-NPU: 768P combined-service configuration
 Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
-degree 8, native tiled VAE patch parallelism degree 8, and distributed
+degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
 layerwise offload:
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -99,7 +98,7 @@ vllm serve "${MODEL}" \
  --usp 8 \
  --ring 1 \
  --text-encoder-tp-size 8 \
-  --enable-distributed-layerwise-offload \
+  --enable-layerwise-offload \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --vae-patch-parallel-size 8 \
@@ -113,11 +112,6 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
 partition restart is required. Layerwise offload applies to both DiTs.
 On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
 768P without offloading or sharding: use distributed layerwise offload (as
 above) or HSDP — see
 [§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
 ### Optional optimizations
 Two independent optimizations may be enabled on top of the configuration
@@ -140,63 +134,6 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
 over the whole sequence, so ring parallelism would split away the keys it
 needs. Scale with `--usp` instead.
 ## Memory and attention optimizations (A3)
 ### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
 Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
 combined MiniMax-H3 service at 768P. Enable **one** of the following at
 server startup:
 **Distributed layerwise offload** — DiT layers are offloaded with parameters
 gathered across the parallel group instead of replicated per rank:
 ```bash
  --enable-distributed-layerwise-offload
 ```
 Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
 input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
 usage is high with this option.
 **HSDP** — hybrid sharded data parallelism for the DiT parameters. The
 multi-stream memory-reuse knob is mandatory with this flag:
 ```bash
 export MULTI_STREAM_MEMORY_REUSE=2
 ```
 ```bash
  --use-hsdp
 ```
 Host memory usage is small, but large shapes may still OOM; HBM usage
 optimizations for this path are ongoing.
 ### FLASH_ATTN backend with MindIE-SD
 Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
 [§ Environment](#environment)). On NPU this backend carries most of the
 memory-reduction and performance work: a mask-free packed varlen path that
 never materializes the quadratic `full_qk` padding mask, and K/V prefix
 slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
 transformer.
 ### Optional: LaserAttention fused kernel
 For an additional attention speedup, select the Ascend LaserAttention fused
 kernel before starting the server:
 ```bash
 export MINDIE_SD_FA_TYPE="ascend_laser_attention"
 ```
 This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
 exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
 kernel's fp16 workspace cannot overflow on outlier activations. Measured
 speedup numbers will be added here.
 ## HTTP API examples
 The request format is identical to the GPU recipe; see
@@ -217,8 +154,8 @@ configuration above:
 | Workload | Configuration |
 |----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
 These measurements describe the validated shapes rather than a general
 throughput guarantee.
@@ -1,346 +0,0 @@
 # MiniMax-H3 on RTX PRO 6000 Blackwell GPUs
 This recipe runs MiniMax-H3 in BF16 on 96 GiB RTX PRO 6000 Blackwell GPUs.
 The additional HBM over the RTX PRO 5000 profile removes the four-GPU
 minimum: two GPUs are enough to keep the model resident with TP2. Four and
 eight GPUs raise the Ulysses degree to shard the attention sequence further.
 CPU offload and distributed layerwise offload are not required in any of
 these configurations.
 Validated on:
 - Host: YLX Y762
 - GPUs: 8 × RTX PRO 6000 Blackwell (96 GiB)
 - Device order: default (`CUDA_VISIBLE_DEVICES` not set)
 - Driver / CUDA: 580.105.08 / 13.0
 - Container image: `vllm/vllm-omni:minimax-h3`
 - Workload: T2VA, 1344×768, `duration=5.0`, 50 steps, `flow_shift=12`, `seed=1101`
 ## Capacity requirements
 | Resource | Two GPUs | Four GPUs | Eight GPUs |
 | --- | ---: | ---: | ---: |
 | GPU HBM | 96 GiB per GPU | 96 GiB per GPU | 96 GiB per GPU |
 | Checkpoint storage | 135 GiB per partition | 135 GiB per partition | 135 GiB per partition |
 | Available system RAM | 200 GiB minimum | 200 GiB minimum | 200 GiB minimum |
 | Recommended system RAM | 384 GiB | 384 GiB | 384 GiB |
 | Measured peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
 `FL2VA` and `Ref2VA` are separate checkpoint partitions. Start one server at
 a time on a host sized for the minimum system-memory requirement, or pass
 `--task-type fl2va` or `--task-type ref2va` to download and load a single
 partition. The two-server layout described in the eight-GPU section doubles
 both the storage and the host-memory requirement.
 ## Parallelism summary
 | Configuration | TP | Ulysses | Text-encoder TP | VAE patch parallel |
 | --- | ---: | ---: | ---: | ---: |
 | Two GPUs | 2 | 1 | 2 | 2 |
 | Four GPUs | 2 | 2 | 4 | 4 |
 | Eight GPUs | 2 | 4 | 8 | 8 |
 | Eight GPUs, headroom variant | 4 | 2 | 8 | 8 |
 Per-GPU weight residency is set by the tensor-parallel degree alone. Raising
 the Ulysses degree lowers activation memory and per-step latency but leaves
 weight residency unchanged. Ring parallelism stays at 1, and H3 is
 CFG-distilled, so `--cfg-parallel-size` must also remain 1.
 ## PCIe topology and GPU order
 RTX PRO 6000 Blackwell does not provide NVLink. Before starting the server,
 inspect the host topology:
 ```bash
 nvidia-smi topo -m
 nvidia-smi nvlink -s
 ```
 Tensor-parallel groups are consecutive ranks. Ulysses groups are strided by
 the tensor-parallel degree. Confirm this grouping against the server log
 before tuning the device order on a new host.
 For the two-GPU configuration, pick a single `PXB` pair on one NUMA node.
 For the four-GPU configuration, pick two `PXB` pairs on the same node and
 order them so that each Ulysses group lands inside one pair. If physical
 GPUs `(0,1)` and `(2,3)` are the two local pairs, the order is
 `CUDA_VISIBLE_DEVICES=0,2,1,3`. Do not copy these IDs blindly: reproduce the
 same relationship on the target host.
 Device-order pinning and NUMA binding are tuning knobs, not prerequisites.
 The measurements in this recipe were taken without either, so any gain from
 them is additional to the numbers reported below and is host-specific.
 ## Two-GPU serving configuration
 Two 96 GiB GPUs hold the BF16 model with TP2 alone. There is no Ulysses
 group, so the DiT attention sequence is not sharded and activation memory is
 higher than in the four-GPU profile at the same output shape. Re-measure
 peak memory before increasing the output shape, the reference-input length,
 or concurrency.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 2 \
  --tensor-parallel-size 2 \
  --usp 1 \
  --ring 1 \
  --text-encoder-tp-size 2 \
  --vae-patch-parallel-size 2 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ## Four-GPU serving configuration
 Four GPUs keep TP2 for the weight shard and add Ulysses2, which shards the
 attention sequence and lowers per-GPU activation memory. Text-encoder TP4
 and VAE patch parallelism 4 spread the Qwen3-VL encoder and tiled decode
 across all four GPUs. TP1 is not an option on this card: an unsharded BF16
 partition does not fit in 96 GiB.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 4 \
  --tensor-parallel-size 2 \
  --usp 2 \
  --ring 1 \
  --text-encoder-tp-size 4 \
  --vae-patch-parallel-size 4 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ## Eight-GPU serving configuration
 Eight GPUs extend the Ulysses degree to 4 while keeping TP2. This does not
 lower per-GPU weight residency relative to the four-GPU profile; it lowers
 activation memory and per-step latency by sharding the attention sequence
 four ways.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 8 \
  --tensor-parallel-size 2 \
  --usp 4 \
  --ring 1 \
  --text-encoder-tp-size 8 \
  --vae-patch-parallel-size 8 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ### Rank ordering across two sockets
 With TP2 and Ulysses4 the two Ulysses groups are ranks `(0,2,4,6)` and
 `(1,3,5,7)`, and the four tensor-parallel pairs are `(0,1)`, `(2,3)`,
 `(4,5)`, and `(6,7)`. On a dual-socket host these two collectives cannot
 both be socket-local. If you choose to pin the device order, the trade-off
 is:
 | Device order | Socket-local collective | Crosses sockets |
 | --- | --- | --- |
 | `0,4,1,5,2,6,3,7` | Ulysses all-to-all | TP all-reduce |
 | `0,1,2,3,4,5,6,7` | TP all-reduce | Ulysses all-to-all |
 The first order assumes physical GPUs `0-3` are on NUMA node 0 and `4-7` on
 node 1. Reproduce that relationship on the target host rather than copying
 the IDs. The reported measurements use the default order and no NUMA
 binding; the two orders were not compared on this host.
 ### Headroom variant: TP4 with Ulysses2
 Raising the tensor-parallel degree to 4 halves per-GPU weight residency and
 leaves room for long Ref2VA references, larger output shapes, or
 concurrency greater than one. It also adds a four-way all-reduce per DiT
 block over a link without NVLink. Use it when the TP2 profile runs out of
 headroom, not as the default. Replace the parallel flags above with:
 ```bash
  --tensor-parallel-size 4 \
  --usp 2 \
 ```
 ### Two independent four-GPU servers
 An eight-GPU host can instead run one FL2VA server and one Ref2VA server
 side by side, each pinned to its own NUMA node with the four-GPU
 configuration. This removes every cross-socket collective and lifts the
 one-server-at-a-time restriction, at the cost of doubling the requirements
 to 400 GiB available system RAM and 270 GiB of model storage.
 Use `CUDA_VISIBLE_DEVICES=0,2,1,3`, `--cpunodebind=0 --membind=0`, and
 `--port 8091` for the FL2VA server. Use `CUDA_VISIBLE_DEVICES=4,6,5,7`,
 `--cpunodebind=1 --membind=1`, and `--port 8092` for the Ref2VA server.
 This layout is described for completeness and was not part of the measured
 runs.
 ## Shared serving notes
 Do not add `--enforce-eager` for a performance run. Warm the server before
 measuring so regional compilation is outside the measured request. On this
 host the first request after startup ran 19% slower than the steady state
 (107.58 s versus 90.48 s on eight GPUs); two warmup requests were enough to
 converge.
 For Ref2VA on a single-server layout, stop the FL2VA server and restart the
 same command with `MODEL="${MODEL_ROOT}/Ref2VA"`.
 ## Attention backend
 RTX PRO 6000 Blackwell is an SM120 part, so the datacenter-Blackwell
 `TRTLLM_ATTN` default does not apply here. Every configuration above selects
 cuDNN BF16 attention explicitly. This keeps the recipe independent of
 platform-default backend changes and does not depend on the experimental
 SM120 kernel.
 ## Optional: online FP8
 `--quantization fp8` quantizes eligible DiT linears at load time and is
 compatible with tensor parallelism and VAE tiling. Use it when the resident
 BF16 peak leaves too little headroom for long Ref2VA references, a larger
 output shape, or concurrency greater than one. It cannot be combined with
 layerwise offload.
 ## Target-hardware validation
 T2VA, 1344×768, `duration=5.0`, `fps=24`, 50 steps, `flow_shift=12`,
 `seed=1101`, BF16, `CUDNN_ATTN`, tiled VAE, one request at a time. Servers
 were started with the commands above — default device order, no NUMA
 binding — plus `--enable-diffusion-pipeline-profiler`. Two warmup requests
 preceded each measured request.
 | Measurement | Two GPUs | Four GPUs | Eight GPUs |
 | --- | ---: | ---: | ---: |
 | Topology | TP2 × USP1 | TP2 × USP2 | TP2 × USP4 |
 | Text encode | 0.040 s | 0.044 s | 0.035 s |
 | Denoise, 50 steps | 278.55 s | 168.73 s | 87.90 s |
 | Per step | 5.571 s | 3.375 s | 1.758 s |
 | VAE decode | 5.396 s | 2.791 s | 1.798 s |
 | Client E2E | 284.76 s | 172.32 s | 90.48 s |
 | Peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
 | Headroom below 96 GiB | 18.5 GiB | 29.6 GiB | 34.9 GiB |
 Stage times are read from the `X-Stage-Durations` response header of
 `/v1/videos/sync`. Peak memory is the maximum of `nvidia-smi
 --query-gpu=memory.used` sampled at 1 Hz across every device for the
 duration of the measured request. Per step is denoise wall time divided by
 the 50 requested steps. Stage times sum to roughly 0.75 s less than
 end-to-end in all three configurations; queueing, result transfer, and MP4
 muxing sit outside the profiled stages.
 ### Scaling behaviour
 Denoise speedup is not uniform across the sweep: 1.65× from two to four
 GPUs (83% efficiency) but 1.92× from four to eight (96%). Two GPUs run
 Ulysses1, so there is no sequence parallelism at all; going to four GPUs
 pays the one-time cost of introducing the all-to-all, while going from four
 to eight only widens an all-to-all that already exists. Expect the second
 doubling to pay off better than the first on this topology.
 VAE decode behaves in the opposite direction — 1.93× from two to four GPUs
 but only 1.55× from four to eight — as patch parallelism produces smaller
 tiles and tile-boundary overhead grows.
 Text encoding is 0.04 s in every configuration and is not worth optimising;
 denoise is 96-98% of end-to-end throughout.
 ### Per-GPU memory model
 Peak memory falls by 11.05 GiB from two to four GPUs and by a further
 5.37 GiB from four to eight — halving each time the Ulysses degree doubles,
 because all three configurations shard the DiT with TP2 and hold identical
 weights per GPU. Fitting the first two points gives:
 ```
 peak HBM per GPU ≈ 55.4 GiB + 22.1 GiB / ulysses_degree
 ```
 That fit predicts 60.9 GiB at Ulysses4 against 61.07 GiB measured, a 0.25%
 error. The practical consequence is a floor near 55 GiB per GPU: adding
 GPUs past eight will not reduce the per-GPU requirement further. To serve
 on smaller cards, raise the tensor-parallel degree instead.
 Re-measure memory for longer reference inputs, concurrency greater than
 one, or a different output shape.
 ### Not yet measured
 Ref2VA latency and memory have not been measured on this host. Ref2VA
 drives the Qwen3-VL presentation from tens of tokens for text-only prompts
 to several thousand with image or video references, so T2VA numbers do not
 transfer to it. The eight-GPU device-order comparison in the section above
 was also not run.
 ## T2VA request example
 `t2va` requires an explicit `aspect_ratio` even when `width` and `height`
 are supplied.
 ```bash
 export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
 curl -sS --max-time 1800 -X POST "${API_URL}" \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=1344' \
  -F 'height=768' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":5.0,"audio_flow_shift":3.0}' \
  -o t2va.mp4
 ```
 To collect the stage breakdown and peak memory for a request, add
 `--enable-diffusion-pipeline-profiler` to the server command and `-D
 headers.txt` to the curl invocation, then read `X-Stage-Durations` and
 `X-Peak-Memory-MB` from the saved headers.
 ## Ref2VA request example
 `ref2va` defaults to a 16:9 aspect ratio, so `aspect_ratio` is optional
 here.
 ```bash
 curl -X POST "http://127.0.0.1:${PORT}/v1/videos/sync" \
  --fail-with-body -w '\nHTTP %{http_code}\n' \
  --max-time 1200 \
  -F "input_reference=@/root/hand.jpg;type=image/jpeg" \
  -F "audio_reference=</root/audio_ref.json" \
  -F "prompt=2D动画融合在一起的影像。夕阳余晖残留在窗边，生活感十足的小厨房里有旧木桌、洗到一半的马克杯、起雾的玻璃瓶、悬挂的抹布。画面带有智能手机单手拍摄的手抖、近距离对焦的犹豫、逆光曝光波动。要像在家中慌忙拍下某个不可思议事件的自然质感，不要广告影像的精心整理。声音只用厨房环境声与手绘生物柔和的电子音、小小的叫声。" \
  -F 'width=1344' -F 'height=768' -F 'fps=24' \
  -F 'num_inference_steps=60' -F 'flow_shift=12' -F 'seed=1101' \
  -F 'extra_params={"task":"ref2va","duration":8,"audio_flow_shift":3.0}' \
  -o /root/out_ref2va.mp4
 ```
@@ -1,262 +0,0 @@
 # MiniMax-H3 on DGX Spark (GB10)
 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) |
 | --- | ---: |
 | Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
 | GPUs | 1 |
 | Checkpoint storage | 135 GiB per partition |
 | Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
 | Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
 On GB10 the CPU and GPU share one physical memory pool, which changes the capacity
 math relative to the RTX and datacenter recipes:
 - **Do not use `--enable-distributed-layerwise-offload`.** DLO stages rank-local
  weights into pinned host memory, which on GB10 is the same pool the weights are
  already in. Peak usage roughly doubles and the process is killed by the OOM
  killer (`Exit code: -9`) shortly after `Enabling offloader backend`.
 - **Do not use `--enable-cpu-offload`.** Moving weights from "VRAM" to "host RAM"
  frees nothing here.
 - **`--quantization fp8` is mandatory.** A BF16 partition is 135 GiB and does not
  fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
  Qwen3-VL text encoder and both VAEs stay BF16.
 With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
 23 GiB of headroom for the OS, page cache, and any other process on the box.
 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
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
 FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=7200 \
 vllm serve /path/to/MiniMax-H3/FL2VA \
  --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
 ```
 Wait for `Application startup complete`, then run a short smoke test before
 committing to a full 50-step request.
 `t2va` requires an explicit `aspect_ratio` even when `width` and `height` are
 supplied; omitting it fails the request with
 `OmniClientError: t2va requires an explicit aspect_ratio`.
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=10' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":4.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-smoke.mp4
 ```
 The full-quality request uses 50 steps:
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":8.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-t2va.mp4
 ffprobe -v error -show_entries \
  stream=index,codec_name,width,height,r_frame_rate,sample_rate,channels \
  -of json h3-gb10-t2va.mp4
 ```
 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:
 | Workload | Steps | Text encode | Denoise | Per step | VAE decode | MP4 mux | Client E2E | Peak memory |
 | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
 | T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
 | T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
 | Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
 | Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
 How these were measured:
 - Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
  `X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
  `MiniMaxH3Pipeline.encode_prompt` (text encode),
  `MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
  (VAE decode).
 - MP4 mux is not a header field. It is the
  `Video response encoding (MP4 bytes)` line in the server log.
 - Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
  high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
  the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
  higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
  memory has no discrete VRAM to report and the field returns `[N/A]`. See the
  capacity section for a `/proc/meminfo` cross-check.
 - For T2VA, text encode and VAE decode were measured on a separate 10-step run at
  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.
 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, and 7200 s is
 shorter than a 50-step Ref2VA run.
 ## Known limitations
 - One GPU means no Ulysses, ring, TP, or VAE patch parallelism. Pass
  `--usp 1 --ring 1 --vae-patch-parallel-size 1` explicitly if a profile or
  script would otherwise supply higher degrees.
 - `--enforce-eager` is used because regional compile adds startup time without a
  parallelism win on a single rank.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
  `height` are both supplied they take precedence and `aspect_ratio` only has to
  name one of the supported ratios.
 - Sustained load degrades throughput. See the variance notes above.
 - Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
  with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
  above; reference videos and additional images push the token count higher and
  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.
 ```
@@ -750,73 +750,8 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
 resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
 reduction.
 ## TeaCache acceleration
 TeaCache reuses DiT block residuals across denoising steps when consecutive
 timestep embeddings are similar. MiniMax-H3 TeaCache is currently calibrated
 only for the FL2VA partition. In combined serving, FL2VA requests use TeaCache
 while Ref2VA requests run uncached; Ref2VA-only serving rejects TeaCache.
 TeaCache and Cache-DiT are mutually exclusive; pick one cache backend per server.
 The model-specific default and examples use `rel_l1_thresh=0.17`, which
 provided the best conservative speed/quality balance in the validated 107-frame
 T2VA workload. Lower values
 may produce few or no cache hits, while higher values can improve performance
 at the cost of output quality. Validate the threshold on representative
 prompts and generation settings before changing it.
 ### Offline (Python API)
 Export the directory containing the `FL2VA` and `Ref2VA` subdirectories before
 running the Python example, for example `export MODEL_ROOT=/models/MiniMax-H3`.
 ```python
 import os
 from vllm_omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 omni = Omni(
    model=os.path.join(os.environ["MODEL_ROOT"], "FL2VA"),
    cache_backend="tea_cache",
    cache_config={"rel_l1_thresh": 0.17},
    trust_remote_code=True,
    enable_cpu_offload=True,
 )
 outputs = omni.generate(
    "A quiet cinematic night scene with matching ambient sound.",
    OmniDiffusionSamplingParams(
        height=256,
        width=448,
        num_frames=29,
        fps=24,
        num_inference_steps=50,
        seed=42,
        extra_args={
            "task": "t2va",
            "duration": 4.0,
            "aspect_ratio": "16:9",
            "flow_shift": 12.0,
            "audio_flow_shift": 3.0,
        },
    ),
 )
 ```
 ### Online serving
 ```bash
 vllm serve "${MODEL_ROOT}/FL2VA" \
  --omni \
  --trust-remote-code \
  --cache-backend tea_cache \
  --cache-config '{"rel_l1_thresh":0.17}'
 ```
 ## Known limitations
 - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
 - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
  both task-specific DiTs, and loads shared components once from `FL2VA`.
 - H3 currently executes one generation request per diffusion batch.
@@ -835,10 +770,6 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
  pure Ulysses.
 - Online FP8 is not compatible with layerwise offload.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone
  audio reference.
 - Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
  cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
  or model-level CPU offload; text-encoder TP alone is not sufficient.
@@ -19,11 +19,6 @@ add files for this repo. Use one Markdown file per model family by default.
 > Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
 > e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
 The `Recipe` column in the Supported Models documentation links to the
 published recipe when available and otherwise to the corresponding recipe in
 this directory. The recipe files remain the source of truth for commands and
 validation details.
 Example layout:
 ```text
@@ -39,7 +39,6 @@ from vllm_omni.config.stage_config import (
    load_deploy_config,
    merge_pipeline_deploy,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -967,24 +966,20 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
        diffusion_attention_backend="flash_attn",
        fa_deterministic=True,
        kv_cache_dtype="fp8",
        kv_cache_skip_steps="0-1",
        kv_cache_skip_layers=[2],
        static_lora_scale=0.25,
        diffusion_kv_mode="paged_scheduler",
        diffusers_load_kwargs=None,
        diffusers_call_kwargs=None,
    )
    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
    assert cfg.fa_deterministic is True
    assert cfg.diffusion_kv_cache_dtype == "fp8"
    assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
    assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
    assert cfg.lora_scale == 0.25
    assert cfg.cache_backend == "tea_cache"
    assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
    assert cfg.diffusers_load_kwargs == {}
    assert cfg.diffusers_call_kwargs == {}
@@ -999,8 +994,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                "stages:",
                "  - stage_id: 0",
                "    diffusion_attention_backend: flash_attn",
                "    fa_deterministic: true",
                "    diffusion_kv_mode: paged_scheduler",
            ]
        )
    )
@@ -1012,29 +1005,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
    assert isinstance(stage, VllmOmniDiffusionStageConfig)
    assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
    assert stage.diffusion_config.fa_deterministic is True
    assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
    deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
    deploy_path.write_text(
        "\n".join(
            [
                "pipeline: dreamzero",
                "async_chunk: false",
                "stages:",
                "  - stage_id: 0",
                "    diffusion_kv_mode: paged_worker_local",
            ]
        )
    )
    with pytest.raises(ValueError, match="reserved but not implemented"):
        _from_pipeline_key(
            "dreamzero",
            deploy_config_path=str(deploy_path),
        )
 def test_diffusion_config_field_classification_covers_current_fields():
@@ -1055,7 +1025,6 @@ def test_diffusion_config_field_classification_covers_current_fields():
        "enable_session_state_manager",
        "prompt_embed_cache_size",
        "diffusion_kv_cache_dtype",
        "diffusion_kv_mode",
    } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
    assert {
        "revision",
@@ -43,9 +43,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.0361,
+            "throughput_qps": 0.0434,
-            "latency_mean": 27.6206,
+            "latency_mean": 23.0654,
-            "peak_memory_mb_mean": 80548.0
+            "peak_memory_mb_mean": 76360.0
          }
        }
      }
@@ -100,9 +100,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.05263,
+            "throughput_qps": 0.0606,
-            "latency_mean": 18.9415,
+            "latency_mean": 16.4928,
-            "peak_memory_mb_mean": 50053.7
+            "peak_memory_mb_mean": 47164.8
          }
        }
      },
@@ -127,9 +127,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.009595,
+            "throughput_qps": 0.0124,
-            "latency_mean": 104.1058,
+            "latency_mean": 80.7784,
-            "peak_memory_mb_mean": 59718.0
+            "peak_memory_mb_mean": 55374.3429
          }
        }
      }
@@ -9,10 +9,6 @@ This script tests two main scenarios:
 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
 """
 import sys
 from types import SimpleNamespace
 from unittest.mock import Mock
 import pytest
 import torch
@@ -375,264 +371,6 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
        impl.forward_cuda(query, query, query, metadata)
 # ---------------------------------------------------------------------------
 # NPU mask-free packed attention (PR #5891).
 #
 # These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
 #   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
 #     sync, no MindIE-SD import);
 #   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
 #   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
 #
 # They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
 # into ``sys.modules`` (same pattern as ``test_paged_attention`` /
 # ``benchmarks/conftest``). No production code is exercised for real kernels.
 def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
    return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
 def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
    """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
    Restored automatically by monkeypatch. Defaults are no-op callables so the
    top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
    instance methods that would otherwise call them.
    """
    fake = SimpleNamespace(
        attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
        attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
    )
    monkeypatch.setitem(sys.modules, "mindiesd", fake)
    return fake
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 def test_resolve_packed_seq_accepts_real_plus_pad_contract():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
 def test_resolve_packed_seq_single_document_no_padding():
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 8], dtype=torch.int32)
    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
    assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
@pytest.mark.parametrize(
    "case",
    [
        "causal",
        "missing_key",
        "batch_ne_1",
        "cu_too_long",
        "cu_qk_mismatch",
        "max_seqlen_not_int",
        "used_gt_total",
        "used_zero",
        "pad_longer_than_real",
    ],
 )
 def test_resolve_packed_seq_rejects_malformed_metadata(case):
    """Any metadata outside the exact [real, pad] contract returns None so the
    caller falls back to the masked path."""
    cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
    cu2 = torch.tensor([0, 8], dtype=torch.int32)
    cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
    cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
    base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
    q = torch.randn(1, 8, 2, 4)
    impl = _npu_impl(causal=(case == "causal"))
    if case == "missing_key":
        extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5}  # 3 of 4 keys absent
    elif case == "batch_ne_1":
        q = torch.randn(2, 8, 2, 4)
        extra = base
    elif case == "cu_too_long":
        extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
    elif case == "cu_qk_mismatch":
        extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
    elif case == "max_seqlen_not_int":
        extra = {**base, "max_seqlen_q": 5.0}
    elif case == "used_gt_total":
        extra = {**base, "max_seqlen_q": 9}  # 9 > total length 8
    elif case == "used_zero":
        extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
    elif case == "pad_longer_than_real":
        # real=3, pad=5: the real document must be the longer one.
        extra = {
            "cu_seqlens_q": cu_short,
            "cu_seqlens_k": cu_short,
            "max_seqlen_q": 3,
            "max_seqlen_k": 3,
        }
    else:
        extra = base
    assert impl._resolve_packed_seq_npu(q, q, extra) is None
 # --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
 def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
    """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
    Lets a test assert which branch ``forward_fa_npu`` dispatched to without
    running any real kernel.
    """
    impl = _npu_impl()
    sentinel = torch.tensor([1.0])
    impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
    impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
    return impl, sentinel
 def test_npu_env_unset_uses_varlen(monkeypatch):
    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_env_laser_uses_slice(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_prefix_kv_slice_npu.assert_called_once()
    impl._forward_varlen_packed_npu.assert_not_called()
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
 def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
    """Locks in the CURRENT dispatch behavior (PR #5891).
    Any value other than ``ascend_laser_attention`` — including the valid
    ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
    the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
    comment #1 on the PR and is intentionally captured here as a regression
    net: update this test when the dispatch starts honoring/rejecting these.
    """
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
    _fake_mindiesd(monkeypatch)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
    impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
    impl._forward_varlen_packed_npu.assert_called_once()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
 def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
    """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
    the existing masked attention_forward path; the new branches are skipped."""
    fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
    _fake_mindiesd(monkeypatch, attention_forward=fake_forward)
    impl, _ = _npu_impl_with_mocked_paths()
    metadata = AttentionMetadata(
        attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
        extra={},  # no npu_attn_varlen opt-in
    )
    q = torch.randn(1, 8, 2, 4)
    out = impl.forward_fa_npu(q, q, q, metadata)
    impl._forward_varlen_packed_npu.assert_not_called()
    impl._forward_prefix_kv_slice_npu.assert_not_called()
    fake_forward.assert_called_once()
    assert out is fake_forward.return_value
 # --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
 def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured["key"] = key
        captured["value"] = value
        captured.update(kwargs)
        return torch.ones_like(query)  # deterministic stand-in kernel output
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()  # softmax_scale == 0.5
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        "laser_input_scale": 256.0,
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    # K/V sliced to the valid prefix (used_k == 5); q keeps full length.
    assert captured["key"].shape == (1, 5, 2, 4)
    assert captured["value"].shape == (1, 5, 2, 4)
    assert captured["query"].shape == (1, 8, 2, 4)
    assert captured["attn_mask"] is None
    # Pre-divide q/k/v by the factor and compensate the kernel scale by its
    # square (exact for power-of-two); output scaled back by the factor.
    torch.testing.assert_close(captured["query"], q / 256.0)
    torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
    assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
 def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
    captured: dict = {}
    def fake_attention_forward(query, key, value, **kwargs):
        captured["query"] = query
        captured.update(kwargs)
        return torch.ones_like(query)
    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
    impl = _npu_impl()
    q = torch.randn(1, 8, 2, 4)
    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
    extra = {
        "cu_seqlens_q": cu,
        "cu_seqlens_k": cu,
        "max_seqlen_q": 5,
        "max_seqlen_k": 5,
        # no laser_input_scale → no pre-scaling
    }
    out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
    torch.testing.assert_close(captured["query"], q)
    assert captured["scale"] == pytest.approx(0.5)
    torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
 if __name__ == "__main__":
    print("Running FlashAttention Padding Tests...")
    print("=" * 60)
@@ -5,9 +5,10 @@
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
-handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
+first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
-to the always-kept prefix segment before generating the block mask.
+mask from the video and prefix segments separately, so the two only agree on the
 block count when the video segment ends on a boundary.
 """
 import dataclasses
@@ -56,13 +57,10 @@ def test_block_aligned_video_segment_runs_sparse():
    assert plan.latent_shape == list(ALIGNED_GRID)
-def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
+def test_misaligned_video_segment_stays_dense():
-    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
+    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
-
+    # pad keys would take a share of every real query's softmax denominator.
-    assert plan is not None
+    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
    assert plan.prefix_len == PREFIX_ROWS
    assert plan.used_len == PREFIX_ROWS + 62496
    assert plan.latent_shape == list(MISALIGNED_GRID)
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -157,28 +155,3 @@ def test_fully_populated_mask_reproduces_dense_attention():
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
 def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
    torch.manual_seed(1)
    prefix_len, grid = 384, (4, 16, 10)
    video_len = grid[0] * grid[1] * grid[2]
    used = prefix_len + video_len
    heads, head_dim = 4, 128
    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
    impl = make_impl()
    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
    out = impl._forward_sparse_npu(q, k, v, plan)
    reference = torch.nn.functional.scaled_dot_product_attention(
        q.transpose(1, 2),
        k.transpose(1, 2),
        v.transpose(1, 2),
        scale=impl.softmax_scale,
    ).transpose(1, 2)
    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@@ -471,66 +471,6 @@ class TestTeaCacheBackend:
        assert backend.enabled is True
        mock_apply_hook.assert_called_once()
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_enable_uses_generic_default_threshold(self, mock_apply_hook):
        pipeline = Mock()
        pipeline.__class__.__name__ = "QwenImagePipeline"
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "QwenImageTransformer2DModel"
        TeaCacheBackend(DiffusionCacheConfig()).enable(pipeline)
        assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.2
    @pytest.mark.parametrize("partition", ["fl2va", "combined"])
    @pytest.mark.parametrize(
        ("configured_threshold", "expected_threshold"),
        [(None, 0.17), (0.2, 0.2)],
    )
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_minimax_h3_only_enables_fl2va_teacache(
        self,
        mock_apply_hook,
        partition,
        configured_threshold,
        expected_threshold,
    ):
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline.partition = partition
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
        pipeline.transformers_ref = Mock()
        config = (
            DiffusionCacheConfig()
            if configured_threshold is None
            else DiffusionCacheConfig(rel_l1_thresh=configured_threshold)
        )
        backend = TeaCacheBackend(config)
        backend.enable(pipeline)
        mock_apply_hook.assert_called_once()
        transformer, teacache_config = mock_apply_hook.call_args.args
        assert transformer is pipeline.transformer
        assert transformer is not pipeline.transformers_ref
        assert teacache_config.rel_l1_thresh == expected_threshold
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_minimax_h3_rejects_ref2va_teacache(self, mock_apply_hook):
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline.partition = "ref2va"
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
        backend = TeaCacheBackend(DiffusionCacheConfig())
        with pytest.raises(ValueError, match="only supports the MiniMax-H3 FL2VA partition"):
            backend.enable(pipeline)
        assert backend.enabled is False
        mock_apply_hook.assert_not_called()
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_enable_with_coefficients(self, mock_apply_hook):
        """Test enabling TeaCache with custom coefficients."""
@@ -15,23 +15,19 @@ Currently implemented:
 - TestFlux2KleinExtractor: Flux2Klein model extractor
 - TestFlux2Extractor: Flux2 model extractor
 - TestFluxExtractor: Flux model extractor
 - TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
 """
 from abc import ABC, abstractmethod
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 import torch.nn as nn
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import (
    extract_flux2_context,
    extract_flux2_klein_context,
    extract_flux_context,
    extract_minimax_h3_context,
 )
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
@@ -427,288 +423,3 @@ class TestFluxExtractor(BaseExtractorTest):
                img_ids=torch.randint(0, 64, (1, 16, 3)),
                txt_ids=torch.randint(0, 64, (1, 8, 3)),
            )
 class _MiniMaxH3FakeLinear(nn.Module):
    def __init__(
        self,
        in_features,
        out_features,
        *,
        bias=False,
        params_dtype=None,
        quant_config=None,
        prefix="",
        total_num_heads=None,
        total_num_kv_heads=None,
        **kwargs,
    ):
        del kwargs
        super().__init__()
        self.prefix = prefix
        self.quant_config = quant_config
        self.weight = nn.Parameter(torch.randn(out_features, in_features, dtype=params_dtype or torch.float32))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=params_dtype or torch.float32))
        else:
            self.register_parameter("bias", None)
        self.num_heads = total_num_heads
        self.num_kv_heads = total_num_kv_heads
    def forward(self, x):
        return torch.nn.functional.linear(x, self.weight, self.bias), None
 class _MiniMaxH3FakeMergedLinear(_MiniMaxH3FakeLinear):
    def __init__(self, in_features, out_features, **kwargs):
        super().__init__(in_features, sum(out_features), **kwargs)
 class _MiniMaxH3FakeQKVLinear(_MiniMaxH3FakeLinear):
    def __init__(
        self,
        *,
        hidden_size,
        head_size,
        total_num_heads,
        total_num_kv_heads,
        **kwargs,
    ):
        qkv_features = (total_num_heads + 2 * total_num_kv_heads) * head_size
        super().__init__(
            hidden_size,
            qkv_features,
            total_num_heads=total_num_heads,
            total_num_kv_heads=total_num_kv_heads,
            **kwargs,
        )
 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__()
        self.attn_backend = self.FakeBackend
        self.use_ring = False
    def forward(self, q, k, v, metadata=None):
        del k, v, metadata
        return q
 def _minimax_h3_small_od_config():
    arch = {
        "num_layers": 1,
        "token_refiner_num_layers": 1,
        "hidden_size": 8,
        "num_attention_heads": 1,
        "attention_head_dim": 8,
        "ffn_hidden_size": 16,
        "latents_dim": 2,
        "audio_latents_dim": 2,
        "patch_size": (1, 2, 2),
        "text_dim": 6,
        "timestep_input_dim": 4,
        "time_embed_hidden_size": 8,
        "time_embed_dim": 4,
        "adaln_out_features": 18 * 8,
        "final_adaln_out_features": 2 * 8,
        "rope_inv_freq_len": 1,
    }
    return SimpleNamespace(
        tf_model_config=arch,
        parallel_config=SimpleNamespace(ulysses_degree=1),
    )
 def _minimax_h3_sample_kwargs(*, seq_len: int = 8, device: torch.device | str = "cpu") -> dict:
    video_row_dim = 2 * 1 * 2 * 2
    audio_row_dim = 2
    text_len = 2
    img_rows = 2
    img_pos = torch.tensor([2, 3], dtype=torch.long, device=device)
    audio_pos = torch.tensor([4, 5], dtype=torch.long, device=device)
    text_pos = torch.tensor([0, 1], dtype=torch.long, device=device)
    infer_out_pos = img_pos.clone()
    token_tags = torch.full((seq_len,), -1, dtype=torch.long, device=device)
    token_tags[text_pos] = 1
    token_tags[img_pos] = 0
    token_tags[audio_pos] = 2
    inverse_indices = torch.zeros(seq_len, dtype=torch.long, device=device)
    return {
        "x": torch.randn(1, seq_len, video_row_dim, device=device),
        "audio_x": torch.randn(1, seq_len, audio_row_dim, device=device),
        "img_position_ids": torch.zeros(1, seq_len, 3, dtype=torch.float64, device=device),
        "unique_timesteps": torch.tensor([0.5], dtype=torch.float32, device=device),
        "inverse_indices": inverse_indices,
        "update_mask": torch.ones(img_rows, dtype=torch.bool, device=device),
        "token_tags": token_tags,
        "prompt_embeds": torch.randn(text_len, 6, device=device),
        "img_pos_info": {"position_ids": img_pos},
        "audio_pos_info": {"position_ids": audio_pos},
        "text_pos_info": {"position_ids": text_pos},
        "img_pos_for_infer_output_info": {"position_ids": infer_out_pos},
        "packed_seq_params": {
            "cu_seqlens_q": torch.tensor([0, seq_len], dtype=torch.int32, device=device),
            "max_seqlen_q": seq_len,
        },
        "refiner_packed_seq_params": {
            "cu_seqlens_q": torch.tensor([0, text_len], dtype=torch.int32, device=device),
            "max_seqlen_q": text_len,
        },
    }
@pytest.mark.cpu
 class TestMiniMaxH3Extractor(BaseExtractorTest):
    """Test extract_minimax_h3_context function."""
    def get_extractor(self):
        return extract_minimax_h3_context
    @pytest.fixture
    def minimax_h3_module(self, monkeypatch):
        """Create a minimal MiniMaxH3DiTModel with fake parallel linears for CPU tests."""
        from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
        monkeypatch.setattr(h3, "ColumnParallelLinear", _MiniMaxH3FakeLinear)
        monkeypatch.setattr(h3, "MergedColumnParallelLinear", _MiniMaxH3FakeMergedLinear)
        monkeypatch.setattr(h3, "QKVParallelLinear", _MiniMaxH3FakeQKVLinear)
        monkeypatch.setattr(h3, "RowParallelLinear", _MiniMaxH3FakeLinear)
        monkeypatch.setattr(h3, "Attention", _MiniMaxH3FakeAttention)
        monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
        model = h3.MiniMaxH3DiTModel(_minimax_h3_small_od_config(), quant_config=None)
        for submodule in model.modules():
            if isinstance(submodule, h3.MiniMaxH3Attention):
                submodule.rope._forward_method = submodule.rope.forward_native
        model.eval()
        return model
    def get_module(self, minimax_h3_module):
        return minimax_h3_module
    @pytest.fixture
    def sample_inputs(self):
        """Create sample packed forward kwargs for MiniMaxH3DiTModel."""
        return _minimax_h3_sample_kwargs(seq_len=8)
    def get_sample_inputs(self, sample_inputs):
        return sample_inputs
    def test_modulated_input_shape(self, minimax_h3_module, sample_inputs):
        """Test that modulated_input matches packed sequence shape [T, H]."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        seq_len = sample_inputs["x"].shape[1]
        assert context.modulated_input.shape == (seq_len, minimax_h3_module.hidden_size)
    def test_run_transformer_blocks_callable(self, minimax_h3_module, sample_inputs):
        """Test that run_transformer_blocks is callable."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        assert callable(context.run_transformer_blocks)
    def test_postprocess_callable(self, minimax_h3_module, sample_inputs):
        """Test that postprocess is callable."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        assert callable(context.postprocess)
    def test_postprocess_output_shape(self, minimax_h3_module, sample_inputs):
        """Test that postprocess returns video/audio logits for infer rows."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        (hidden,) = context.run_transformer_blocks()
        video_logits, audio_logits = context.postprocess(hidden)
        num_video_rows = sample_inputs["img_pos_info"]["position_ids"].shape[0]
        num_audio_rows = sample_inputs["audio_pos_info"]["position_ids"].shape[0]
        assert video_logits.shape[0] == num_video_rows
        assert audio_logits.shape[0] == num_audio_rows
    def test_extractor_matches_native_forward(self, minimax_h3_module, sample_inputs):
        """The uncached extractor path must preserve the model's forward result."""
        expected_video, expected_audio = minimax_h3_module(**sample_inputs)
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        (hidden,) = context.run_transformer_blocks()
        actual_video, actual_audio = context.postprocess(hidden)
        torch.testing.assert_close(actual_video, expected_video)
        torch.testing.assert_close(actual_audio, expected_audio)
    def test_run_transformer_blocks_forwards_packed_layout(
        self,
        minimax_h3_module,
        sample_inputs,
        monkeypatch,
    ):
        """The extractor block call must mirror the native packed contract."""
        from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
        video_layout = VideoTokenLayout(prefix_len=2, latent_grid=(1, 1, 2))
        inputs = {**sample_inputs, "video_token_layout": video_layout}
        block = minimax_h3_module.blocks[0]
        original_forward = block.forward
        received: list[tuple[int, VideoTokenLayout | None]] = []
        def capture_forward(*args, **kwargs):
            received.append((kwargs["packed_total"], kwargs["video_layout"]))
            return original_forward(*args, **kwargs)
        monkeypatch.setattr(block, "forward", capture_forward)
        context = extract_minimax_h3_context(minimax_h3_module, **inputs)
        context.run_transformer_blocks()
        assert received == [(sample_inputs["x"].shape[1], video_layout)]
    def test_teacache_reuses_h3_block_residual(self, minimax_h3_module, sample_inputs, monkeypatch):
        """A guaranteed cache hit must skip the H3 transformer block."""
        from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
        from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
        block = minimax_h3_module.blocks[0]
        original_forward = block.forward
        block_calls = 0
        def counted_forward(*args, **kwargs):
            nonlocal block_calls
            block_calls += 1
            return original_forward(*args, **kwargs)
        monkeypatch.setattr(block, "forward", counted_forward)
        hook = TeaCacheHook(
            TeaCacheConfig(
                transformer_type="MiniMaxH3DiTModel",
                rel_l1_thresh=0.3,
            )
        )
        hook.initialize_hook(minimax_h3_module)
        hook.new_forward(minimax_h3_module, **sample_inputs)
        video_logits, audio_logits = hook.new_forward(minimax_h3_module, **sample_inputs)
        assert block_calls == 1
        assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0]
        assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0]
    def test_invalid_module_raises_error(self):
        """Test that invalid module without blocks raises ValueError."""
        invalid_module = Mock()
        invalid_module.blocks = []
        with pytest.raises(ValueError, match="Module must have blocks"):
            extract_minimax_h3_context(invalid_module, **_minimax_h3_sample_kwargs())
@@ -1,52 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
 def _fixed_master_port(monkeypatch) -> None:
    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
 def test_dense_legacy_is_default() -> None:
    config = OmniDiffusionConfig.from_kwargs()
    assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
 def test_paged_worker_local_is_rejected_until_implemented() -> None:
    with pytest.raises(ValueError, match="reserved but not implemented"):
        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")
 def test_unknown_cache_mode_is_rejected() -> None:
    with pytest.raises(ValueError, match="Unsupported Diffusion KV diffusion_kv_mode"):
        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="unknown")
 def test_non_mapping_omni_kv_config_is_rejected() -> None:
    with pytest.raises(TypeError, match="omni_kv_config must be a mapping"):
        OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
 def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
    with pytest.raises(ValueError, match="does not support imported AR KV"):
        OmniDiffusionConfig.from_kwargs(
            diffusion_kv_mode="paged_scheduler",
            omni_kv_config={"need_recv_cache": True},
        )
 def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
    config = OmniDiffusionConfig.from_kwargs(
        model_class_name="FutureDiffusionModel",
        diffusion_kv_mode="paged_scheduler",
    )
    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
@@ -1,200 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from vllm.v1.core.kv_cache_manager import KVCacheManager
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 BLOCK_SIZE = 4
 def _request(**overrides) -> DiffusionKVRequest:
    values = dict(
        request_id="req/diffusion-kv/0",
        sequence_id=0,
        prefix_len=4,
        target_len=4,
        seq_len=8,
    )
    values.update(overrides)
    return DiffusionKVRequest(**values)
 def test_request_exposes_native_request_and_diffusion_semantics() -> None:
    request = _request()
    assert request.sequence_id == 0
    assert request.prefix_len == 4
    assert request.target_len == 4
    assert request.num_tokens == request.seq_len == 8
    assert request.num_prompt_tokens == 4
    assert request.num_computed_tokens == 0
    assert request.block_hashes == []
    assert request.kv_contexts == ()
    assert request.skip_reading_prefix_cache is True
    assert request.shared_prefix_boundary == 0
    assert request.status is RequestStatus.WAITING
@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"sequence_id": -1}, "sequence_id"),
        ({"prefix_len": -1}, "prefix_len"),
        ({"target_len": 0}, "target_len"),
        ({"seq_len": 0}, "seq_len"),
        ({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"),
    ],
 )
 def test_request_rejects_invalid_lengths(overrides, message: str) -> None:
    with pytest.raises(ValueError, match=message):
        _request(**overrides)
@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"context_id": ""}, "context_id"),
        ({"cache_role": ""}, "cache_role"),
        ({"num_tokens": 0}, "num_tokens"),
    ],
 )
 def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None:
    values = dict(context_id="text", cache_role="cross.text", num_tokens=8)
    values.update(overrides)
    with pytest.raises(ValueError, match=message):
        DiffusionKVContext(**values)
 def test_request_owns_independent_kv_contexts() -> None:
    text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
    image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4)
    request = _request(kv_contexts=(text, image))
    assert request.kv_contexts == (text, image)
 def test_request_rejects_duplicate_context_ids() -> None:
    contexts = (
        DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8),
        DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4),
    )
    with pytest.raises(ValueError, match="unique context_id"):
        _request(kv_contexts=contexts)
 def test_request_rejects_non_context_values() -> None:
    with pytest.raises(TypeError, match="DiffusionKVContext"):
        _request(kv_contexts=(object(),))
 def _manager(*, enable_caching: bool) -> KVCacheManager:
    # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
    # already registered an out-of-tree spec, making the registry non-empty
    # before FullAttentionSpec is registered. Initialize the native mappings
    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
    if register_all is not None:
        register_all(None)
    spec = FullAttentionSpec(
        block_size=BLOCK_SIZE,
        num_kv_heads=2,
        head_size=8,
        dtype=torch.bfloat16,
    )
    group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
    config = KVCacheConfig(
        num_blocks=8,
        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
        kv_cache_groups=[group],
    )
    return KVCacheManager(
        config,
        max_model_len=64,
        scheduler_block_size=BLOCK_SIZE,
        hash_block_size=BLOCK_SIZE,
        enable_caching=enable_caching,
    )
 def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None:
    manager = _manager(enable_caching=True)
    free_before = manager.block_pool.get_num_free_blocks()
    request = _request()
    # vLLM 0.26 adds a shared-prefix-boundary return value.
    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert blocks is not None
    assert len(manager.get_block_ids(request.request_id)[0]) == 2
    # Empty hashes intentionally mean request-local residency only. Calling
    # native cache_blocks() would require a canonical hash for the full prefix
    # block and is therefore outside this request's contract.
    manager.free(request)
    assert manager.block_pool.get_num_free_blocks() == free_before
 def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None:
    manager = _manager(enable_caching=True)
    prefix_hash = BlockHash(b"prefix".ljust(32, b"\0"))
    target_hash = BlockHash(b"target".ljust(32, b"\0"))
    request = _request(block_hashes=[prefix_hash, target_hash])
    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert blocks is not None
    manager.cache_blocks(request, request.prefix_len)
    manager.free(request)
    cached_request = _request(
        request_id="req/diffusion-kv/cached",
        block_hashes=[prefix_hash, target_hash],
    )
    cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request)
    # Only the stable prefix was published. The dynamic target remains a miss.
    assert num_computed_tokens == cached_request.prefix_len
    assert len(cached_blocks.blocks[0]) == 1
    cached_allocation = manager.allocate_slots(
        cached_request,
        num_new_tokens=cached_request.num_tokens - num_computed_tokens,
        num_new_computed_tokens=num_computed_tokens,
        new_computed_blocks=cached_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert cached_allocation is not None
    manager.free(cached_request)
@@ -1,447 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import pytest
 import torch
 from transformers.utils.generic import ModelOutput
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
    ImageInfo,
    JointImageInfo,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    HunyuanImage3Pipeline,
    get_hunyuan_image_3_pre_process_func,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
    HunyuanPreparedLayout,
    build_hunyuan_diffusion_kv_requests,
    extract_hunyuan_prompt_inputs,
    hunyuan_num_image_tokens,
    normalize_hunyuan_cot_text,
    prepare_hunyuan_layout,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def test_dense_legacy_preprocess_does_not_initialize_layout_tokenizer(monkeypatch) -> None:
    hf_config = SimpleNamespace(vae_downsample_factor=(8, 8), patch_size=2)
    image_processor = SimpleNamespace(vision_encoder_processor=SimpleNamespace(patch_size=1))
    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
    monkeypatch.setattr(
        pipeline_module,
        "TokenizerWrapper",
        lambda _model: pytest.fail("dense_legacy must not initialize the Diffusion KV tokenizer"),
    )
    get_hunyuan_image_3_pre_process_func(
        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
    )
 class _FakeTokenizerWrapper:
    def __init__(self, *, prefix_lens: list[int]) -> None:
        self.prefix_lens = prefix_lens
        self.calls: list[dict] = []
        self.pad_token_id = 0
        self.eos_token_id = 2
        self.boi_token_id = 3
        self.end_recaption_token_id = 4
        self.end_answer_token_id = 5
        self.special_token_map = {f"<img_ratio_{index}>": 100 + index for index in range(33)}
    def apply_chat_template(self, **kwargs):
        self.calls.append(kwargs)
        rows = len(self.prefix_lens)
        has_cond_image = kwargs.get("batch_cond_image_info") is not None
        seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
        padded_seq_len = max(seq_lens)
        positions = torch.arange(padded_seq_len)
        joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
        gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
        all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
        cond_vae_image_mask = (
            torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
        )
        cond_vit_image_mask = (
            torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
        )
        sections = []
        for _ in range(rows):
            row_sections = []
            if has_cond_image:
                row_sections.append(
                    {
                        "type": "joint_image",
                        "token_height": [1, 1],
                        "token_width": [2, 2],
                    }
                )
            row_sections.append({"type": "gen_image", "token_height": 4, "token_width": 4})
            sections.append(row_sections)
        return {
            "output": TokenizerEncodeOutput(
                tokens=torch.arange(rows * padded_seq_len, dtype=torch.long).reshape(rows, padded_seq_len),
                gen_timestep_scatter_index=torch.tensor(self.prefix_lens, dtype=torch.long).reshape(rows, 1),
                cond_timestep_scatter_index=(torch.zeros(rows, 1, dtype=torch.long) if has_cond_image else None),
                all_image_slices=all_image_slices,
                gen_image_mask=torch.stack(
                    [(positions >= length + 1) & (positions < length + 17) for length in self.prefix_lens]
                ),
                cond_vae_image_mask=cond_vae_image_mask,
                cond_vit_image_mask=cond_vit_image_mask,
                joint_image_slices=joint_image_slices,
                gen_image_slices=gen_image_slices,
                real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
            ),
            "sections": sections,
        }
 class _FakeImageProcessor:
    def __init__(self, image_token_length: int = 16) -> None:
        self.image_token_length = image_token_length
        self.image_sizes: list[tuple[int, int]] = []
        self.vision_encoder_processor = SimpleNamespace(patch_size=1)
    def build_image_info(self, image_size):
        self.image_sizes.append(image_size)
        return ImageInfo(
            image_type="gen_image",
            image_width=image_size[1],
            image_height=image_size[0],
            token_width=4,
            token_height=4,
            image_token_length=self.image_token_length,
            base_size=1024,
            ratio_index=0,
        )
 def _components(prefix_lens: list[int]):
    return _FakeTokenizerWrapper(prefix_lens=prefix_lens), _FakeImageProcessor()
 def _prepare(
    request: OmniDiffusionRequest,
    tokenizer: _FakeTokenizerWrapper,
    image_processor: _FakeImageProcessor,
 ) -> HunyuanPreparedLayout:
    prepared_layout = prepare_hunyuan_layout(
        request,
        tokenizer_wrapper=tokenizer,
        image_processor=image_processor,
        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
        image_base_size=1024,
    )
    request.prepared_layout = prepared_layout
    return prepared_layout
 def _request(
    *,
    guidance_scale: float,
    prompt="draw a cat",
    request_id: str = "req",
 ) -> OmniDiffusionRequest:
    return OmniDiffusionRequest(
        prompt=prompt,
        sampling_params=OmniDiffusionSamplingParams(
            height=768,
            width=1024,
            guidance_scale=guidance_scale,
            num_inference_steps=4,
        ),
        request_id=request_id,
    )
 def test_builds_kv_request_lengths_without_model_execution() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=1.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
    assert all(isinstance(item, DiffusionKVRequest) for item in kv_requests)
    assert [(item.prefix_len, item.target_len, item.seq_len) for item in kv_requests] == [(12, 17, 32)]
    assert image_processor.image_sizes == [(768, 1024)]
    assert tokenizer.calls[0]["sequence_template"] == "instruct"
    assert tokenizer.calls[0]["cfg_factor"] == 1
    assert kv_requests[0].block_hashes == []
    assert kv_requests[0].kv_contexts == ()
    assert kv_requests[0].skip_reading_prefix_cache is True
 def test_builds_one_kv_request_per_cfg_row() -> None:
    tokenizer, image_processor = _components([12, 14])
    request = _request(guidance_scale=5.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert [item.sequence_id for item in kv_requests] == [0, 1]
    assert [item.prefix_len for item in kv_requests] == [12, 14]
    assert [item.seq_len for item in kv_requests] == [32, 34]
    assert tokenizer.calls[0]["cfg_factor"] == 2
 def _reference_image() -> JointImageInfo:
    return JointImageInfo(
        vae_image_info=ImageInfo(
            image_type="vae",
            token_width=8,
            token_height=8,
            image_token_length=64,
        ),
        vision_image_info=ImageInfo(
            image_type="siglip2",
            token_width=4,
            token_height=4,
            image_token_length=16,
        ),
        vision_encoder_kwargs={
            "spatial_shapes": torch.tensor([4, 4]),
            "pixel_attention_mask": torch.ones(4, dtype=torch.bool),
        },
    )
 def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
    tokenizer, image_processor = _components([20])
    joint_image = _reference_image()
    prompt = {
        "prompt": "edit this image",
        "additional_information": {"batch_cond_image_info": [joint_image]},
    }
    request = _request(guidance_scale=1.0, prompt=prompt)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert tokenizer.calls[0]["batch_cond_image_info"] == [[joint_image]]
    assert kv_requests[0].prefix_len == 20
    assert kv_requests[0].target_len == 17
    assert kv_requests[0].seq_len == 40
    assert kv_requests[0].kv_contexts == ()
 def _assert_nested_equal(actual, expected) -> None:
    if isinstance(expected, torch.Tensor):
        torch.testing.assert_close(actual, expected)
    elif isinstance(expected, dict):
        assert actual.keys() == expected.keys()
        for key in expected:
            _assert_nested_equal(actual[key], expected[key])
    elif isinstance(expected, list | tuple):
        assert isinstance(actual, type(expected))
        assert len(actual) == len(expected)
        for actual_item, expected_item in zip(actual, expected):
            _assert_nested_equal(actual_item, expected_item)
    else:
        assert actual == expected
@pytest.mark.parametrize(
    ("prefix_lens", "guidance_scale", "with_reference_image"),
    [
        pytest.param([12], 1.0, False, id="cfg-1"),
        pytest.param([12, 14], 5.0, False, id="cfg-2"),
        pytest.param([20], 1.0, True, id="cfg-1-reference-image"),
        pytest.param([20, 22], 5.0, True, id="cfg-2-reference-image"),
    ],
 )
 def test_prepared_model_inputs_match_local_tokenization(
    monkeypatch,
    prefix_lens: list[int],
    guidance_scale: float,
    with_reference_image: bool,
 ) -> None:
    tokenizer, image_processor = _components(prefix_lens)
    prompt: str | dict = "draw a cat"
    if with_reference_image:
        prompt = {
            "prompt": "edit this image",
            "additional_information": {"batch_cond_image_info": [_reference_image()]},
        }
    request = _request(guidance_scale=guidance_scale, prompt=prompt)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    prompts, cot_texts, system_prompt, batch_cond_image_info, bot_task = extract_hunyuan_prompt_inputs(
        [request.prompt],
        request.sampling_params.extra_args or {},
        request_id=request.request_id,
        allow_cond_image=True,
    )
    cot_text = (
        [normalize_hunyuan_cot_text(text) for text in cot_texts]
        if any(text is not None for text in cot_texts)
        else None
    )
    pipeline = object.__new__(HunyuanImage3Pipeline)
    pipeline._tkwrapper = tokenizer
    pipeline.image_processor = image_processor
    pipeline.generation_config = SimpleNamespace(sequence_template="instruct", drop_think=False)
    pipeline.config = SimpleNamespace(
        image_base_size=1024,
        attention_head_dim=2,
        rope_theta=10000.0,
    )
    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda _self: torch.device("cpu")))
    def fake_encode_cond_image(batch_cond_image_info_list, cfg_factor=1, generator=None):
        del generator
        rows = len(batch_cond_image_info_list) * cfg_factor
        return (
            torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 1, 2, 2),
            torch.arange(rows, dtype=torch.float32),
            torch.arange(rows * 3, dtype=torch.float32).reshape(rows, 3),
        )
    pipeline._encode_cond_image = fake_encode_cond_image
    def fake_build_batch_2d_rope(*, image_infos, seq_len, **_kwargs):
        geometry_sum = sum(
            int(image_slice.start or 0) + int(image_slice.stop or 0) + height + width
            for row in image_infos
            for image_slice, (height, width) in row
        )
        values = torch.tensor([len(image_infos), seq_len, geometry_sum], dtype=torch.float32)
        return values, values + 1
    monkeypatch.setattr(pipeline_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
    common_kwargs = dict(
        prompt=prompts,
        cot_text=cot_text,
        system_prompt=system_prompt,
        mode="gen_image",
        guidance_scale=guidance_scale,
        image_size=(request.sampling_params.height, request.sampling_params.width),
        generator=[torch.Generator().manual_seed(0)],
        batch_cond_image_info=batch_cond_image_info,
        bot_task=bot_task,
    )
    local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
    prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
    comparable_fields = (
        "input_ids",
        "position_ids",
        "custom_pos_emb",
        "image_mask",
        "gen_timestep_scatter_index",
        "cond_vae_images",
        "cond_timestep",
        "cond_vae_image_mask",
        "cond_vit_images",
        "cond_vit_image_mask",
        "vit_kwargs",
        "cond_timestep_scatter_index",
    )
    for field in comparable_fields:
        _assert_nested_equal(prepared_inputs[field], local_inputs[field])
    tokenizer_fields = (
        "tokens",
        "gen_image_mask",
        "cond_vae_image_mask",
        "cond_vit_image_mask",
        "gen_timestep_scatter_index",
        "cond_timestep_scatter_index",
        "real_pos",
        "all_image_slices",
        "joint_image_slices",
        "gen_image_slices",
    )
    for field in tokenizer_fields:
        _assert_nested_equal(
            getattr(prepared_inputs["tokenizer_output"], field),
            getattr(local_inputs["tokenizer_output"], field),
        )
    local_mask = pipeline._prepare_attention_mask_for_generation(
        local_inputs["input_ids"],
        pipeline.generation_config,
        local_inputs,
    )
    prepared_mask = pipeline._prepare_attention_mask_for_generation(
        prepared_inputs["input_ids"],
        pipeline.generation_config,
        prepared_inputs,
    )
    torch.testing.assert_close(prepared_mask, local_mask)
    assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
    num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
    local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
    prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
    local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
    prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
    for field in ("position_ids", "attention_mask", "gen_timestep_scatter_index", "full_attn_spans"):
        _assert_nested_equal(prepared_step_inputs[field], local_step_inputs[field])
 def test_rejects_tokenizer_cfg_row_mismatch() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=5.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    with pytest.raises(ValueError, match="sequence count does not match"):
        build_hunyuan_diffusion_kv_requests(request, prepared_layout)
 def test_sender_endpoint_does_not_change_local_kv_request() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=1.0)
    request.kv_sender_info = {"host": "127.0.0.1", "zmq_port": 5000}
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert kv_requests[0].seq_len == 32
    assert kv_requests[0].num_computed_tokens == 0
 def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch) -> None:
    tokenizer, image_processor = _components([12])
    hf_config = SimpleNamespace(
        vae_downsample_factor=(8, 8),
        patch_size=2,
        image_base_size=1024,
    )
    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
    monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
    monkeypatch.setattr(
        pipeline_module.GenerationConfig,
        "from_pretrained",
        lambda _model: SimpleNamespace(sequence_template="instruct", drop_think=False),
    )
    preprocess = get_hunyuan_image_3_pre_process_func(
        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
    )
    request = _request(guidance_scale=1.0)
    prepared_request = preprocess(request)
    assert prepared_request is request
    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
    assert request.diffusion_kv_requests is not None
    assert [(item.prefix_len, item.target_len, item.seq_len) for item in request.diffusion_kv_requests] == [
        (12, 17, 32)
    ]
    assert len(tokenizer.calls) == 1
@@ -7,10 +7,7 @@ import pytest
 import torch
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module
 import vllm_omni.diffusion.models.hunyuan_image3.request_layout as hy3_layout_module
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageInfo
 from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    _STEP_AR_KV,
    _STEP_CFG_FACTOR,
@@ -21,7 +18,6 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    _STEP_PROMPT_KV,
    HunyuanImage3Pipeline,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.request_layout import HunyuanPreparedLayout
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import StepRequestState
@@ -80,72 +76,6 @@ def _sampling_params(**extra_args):
    )
 def _prepared_layout() -> HunyuanPreparedLayout:
    return HunyuanPreparedLayout(
        tokenizer_output=TokenizerEncodeOutput(
            tokens=torch.arange(21).reshape(1, 21),
            gen_image_mask=torch.ones(1, 21, dtype=torch.bool),
            gen_timestep_scatter_index=torch.tensor([[4]]),
            all_image_slices=[[slice(5, 21)]],
            joint_image_slices=[[]],
            gen_image_slices=[[slice(5, 21)]],
            real_pos=torch.tensor([[21]]),
        ),
        rope_image_info=[[(slice(5, 21), (4, 4))]],
        generated_image_info=ImageInfo(
            image_type="gen_image",
            image_width=512,
            image_height=512,
            token_width=4,
            token_height=4,
            image_token_length=16,
        ),
    )
 def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
    pipeline = _pipeline()
    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
    rope_kwargs = {}
    def fake_build_batch_2d_rope(**kwargs):
        rope_kwargs.update(kwargs)
        return torch.zeros(1), torch.zeros(1)
    monkeypatch.setattr(hy3_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
    def fail_apply_chat_template(**_kwargs):
        pytest.fail("prepared Hunyuan execution must not apply the chat template again")
    pipeline._tkwrapper = SimpleNamespace(
        apply_chat_template=fail_apply_chat_template,
        eos_token_id=2,
        boi_token_id=3,
        end_recaption_token_id=4,
        end_answer_token_id=5,
    )
    pipeline.config = SimpleNamespace(
        attention_head_dim=2,
        rope_theta=10000.0,
    )
    prepared = _prepared_layout()
    model_inputs = pipeline.prepare_model_inputs(
        prompt="prompt",
        mode="gen_image",
        guidance_scale=1.0,
        image_size=(512, 512),
        device=torch.device("cpu"),
        generator=[torch.Generator().manual_seed(0)],
        prepared_layout=prepared,
    )
    assert model_inputs["tokenizer_output"] is prepared.tokenizer_output
    assert model_inputs["batch_gen_image_info"] == [prepared.generated_image_info]
    assert rope_kwargs["image_infos"] is prepared.rope_image_info
    torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
    pipeline = _pipeline()
    states = [_state("req-0", 1), _state("req-1", 3)]
@@ -201,13 +131,12 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
        captured.update(kwargs)
        raise RuntimeError("stop after prepare_model_inputs")
-    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
    pipeline.prepare_model_inputs = fake_prepare_model_inputs
    state = StepRequestState(
        request_id="req-bot-task",
        sampling=sampling,
        prompt=prompt_item,
        prepared_layout=_prepared_layout(),
    )
    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
@@ -215,7 +144,6 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
    assert captured["bot_task"] == expected_model_bot_task
    assert captured["system_prompt_bot_task"] == expected_system_bot_task
    assert captured["prepared_layout"] is state.prepared_layout
 def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
@@ -231,22 +159,23 @@ def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
        captured.update(kwargs)
        raise RuntimeError("stop after prepare_model_inputs")
-    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
    pipeline.prepare_model_inputs = fake_prepare_model_inputs
-    request = SimpleNamespace(
+    req = DiffusionRequestBatch(
-        request_id="req-forward-bot-task",
+        requests=[
-        sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
+            SimpleNamespace(
-        prompt={"prompt": "prompt", "bot_task": "vanilla"},
+                request_id="req-forward-bot-task",
-        prepared_layout=_prepared_layout(),
+                sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
                prompt={"prompt": "prompt", "bot_task": "vanilla"},
            )
        ]
    )
    req = DiffusionRequestBatch(requests=[request])
    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
        pipeline.forward(req)
    assert captured["bot_task"] == "think"
    assert captured["system_prompt_bot_task"] == "think"
    assert captured["prepared_layout"] is request.prepared_layout
 def test_grouped_denoise_rejects_non_sdpa_attention_backend():
@@ -1,133 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 from vllm_omni.diffusion.models.lingbot_video.image_condition import (
    apply_clean_prefix,
    geometry_align_image,
    prepare_ti2v_image_condition,
 )
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _gradient_image(width: int, height: int) -> Image.Image:
    x = np.linspace(0, 255, width, dtype=np.uint8)[None, :]
    y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
    rgb = np.stack(
        [
            np.broadcast_to(x, (height, width)),
            np.broadcast_to(y, (height, width)),
            np.full((height, width), 127, dtype=np.uint8),
        ],
        axis=-1,
    )
    return Image.fromarray(rgb, mode="RGB")
@pytest.mark.parametrize("source_size", [(160, 64), (64, 160), (96, 96)])
 def test_geometry_align_image_supports_source_aspect_ratios(source_size):
    pixel = geometry_align_image(_gradient_image(*source_size), 64, 96)
    assert pixel.shape == (1, 3, 1, 64, 96)
    assert pixel.dtype == torch.float32
    assert 0 <= pixel.min() <= pixel.max() <= 1
@pytest.mark.parametrize(
    ("source_size", "target_size"),
    [((160, 64), (64, 96)), ((64, 160), (96, 64))],
 )
 def test_geometry_align_image_center_crops_without_stretching(source_size, target_size):
    width, height = source_size
    target_height, target_width = target_size
    rgb = np.full((height, width, 3), (255, 0, 0), dtype=np.uint8)
    if width > target_width:
        left = (width - target_width) // 2
        rgb[:, left : left + target_width] = (0, 255, 0)
    else:
        top = (height - target_height) // 2
        rgb[top : top + target_height] = (0, 255, 0)
    pixel = geometry_align_image(
        Image.fromarray(rgb, mode="RGB"),
        target_height,
        target_width,
    )
    assert pixel.shape == (1, 3, 1, target_height, target_width)
    expected_green = torch.tensor((0.0, 1.0, 0.0)).view(1, 3, 1, 1, 1)
    assert torch.equal(pixel, expected_green.expand_as(pixel))
 class _LatentDistribution:
    def __init__(self, latent: torch.Tensor):
        self.latent = latent
    def sample(self, generator=None):
        del generator
        return self.latent
 class _CapturingVAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.anchor = nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(
            latents_mean=[0.0, 0.0, 0.0],
            latents_std=[1.0, 1.0, 1.0],
        )
        self.encoded_pixel = None
    def encode(self, pixel):
        self.encoded_pixel = pixel.detach().clone()
        return SimpleNamespace(latent_dist=_LatentDistribution(pixel))
 def test_ti2v_vlm_and_vae_paths_share_geometry_aligned_image():
    vae = _CapturingVAE()
    condition = prepare_ti2v_image_condition(
        _gradient_image(160, 64),
        height=64,
        width=96,
        vae=vae,
        vision_patch_size=16,
        device=torch.device("cpu"),
    )
    vae_pixel = (vae.encoded_pixel + 1.0) / 2.0
    vae_image = vae_pixel[0, :, 0].permute(1, 2, 0).mul(255).byte().numpy()
    assert condition.vlm_image.size == (96, 64)
    np.testing.assert_allclose(
        np.asarray(condition.vlm_image),
        vae_image,
        rtol=0,
        atol=1,
    )
 def test_apply_clean_prefix_only_overwrites_condition_frames():
    latents = torch.arange(20, dtype=torch.float32).reshape(1, 1, 5, 2, 2)
    original_tail = latents[:, :, 2:].clone()
    clean = torch.full((1, 1, 2, 2, 2), -3.0)
    returned = apply_clean_prefix(latents, clean)
    assert returned is latents
    assert torch.equal(latents[:, :, :2], clean)
    assert torch.equal(latents[:, :, 2:], original_tail)
 def test_apply_clean_prefix_rejects_incompatible_shapes():
    with pytest.raises(ValueError, match="shape must match"):
        apply_clean_prefix(
            torch.zeros(1, 2, 3, 4, 4),
            torch.zeros(1, 3, 1, 4, 4),
        )
@@ -6,10 +6,9 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
-from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -23,8 +22,6 @@ def _make_pipeline():
    nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.default_negative_prompt = "default negative"
    pipeline.default_image_negative_prompt = "default image negative"
    pipeline.img_prompt_template = "<image>"
    pipeline.od_config = SimpleNamespace(flow_shift=None)
    return pipeline
@@ -40,24 +37,17 @@ def test_lingbot_video_pipeline_import_and_registry():
        LingBotVideoPipeline,
        LingBotVideoTransformer3DModel,
        get_lingbot_video_post_process_func,
        get_lingbot_video_pre_process_func,
    )
    from vllm_omni.diffusion.registry import (
        _DIFFUSION_MODELS,
        _DIFFUSION_POST_PROCESS_FUNCS,
        _DIFFUSION_PRE_PROCESS_FUNCS,
    )
    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
    assert LingBotVideoPipeline is not None
    assert LingBotVideoTransformer3DModel is not None
    assert get_lingbot_video_pre_process_func is not None
    assert get_lingbot_video_post_process_func is not None
    assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
        "lingbot_video",
        "pipeline_lingbot_video",
        "LingBotVideoPipeline",
    )
    assert _DIFFUSION_PRE_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_pre_process_func"
    assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
@@ -68,7 +58,6 @@ def test_component_discovery_declarations():
    assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
    assert LingBotVideoPipeline._vae_modules == ["vae"]
    assert LingBotVideoPipeline.supports_step_execution is False
    assert LingBotVideoPipeline.max_outputs_per_prompt == 1
 def test_extra_body_params_include_video_flow_shift_alias():
@@ -80,47 +69,17 @@ def test_extra_body_params_include_video_flow_shift_alias():
    assert "batch_cfg" in params
-def test_preprocess_preserves_geometry_and_converts_rgb(tmp_path):
+def test_extract_prompt_string_and_mapping():
-    from vllm_omni.diffusion.models.lingbot_video import (
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
        get_lingbot_video_pre_process_func,
    )
-    image_path = tmp_path / "reference.png"
+    string_req = SimpleNamespace(prompt="a robot waves")
-    Image.new("L", (48, 32), color=127).save(image_path)
+    assert _extract_prompt(string_req) == ("a robot waves", None)
    request = OmniDiffusionRequest(
        prompt={
            "prompt": "a robot walks",
            "modalities": ["video"],
            "multi_modal_data": {"image": str(image_path)},
        },
        sampling_params=OmniDiffusionSamplingParams(),
        request_id="lingbot-preprocess-test",
    )
-    result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
+    mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
    assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
-    image = result.prompt["multi_modal_data"]["image"]
+    with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
-    assert image.mode == "RGB"
+        _extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
    assert image.size == (48, 32)
 def test_preprocess_marks_internal_image_warmup_as_video():
    from vllm_omni.diffusion.models.lingbot_video import (
        get_lingbot_video_pre_process_func,
    )
    request = OmniDiffusionRequest(
        prompt={
            "prompt": "dummy run",
            "multi_modal_data": {"image": Image.new("RGB", (32, 32))},
        },
        sampling_params=OmniDiffusionSamplingParams(),
        request_id=DUMMY_DIFFUSION_REQUEST_ID,
    )
    result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
    assert result.prompt["modalities"] == ["video"]
 def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
@@ -145,59 +104,6 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
    array = post(frames, SimpleNamespace(output_type="np"))
    assert isinstance(array, np.ndarray)
    assert array.shape == (2, 4, 4, 3)
    video = post({"video": frames}, SimpleNamespace(output_type="pt"))
    assert set(video) == {"video"}
    assert video["video"] is frames
    image = post(
        {"image": torch.ones(4, 4, 3)},
        SimpleNamespace(output_type="pt"),
    )
    assert set(image) == {"image"}
    assert isinstance(image["image"], np.ndarray)
    assert image["image"].shape == (4, 4, 3)
    latent = torch.ones(1, 4, 1, 2, 2)
    latent_output = post({"image": latent}, SimpleNamespace(output_type="latent"))
    assert set(latent_output) == {"image"}
    assert latent_output["image"] is latent
 def test_postprocess_preserves_image_contract_for_serving():
    from vllm_omni.diffusion.data import DiffusionOutput
    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
    from vllm_omni.diffusion.output_formatter import (
        format_diffusion_outputs,
        normalize_diffusion_postprocess_output,
    )
    from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result
    sampling = OmniDiffusionSamplingParams(
        num_inference_steps=1,
        output_type="pt",
        resolution=64,
    )
    request = OmniDiffusionRequest(
        prompt={"prompt": "a robot", "modalities": ["image"]},
        sampling_params=sampling,
        request_id="lingbot-image-contract-test",
    )
    raw_output = {"image": torch.ones(8, 12, 3)}
    processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
    normalized = normalize_diffusion_postprocess_output(processed)
    assert normalized.primary_key == "image"
    [result] = format_diffusion_outputs(
        request=request,
        od_config=SimpleNamespace(model_class_name="LingBotVideoPipeline"),
        diffusion_output=DiffusionOutput(output=raw_output),
        output_data=raw_output,
        postprocess_output=normalized,
    )
    assert len(result.images) == 1
    assert isinstance(result.images[0], np.ndarray)
    endpoint_images = _extract_images_from_result(result)
    assert len(endpoint_images) == 1
    assert endpoint_images[0].size == (12, 8)
 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
@@ -210,11 +116,7 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
    pipeline._generate = fake_generate
    req = _make_request_batch(
-        {
+        {"prompt": "a robot arm", "negative_prompt": "bad hands"},
            "prompt": "a robot arm",
            "negative_prompt": "bad hands",
            "modalities": ["video"],
        },
        height=192,
        width=320,
        num_frames=9,
@@ -227,12 +129,10 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
    out = pipeline.forward(req)
-    assert torch.equal(out.output["video"], torch.zeros(9, 192, 320, 3))
+    assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
    assert len(calls) == 1
    call = calls[0]
    assert call["prompt"] == "a robot arm"
    assert call["mode"].value == "t2v"
    assert call["input_image"] is None
    assert call["negative_prompt"] == "bad hands"
    assert call["height"] == 192
    assert call["width"] == 320
@@ -294,222 +194,6 @@ def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
    assert calls[0]["negative_prompt"] == "default negative"
@pytest.mark.parametrize(
    ("prompt", "num_frames", "expected_mode", "output_key", "default_negative"),
    [
        (
            {"prompt": "a still robot", "modalities": ["image"]},
            1,
            "t2i",
            "image",
            "default image negative",
        ),
        (
            {
                "prompt": "a moving robot",
                "modalities": ["video"],
                "multi_modal_data": {"image": Image.new("RGB", (48, 32), color="blue")},
            },
            9,
            "ti2v",
            "video",
            "default negative",
        ),
    ],
 )
 def test_forward_dispatches_generation_mode_and_output(
    prompt,
    num_frames,
    expected_mode,
    output_key,
    default_negative,
 ):
    pipeline = _make_pipeline()
    calls = []
    def fake_generate(**kwargs):
        calls.append(kwargs)
        return torch.zeros(192, 320, 3) if output_key == "image" else torch.zeros(num_frames, 192, 320, 3)
    pipeline._generate = fake_generate
    output = pipeline.forward(
        _make_request_batch(
            prompt,
            height=192,
            width=320,
            num_frames=num_frames,
            num_inference_steps=2,
            guidance_scale=3.0,
            seed=42,
            output_type="pt",
        )
    )
    assert set(output.output) == {output_key}
    assert calls[0]["mode"].value == expected_mode
    assert calls[0]["negative_prompt"] == default_negative
    if expected_mode == "ti2v":
        assert isinstance(calls[0]["input_image"], Image.Image)
 def test_batched_cfg_pads_image_conditioned_embeddings_and_masks():
    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
        _batch_cfg_prompt_inputs,
    )
    positive = torch.ones(1, 7, 4)
    positive_mask = torch.ones(1, 7, dtype=torch.long)
    negative = torch.full((1, 9, 4), 2.0)
    negative_mask = torch.ones(1, 9, dtype=torch.long)
    embeds, mask = _batch_cfg_prompt_inputs(
        positive,
        positive_mask,
        negative,
        negative_mask,
    )
    assert embeds.shape == (2, 9, 4)
    assert mask.shape == (2, 9)
    assert torch.equal(mask[0], torch.tensor([1, 1, 1, 1, 1, 1, 1, 0, 0]))
    assert torch.equal(mask[1], torch.ones(9, dtype=torch.long))
 class _RecordingTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.anchor = nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(in_channels=1)
        self.latent_inputs = []
    def forward(
        self,
        latents,
        timestep,
        prompt_embeds,
        *,
        encoder_attention_mask,
        return_dict,
    ):
        del timestep, prompt_embeds, encoder_attention_mask, return_dict
        self.latent_inputs.append(latents.detach().clone())
        return (torch.zeros_like(latents),)
 class _CorruptingScheduler:
    sigma_max = 1.0
    sigma_min = 0.0
    def set_timesteps(self, num_inference_steps, *, device, shift, **kwargs):
        del num_inference_steps, shift, kwargs
        self.timesteps = torch.tensor([1000.0, 500.0], device=device)
    def step(
        self,
        noise_pred,
        timestep,
        latents,
        *,
        return_dict,
        generator,
    ):
        del noise_pred, timestep, return_dict, generator
        return (latents + 7.0,)
@pytest.mark.parametrize(
    ("guidance_scale", "batch_cfg", "expected_prompts", "expected_transformer_calls"),
    [
        pytest.param(2.0, False, ["positive", "negative"], 4, id="sequential-cfg"),
        pytest.param(2.0, True, ["positive", "negative"], 2, id="batched-cfg"),
        pytest.param(1.0, False, ["positive"], 2, id="cfg-disabled"),
    ],
 )
 def test_ti2v_reinjects_clean_prefix_across_cfg_modes(
    guidance_scale,
    batch_cfg,
    expected_prompts,
    expected_transformer_calls,
 ):
    from vllm_omni.diffusion.models.lingbot_video import (
        LingBotGenerationMode,
        LingBotImageCondition,
    )
    pipeline = _make_pipeline()
    pipeline.transformer = _RecordingTransformer()
    pipeline.scheduler = _CorruptingScheduler()
    pipeline.progress_bar = lambda timesteps: timesteps
    pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 3, 1, 1)
    vlm_image = Image.new("RGB", (32, 32), color="red")
    condition = LingBotImageCondition(
        vlm_image=vlm_image,
        clean_latent=torch.full((1, 1, 1, 1, 1), 5.0),
    )
    pipeline.prepare_ti2v_image_condition = lambda *args, **kwargs: condition
    encode_calls = []
    def fake_encode(prompt, *, images, device):
        del device
        encode_calls.append((prompt, images))
        return torch.ones(1, 2, 4), torch.ones(1, 2, dtype=torch.long)
    pipeline.encode_prompt = fake_encode
    latents = pipeline._generate(
        prompt="positive",
        mode=LingBotGenerationMode.TI2V,
        input_image=vlm_image,
        negative_prompt="negative",
        height=16,
        width=16,
        num_frames=9,
        num_inference_steps=2,
        guidance_scale=guidance_scale,
        shift=3.0,
        output_type="latent",
        batch_cfg=batch_cfg,
    )
    assert [call[0] for call in encode_calls] == expected_prompts
    assert all(call[1] == [vlm_image] for call in encode_calls)
    assert len(pipeline.transformer.latent_inputs) == expected_transformer_calls
    for item in pipeline.transformer.latent_inputs:
        expected_prefix = condition.clean_latent.expand(item.shape[0], -1, -1, -1, -1)
        assert torch.equal(item[:, :, :1], expected_prefix)
    assert torch.equal(latents[:, :, :1], condition.clean_latent)
    assert torch.equal(latents[:, :, 1:], torch.full((1, 1, 2, 1, 1), 14.0))
 def test_t2i_decodes_and_returns_the_unique_frame():
    from vllm_omni.diffusion.models.lingbot_video import LingBotGenerationMode
    pipeline = _make_pipeline()
    pipeline.transformer = _RecordingTransformer()
    pipeline.scheduler = _CorruptingScheduler()
    pipeline.progress_bar = lambda timesteps: timesteps
    pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 1, 2, 2)
    pipeline.encode_prompt = lambda *args, **kwargs: (
        torch.ones(1, 2, 4),
        torch.ones(1, 2, dtype=torch.long),
    )
    pipeline._decode_latents = lambda latents: torch.ones(1, 16, 16, 3)
    image = pipeline._generate(
        prompt="a still robot",
        mode=LingBotGenerationMode.T2I,
        height=16,
        width=16,
        num_frames=1,
        num_inference_steps=2,
        guidance_scale=1.0,
        shift=3.0,
        output_type="pt",
    )
    assert image.shape == (16, 16, 3)
 def test_forward_rejects_multi_request_batches():
    pipeline = _make_pipeline()
    first = OmniDiffusionRequest(
@@ -527,45 +211,6 @@ def test_forward_rejects_multi_request_batches():
        pipeline.forward(DiffusionRequestBatch([first, second]))
 def test_forward_marks_invalid_request_contract_as_client_error():
    from vllm_omni.errors import OmniClientError
    pipeline = _make_pipeline()
    request = _make_request_batch(
        {
            "prompt": "unsupported image edit",
            "modalities": ["image"],
            "multi_modal_data": {"image": Image.new("RGB", (32, 32), color="blue")},
        },
        height=192,
        width=320,
        num_frames=1,
    )
    with pytest.raises(OmniClientError, match="does not accept") as exc_info:
        pipeline.forward(request)
    assert exc_info.value.status_code == 400
 def test_forward_marks_unsupported_output_count_as_client_error():
    from vllm_omni.errors import OmniClientError
    pipeline = _make_pipeline()
    request = _make_request_batch(
        {"prompt": "motion", "modalities": ["video"]},
        height=192,
        width=320,
        num_frames=9,
        num_outputs_per_prompt=2,
    )
    with pytest.raises(OmniClientError, match="one output per prompt") as exc_info:
        pipeline.forward(request)
    assert exc_info.value.status_code == 400
 def test_load_weights_rejects_external_weight_stream():
    pipeline = _make_pipeline()
@@ -1,350 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import pytest
 from PIL import Image
 from vllm_omni.diffusion.models.lingbot_video.request_utils import (
    LingBotGenerationMode,
    caption_from_lingbot_prompt,
    normalize_lingbot_num_frames,
    normalize_lingbot_request,
    resolve_lingbot_mode,
    resolve_lingbot_num_frames,
    resolve_lingbot_output_dimensions,
    resolve_lingbot_size,
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _request(prompt, **sampling_overrides):
    return SimpleNamespace(
        prompt=prompt,
        sampling_params=OmniDiffusionSamplingParams(**sampling_overrides),
    )
 def _normalize(prompt, **sampling_overrides):
    return normalize_lingbot_request(
        _request(prompt, **sampling_overrides),
        default_negative_prompt="video default",
        default_image_negative_prompt="image default",
    )
@pytest.mark.parametrize(
    ("prompt", "expected"),
    [
        ({"prompt": "still", "modalities": ["image"]}, LingBotGenerationMode.T2I),
        ({"prompt": "motion", "modalities": ["video"]}, LingBotGenerationMode.T2V),
        (
            {
                "prompt": "motion",
                "modalities": ["video"],
                "multi_modal_data": {"image": Image.new("RGB", (8, 8))},
            },
            LingBotGenerationMode.TI2V,
        ),
    ],
 )
 def test_resolve_lingbot_mode(prompt, expected):
    assert resolve_lingbot_mode(prompt) is expected
@pytest.mark.parametrize(
    ("prompt", "message"),
    [
        (
            {"prompt": "x", "modalities": ["image"], "multi_modal_data": {"image": Image.new("RGB", (8, 8))}},
            "does not accept",
        ),
        (
            {
                "prompt": "x",
                "modalities": ["video"],
                "multi_modal_data": {"image": [Image.new("RGB", (8, 8)), Image.new("RGB", (8, 8))]},
            },
            "exactly one",
        ),
        ({"prompt": "x", "modalities": ["video"], "multi_modal_data": {"video": [1]}}, "reference modalities"),
        ({"prompt": "x", "multi_modal_data": {"image": Image.new("RGB", (8, 8))}}, "explicitly set"),
        ({"prompt": "x", "modalities": ["audio"]}, "output modality"),
    ],
 )
 def test_resolve_lingbot_mode_rejects_invalid_media_combinations(prompt, message):
    with pytest.raises(ValueError, match=message):
        resolve_lingbot_mode(prompt)
 def test_caption_from_lingbot_prompt_preserves_text_and_compact_json():
    assert caption_from_lingbot_prompt("plain text") == "plain text"
    assert caption_from_lingbot_prompt({"caption": "caption only", "duration": 4}) == "caption only"
    assert caption_from_lingbot_prompt({"caption": {"中文": True, "nested": {"value": 2}}}) == (
        '{"中文":true,"nested":{"value":2}}'
    )
    assert (
        caption_from_lingbot_prompt(
            {"scene": "海边", "camera": {"move": "pan"}, "duration": 5, "fps": 24, "size": "320x192"}
        )
        == '{"scene":"海边","camera":{"move":"pan"}}'
    )
@pytest.mark.parametrize("prompt", [{"duration": 4, "fps": 24}, {"caption": ""}, ""])
 def test_caption_from_lingbot_prompt_rejects_empty_caption(prompt):
    with pytest.raises(ValueError, match="empty|must not be empty"):
        caption_from_lingbot_prompt(prompt)
@pytest.mark.parametrize(
    ("num_frames", "expected"),
    [(1, 1), (2, 5), (4, 5), (5, 5), (96, 97), (98, 101), (121, 121)],
 )
 def test_normalize_lingbot_num_frames_rounds_up_to_4n_plus_1(num_frames, expected):
    assert normalize_lingbot_num_frames(num_frames) == expected
@pytest.mark.parametrize("num_frames", [0, -1, 1.5, "invalid"])
 def test_normalize_lingbot_num_frames_rejects_invalid_values(num_frames):
    with pytest.raises(ValueError, match="num_frames"):
        normalize_lingbot_num_frames(num_frames)
@pytest.mark.parametrize(
    ("duration", "fps", "expected"),
    [(4, 24, 97), (5, 24, 121), (8, 24, 193), (4, 12, 49), (4, 30, 121)],
 )
 def test_resolve_lingbot_num_frames_rounds_up(duration, fps, expected):
    assert resolve_lingbot_num_frames(duration, fps) == expected
@pytest.mark.parametrize(("duration", "fps"), [(0, 24), (-1, 24), (0.01, 24), (1, 0), (1, -1)])
 def test_resolve_lingbot_num_frames_rejects_non_positive_or_empty_results(duration, fps):
    with pytest.raises(ValueError, match="duration|fps|at least one"):
        resolve_lingbot_num_frames(duration, fps)
@pytest.mark.parametrize(
    ("resolution", "ratio", "expected"),
    [
        ("192p", "1:1", (192, 192)),
        ("192p", "9:16", (192, 320)),
        ("192p", "16:9", (320, 192)),
        ("192p", "3:4", (192, 256)),
        ("192p", "4:3", (256, 192)),
        ("480p", "1:1", (480, 480)),
        ("480p", "9:16", (480, 832)),
        ("480p", "16:9", (832, 480)),
        ("480p", "3:4", (480, 640)),
        ("480p", "4:3", (640, 480)),
        ("720p", "1:1", (736, 736)),
        ("720p", "9:16", (736, 1280)),
        ("720p", "16:9", (1280, 736)),
        ("720p", "3:4", (736, 960)),
        ("720p", "4:3", (960, 736)),
        ("1080p", "1:1", (1088, 1088)),
        ("1080p", "9:16", (1088, 1920)),
        ("1080p", "16:9", (1920, 1088)),
        ("1080p", "3:4", (1088, 1440)),
        ("1080p", "4:3", (1440, 1088)),
        ("2k", "1:1", (1440, 1440)),
        ("2k", "9:16", (1440, 2560)),
        ("2k", "16:9", (2560, 1440)),
        ("2k", "3:4", (1440, 1920)),
        ("2k", "4:3", (1920, 1440)),
        ("4k", "1:1", (2176, 2176)),
        ("4k", "9:16", (2176, 3840)),
        ("4k", "16:9", (3840, 2176)),
        ("4k", "3:4", (2176, 2880)),
        ("4k", "4:3", (2880, 2176)),
    ],
 )
 def test_resolve_lingbot_size_supports_all_official_presets(resolution, ratio, expected):
    assert resolve_lingbot_size(resolution=resolution, ratio=ratio) == expected
 def test_resolve_lingbot_size_validates_sources_and_alignment():
    assert resolve_lingbot_size(size="320x192") == (192, 320)
    assert resolve_lingbot_size(width=320, height=192) == (192, 320)
    with pytest.raises(ValueError, match="provided together"):
        resolve_lingbot_size(width=320)
    with pytest.raises(ValueError, match="provided together"):
        resolve_lingbot_size(resolution="192p")
    with pytest.raises(ValueError, match="mutually exclusive"):
        resolve_lingbot_size(width=320, height=192, resolution="192p", ratio="9:16")
    with pytest.raises(ValueError, match="multiples of 16"):
        resolve_lingbot_size(width=321, height=192)
 def test_resolve_lingbot_output_dimensions_matches_request_priority():
    assert resolve_lingbot_output_dimensions(
        sampling_width=480,
        sampling_height=480,
        prompt_fields={"resolution": "720p", "ratio": "16:9"},
    ) == (736, 1280)
    assert resolve_lingbot_output_dimensions(
        sampling_width=480,
        sampling_height=480,
        prompt_fields={"width": 320, "height": 192},
    ) == (320, 192)
    assert resolve_lingbot_output_dimensions(
        sampling_width=320,
        sampling_height=192,
    ) == (320, 192)
    assert resolve_lingbot_output_dimensions() == (480, 480)
 def test_normalize_lingbot_request_uses_mode_defaults_and_preserves_empty_negative():
    t2i = _normalize(
        {"prompt": "still", "modalities": ["image"], "negative_prompt": ""},
        width=320,
        height=192,
        num_frames=1,
    )
    assert t2i.mode is LingBotGenerationMode.T2I
    assert t2i.num_frames == 1
    assert t2i.negative_prompt == ""
    t2v = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=9,
    )
    assert t2v.mode is LingBotGenerationMode.T2V
    assert t2v.negative_prompt == "video default"
 def test_normalize_lingbot_request_rejects_explicit_zero_flow_shift():
    with pytest.raises(ValueError, match="shift.*positive"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=1,
            extra_args={"flow_shift": 0},
        )
 def test_normalize_lingbot_request_rejects_explicit_zero_width():
    with pytest.raises(ValueError, match="width"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=0,
            height=192,
            num_frames=1,
        )
 def test_normalize_lingbot_request_rounds_all_video_frame_counts():
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=96,
        fps=24,
        frame_rate=24,
    )
    assert config.num_frames == 97
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=121,
        fps=24,
        frame_rate=24,
    )
    assert config.num_frames == 121
 def test_normalize_lingbot_request_prefers_explicit_frames_over_duration():
    online = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=9,
        fps=24,
        extra_args={"duration": 5},
    )
    assert online.num_frames == 9
    offline = _normalize(
        {
            "prompt": {"caption": "motion", "duration": 5, "num_frames": 9},
            "modalities": ["video"],
        },
        width=320,
        height=192,
        fps=24,
    )
    assert offline.num_frames == 9
 def test_normalize_lingbot_request_uses_duration_then_model_default():
    duration = _normalize(
        {"prompt": {"caption": "motion", "duration": 5}, "modalities": ["video"]},
        width=320,
        height=192,
        fps=24,
    )
    assert duration.num_frames == 121
    default = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
    )
    assert default.num_frames == 81
 def test_normalize_lingbot_request_rejects_contract_conflicts():
    with pytest.raises(ValueError, match="text-to-image.*duration"):
        _normalize(
            {"prompt": {"caption": "still", "duration": 4}, "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=1,
        )
    with pytest.raises(ValueError, match="text-to-image requires"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=5,
        )
    with pytest.raises(ValueError, match="seconds.*duration"):
        _normalize(
            {"prompt": "motion", "modalities": ["video"]},
            width=320,
            height=192,
            num_frames=96,
            extra_args={"seconds": 4, "duration": 4},
        )
 def test_normalize_lingbot_request_model_size_overrides_sampling_dimensions():
    # A prompt-provided size overrides the sampling width/height.
    config = _normalize(
        {"prompt": {"caption": "motion", "size": "320x192"}, "modalities": ["video"]},
        width=480,
        height=480,
        num_frames=81,
    )
    assert (config.height, config.width) == (192, 320)
    # Without a prompt size, the sampling dimensions are used.
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=480,
        height=480,
        num_frames=81,
    )
    assert (config.height, config.width) == (480, 480)
@@ -492,10 +492,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    class FakeBackend:
        supports_prefix_kv_slicing = True
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
@@ -530,48 +526,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
    assert attention.attention.metadata.extra["valid_kv_length"] == 5
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
    )
    class FakeBackend:
        supports_prefix_kv_slicing = False
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return True
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
        def __init__(self):
            super().__init__()
            self.metadata = None
        def forward(self, query, key, value, metadata):
            self.metadata = metadata
            return query
    attention = object.__new__(MiniMaxH3Attention)
    torch.nn.Module.__init__(attention)
    attention.attention = FakeAttention()
    q = torch.randn(8, 2, 4)
    attention._run_packed_attention(
        q,
        q,
        q,
        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
        max_seqlen=5,
        packed_total=8,
    )
    assert attention.attention.metadata.attn_mask is None
    assert attention.attention.metadata.extra["valid_kv_length"] == 5
    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
 def test_packed_attention_keeps_padding_mask_for_other_backends():
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
@@ -580,10 +534,6 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
    class FakeBackend:
        supports_prefix_kv_slicing = False
        @classmethod
        def supports_packed_mask_free(cls) -> bool:
            return False
    class FakeAttention(torch.nn.Module):
        attn_backend = FakeBackend
@@ -19,7 +19,7 @@ MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
 def _assert_joint_output(outputs):
    assert len(outputs) == 1
    frames = np.asarray(outputs[0].images[0])
-    assert frames.shape == (107, 256, 448, 3)
+    assert frames.shape == (39, 256, 448, 3)
    multimodal = outputs[0].multimodal_output
    assert multimodal is not None
    assert np.asarray(multimodal["audio"]).shape[1] == 2
@@ -122,58 +122,6 @@ def test_minimax_h3_t2va_fp8_single_gpu_smoke():
                output_type="np",
                extra_args={
                    "task": "t2va",
                    "duration": 4.0,
                    "aspect_ratio": "16:9",
                    "flow_shift": 12.0,
                    "audio_flow_shift": 3.0,
                },
            ),
            use_tqdm=False,
        )
    finally:
        engine.close()
    _assert_joint_output(outputs)
@pytest.mark.skipif(
    not os.environ.get(MODEL_ENV),
    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
 )
 def test_minimax_h3_t2va_teacache_single_gpu_smoke():
    import torch
    from vllm_omni.diffusion.data import DiffusionParallelConfig
    from vllm_omni.entrypoints.omni import Omni
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
    if torch.accelerator.device_count() < 1:
        pytest.skip("MiniMax H3 TeaCache smoke requires a CUDA device")
    engine = Omni(
        model=os.environ[MODEL_ENV],
        cache_backend="tea_cache",
        cache_config={"rel_l1_thresh": 0.17},
        parallel_config=DiffusionParallelConfig(ulysses_degree=1),
        trust_remote_code=True,
        enable_cpu_offload=True,
        enforce_eager=True,
    )
    try:
        outputs = engine.generate(
            "A quiet cinematic night scene with matching ambient sound.",
            OmniDiffusionSamplingParams(
                height=256,
                width=448,
                num_frames=29,
                fps=24,
                num_inference_steps=2,
                seed=42,
                output_type="np",
                extra_args={
                    "task": "t2va",
                    "duration": 4.0,
                    "aspect_ratio": "16:9",
                    "flow_shift": 12.0,
                    "audio_flow_shift": 3.0,
                },
@@ -69,7 +69,6 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    attention = object.__new__(MiniMaxH3Attention)
    nn.Module.__init__(attention)
    attention.rot_dim = 96
    attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
    attention.rope._forward_method = attention.rope.forward_native
@@ -15,7 +15,6 @@ from vllm_omni.diffusion.data import (
    DiffusionParallelConfig,
    OmniDiffusionConfig,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.model_metadata import (
    HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
    QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
@@ -113,14 +112,6 @@ class TestCreateDefaultDiffusion:
        stages = StageConfigFactory.create_default_diffusion({"model": "x"})
        assert stages[0]["engine_args"]["cache_backend"] == "none"
    def test_vae_decode_precision_roundtrip(self):
        od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16")
        assert od.vae_decode_precision == "bf16"
    def test_invalid_vae_decode_precision_is_rejected(self):
        with pytest.raises(ValueError, match="vae_decode_precision"):
            OmniDiffusionConfig(model="x", vae_decode_precision="int8")
    def test_single_gpu_default_devices(self):
        stages = StageConfigFactory.create_default_diffusion({"model": "x"})
        assert stages[0]["runtime"]["devices"] == "0"
@@ -133,11 +124,6 @@ class TestCreateDefaultDiffusion:
        assert ea["enforce_eager"] is True
        assert ea["lora_path"] == "/tmp/lora"
    def test_diffusion_kv_mode_roundtrip(self):
        od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
        assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
@@ -25,7 +25,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
    CachedRequestData,
    NewRequestData,
    _AdmissionWaitDecision,
 )
 from vllm_omni.diffusion.sched.interface import (
    DiffusionSchedulerOutput as RealDiffusionSchedulerOutput,
@@ -76,34 +75,6 @@ class MockScheduler:
    def num_running_requests(self) -> int:
        return 0
    def get_admission_wait_decision(
        self,
        *,
        now: float,
        dp_concurrent: bool = False,
    ) -> _AdmissionWaitDecision:
        del dp_concurrent
        return _AdmissionWaitDecision(
            should_wait=True,
            deadline=now + 0.5,
            stable_window_s=0.05,
            max_batch=self.max_num_running_reqs,
        )
    def should_end_admission_wait(
        self,
        decision: _AdmissionWaitDecision,
        *,
        now: float,
        stable_since: float,
    ) -> bool:
        waiting = self.num_waiting_requests()
        return (
            waiting >= decision.max_batch
            or (waiting > 0 and now - stable_since >= decision.stable_window_s)
            or (decision.deadline is not None and now >= decision.deadline)
        )
    def schedule(self) -> DiffusionSchedulerOutput:
        if not self._waiting_queue:
            return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -137,17 +108,6 @@ class MockScheduler:
        pass
 def _make_admission_engine(pre_process_func) -> DiffusionEngine:
    engine = object.__new__(DiffusionEngine)
    engine.scheduler = MockScheduler()
    engine.pre_process_func = pre_process_func
    engine._out_streams = {}
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._closed = False
    return engine
 class _BatchCapablePipeline:
    supports_request_batch = True
@@ -565,11 +525,6 @@ class TestRequestBatchAdmission:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0)
    @pytest.mark.parametrize("bad_wait", [float("nan"), float("inf"), float("-inf")])
    def test_config_rejects_nonfinite_request_batch_max_wait_ms(self, bad_wait: float) -> None:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=bad_wait)
    def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None:
        config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5)
@@ -637,7 +592,7 @@ class TestRequestBatchAdmission:
        start = time.monotonic()
        with engine._cv:
-            engine._wait_for_admission_if_needed_locked()
+            engine._wait_for_request_batch_admission_locked()
        waited_s = time.monotonic() - start
        # Stable-window exit (~50ms), not the full 1000ms deadline.
@@ -728,34 +683,6 @@ async def _consume_final_output(generator):
    return final_output
@pytest.mark.cpu
@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"])
 def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None:
    raw_request = OmniDiffusionRequest(
        prompt="raw",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
    )
    prepared_request = OmniDiffusionRequest(
        prompt="prepared",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
        prepared_layout=object(),
    )
    preprocess_calls = []
    def preprocess(request):
        preprocess_calls.append(request)
        return prepared_request
    engine = _make_admission_engine(preprocess)
    getattr(engine, entrypoint)(raw_request)
    assert preprocess_calls == [raw_request]
    assert engine.scheduler._waiting_queue == [prepared_request]
@pytest.mark.cpu
@pytest.mark.asyncio
 async def test_async_add_req_and_stream_response():
@@ -127,7 +127,6 @@ def _make_request(tag: str):
        request_id=tag,
        prompt=f"prompt_{tag}",
        sampling_params=SimpleNamespace(**sampling_params),
        diffusion_kv_requests=None,
    )
@@ -146,7 +145,7 @@ def _make_engine_with_loop(
    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
    sched = RequestScheduler()
-    sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+    sched.initialize(SimpleNamespace(max_num_seqs=1))
    engine.scheduler = sched
    engine.step_execution = False
    engine.supports_request_batch = False
@@ -218,25 +218,6 @@ def _make_compile_runner(
    return runner
@pytest.mark.core_model
@pytest.mark.cpu
 def test_update_states_carries_prepared_layout() -> None:
    runner = _make_runner(cache_backend=None, cache_backend_name=None)
    request = _make_request()
    prepared_layout = object()
    request.prepared_layout = prepared_layout
    scheduler_output = SimpleNamespace(
        finished_req_ids=set(),
        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request)],
        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
    )
    states, new_request_ids = DiffusionModelRunner._update_states(runner, scheduler_output)
    assert new_request_ids == [request.request_id]
    assert states[0].prepared_layout is prepared_layout
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("use_hsdp", [False, True])
@@ -512,35 +493,6 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
    assert runner.pipeline.last_req.num_reqs == 1
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
    ("quality", "expected_enabled"),
    [(None, False), ("lossless", False), ("high", True)],
 )
 def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled):
    runner = _make_runner(cache_backend=None, cache_backend_name="none")
    req = _make_request()
    req.sampling_params.quality = quality
    gate_calls = []
    @contextmanager
    def _record_gate(pipeline, enabled):
        gate_calls.append((pipeline, enabled, "enter"))
        yield
        gate_calls.append((pipeline, enabled, "exit"))
    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
    monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate)
    DiffusionModelRunner.execute_model(runner, req)
    assert gate_calls == [
        (runner.pipeline, expected_enabled, "enter"),
        (runner.pipeline, expected_enabled, "exit"),
    ]
 class _BatchPipeline:
    """Pipeline returning a configurable list of outputs from forward()."""
@@ -608,18 +560,6 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
    sched = _make_scheduler_output(num_reqs=2)
    sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
    sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
    with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -769,19 +709,6 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.core_model
@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    memory_pool_active = False
    optimization_pool_states = []
    @contextmanager
    def _memory_pool_context(tag):
        nonlocal memory_pool_active
        assert tag == "weights"
        memory_pool_active = True
        try:
            yield
        finally:
            memory_pool_active = False
    class _DummyLoader:
        def __init__(self, load_config, od_config=None):
            del load_config, od_config
@@ -829,22 +756,16 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
    monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
    monkeypatch.setattr(
        model_runner_module,
        "optimize_pipeline_vaes",
        lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
    )
    monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
    monkeypatch.setattr(
        model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
    )
-    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
+    DiffusionModelRunner.load_model(runner)
    assert runner.cache_backend is None
    assert runner.od_config.cache_backend is None
    assert dummy_cache_backend.enabled is False
    assert optimization_pool_states == [True]
@pytest.mark.core_model
@@ -12,8 +12,6 @@ from pytest_mock import MockerFixture
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
    BaseScheduler,
@@ -244,86 +242,7 @@ class TestGetRequestBatchSamplingParamsKey:
 class TestRequestScheduler:
    def setup_method(self) -> None:
        self.scheduler: RequestScheduler = RequestScheduler()
-        self.scheduler.initialize(SimpleNamespace(request_batch_max_wait_ms=0.0))
+        self.scheduler.initialize(SimpleNamespace())
    def test_admission_wait_disabled_with_zero_max_wait(self) -> None:
        self.scheduler.initialize(SimpleNamespace(request_batch_max_wait_ms=0.0))
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        assert decision.should_wait is False
    @pytest.mark.parametrize(
        ("dp_concurrent", "expected_stable_window_s"),
        [(False, 0.05), (True, 0.3)],
    )
    def test_admission_wait_decision_encodes_coalescing_policy(
        self,
        dp_concurrent: bool,
        expected_stable_window_s: float,
    ) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=4,
                request_batch_max_wait_ms=1000.0,
            )
        )
        decision = self.scheduler.get_admission_wait_decision(
            now=10.0,
            dp_concurrent=dp_concurrent,
        )
        assert decision.should_wait is True
        assert decision.deadline == 11.0
        assert decision.stable_window_s == expected_stable_window_s
        assert decision.max_batch == 4
    def test_admission_wait_disabled_while_wave_is_running(self) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=1,
                request_batch_max_wait_ms=1000.0,
            )
        )
        self.scheduler.add_request(_make_request("running"))
        self.scheduler.schedule()
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        assert decision.should_wait is False
    def test_admission_wait_end_conditions(self) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=2,
                request_batch_max_wait_ms=1000.0,
            )
        )
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        self.scheduler.add_request(_make_request("a"))
        assert not self.scheduler.should_end_admission_wait(
            decision,
            now=10.01,
            stable_since=10.0,
        )
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=10.05,
            stable_since=10.0,
        )
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=11.0,
            stable_since=11.0,
        )
        self.scheduler.add_request(_make_request("b"))
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=10.01,
            stable_since=10.01,
        )
    def test_single_request_success_lifecycle(self) -> None:
        req_id = self.scheduler.add_request(_make_request("a"))
@@ -340,90 +259,6 @@ class TestRequestScheduler:
        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
        assert self.scheduler.has_requests() is False
    def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        prepared_layout = object()
        request.prepared_layout = prepared_layout
        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
        kv_request = DiffusionKVRequest(
            "diffusion-kv/diffusion-kv/0",
            sequence_id=0,
            prefix_len=4,
            target_len=8,
            seq_len=16,
            kv_contexts=(text_context,),
        )
        request.diffusion_kv_requests = (kv_request,)
        self.scheduler.add_request(request)
        state = self.scheduler.get_request_state(request.request_id)
        scheduler_output = self.scheduler.schedule()
        assert state is not None
        assert state.diffusion_kv_requests == (kv_request,)
        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
        assert request.diffusion_kv_requests is None
        assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
    def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
            self.scheduler.add_request(request)
    @pytest.mark.parametrize(
        ("owner_name", "field_name"),
        [
            ("request", "past_key_values"),
            ("sampling_params", "past_key_values"),
            ("sampling_params", "cfg_text_past_key_values"),
            ("sampling_params", "cfg_img_past_key_values"),
            ("sampling_params", "cfg_branch_past_key_values"),
        ],
    )
    def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        owner = request if owner_name == "request" else request.sampling_params
        setattr(owner, field_name, object())
        request.diffusion_kv_requests = (
            DiffusionKVRequest(
                "diffusion-kv/diffusion-kv/0",
                sequence_id=0,
                prefix_len=4,
                target_len=8,
                seq_len=16,
            ),
        )
        with pytest.raises(ValueError, match=rf"{owner_name}\.{field_name}"):
            self.scheduler.add_request(request)
    def test_dense_scheduler_rejects_scheduler_only_kv_request(self) -> None:
        request = _make_request("dense-kv-request")
        request.diffusion_kv_requests = (
            DiffusionKVRequest(
                "dense-kv-request/diffusion-kv/0",
                sequence_id=0,
                prefix_len=4,
                target_len=8,
                seq_len=16,
            ),
        )
        with pytest.raises(ValueError, match="dense_legacy request unexpectedly contains"):
            self.scheduler.add_request(request)
    def test_dense_scheduler_accepts_legacy_dense_kv(self) -> None:
        request = _make_request("dense-kv")
        request.sampling_params.past_key_values = object()
        request_id = self.scheduler.add_request(request)
        assert request_id == request.request_id
    def test_error_output_marks_finished_error(self) -> None:
        req_id = self.scheduler.add_request(_make_request("err"))
@@ -727,20 +562,12 @@ class TestDiffusionEngine:
        engine.abort_queue = queue.Queue()
        request = _make_request("engine")
        prepared_layout = object()
        def preprocess(req):
            req.prepared_layout = prepared_layout
            return req
        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
        runner_output = _make_request_output("engine")
        engine.execute_fn = mocker.Mock(return_value=runner_output)
        output = engine.add_req_and_wait_for_response(request)
        assert output is runner_output.result
        engine.pre_process_func.assert_called_once_with(request)
        engine.execute_fn.assert_called_once()
    def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None:
@@ -804,29 +631,17 @@ class TestDiffusionEngine:
    async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine._check_and_start_background_loop = mocker.AsyncMock()
-        request = _make_request("req-abort")
+        engine.pre_process_func = None
        prepared_layout = object()
        def preprocess(req):
            req.prepared_layout = prepared_layout
            return req
        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
        async def _stream(_request):
            yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
-        engine._add_prepared_request = mocker.Mock(return_value=request.request_id)
+        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
        engine.get_output_stream = mocker.Mock(return_value=_stream(None))
        with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            async for _ in engine.step_streaming(request):
+            async for _ in engine.step_streaming(_make_request("req-abort")):
                pass
        engine.pre_process_func.assert_called_once_with(request)
        engine._add_prepared_request.assert_called_once_with(request)
        engine.get_output_stream.assert_called_once_with(request.request_id)
    def test_abort_queue_marks_request_finished_aborted(self) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine._rpc_lock = threading.RLock()
@@ -959,32 +774,12 @@ class TestDiffusionEngine:
        with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
            engine._dummy_run()
    def test_dummy_run_delegates_preprocessing_to_sync_admission(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine.od_config = SimpleNamespace(model_class_name="mock_model", diffusion_load_format="default")
        engine.pre_process_func = mocker.Mock()
        engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(output=None))
        engine._dummy_run()
        engine.pre_process_func.assert_not_called()
        admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
        assert admitted_request.request_id == "dummy_req_id"
 class TestStepScheduler:
    def setup_method(self) -> None:
        self.scheduler: StepScheduler = StepScheduler()
        self.scheduler.initialize(SimpleNamespace())
    def test_admission_wait_is_not_supported(self) -> None:
        decision = self.scheduler.get_admission_wait_decision(
            now=10.0,
            dp_concurrent=True,
        )
        assert decision.should_wait is False
    def test_single_request_step_lifecycle(self) -> None:
        request = _make_step_request("step", num_inference_steps=3)
        req_id = self.scheduler.add_request(request)
@@ -102,7 +102,6 @@ class TestDiffusionWorkerSleep:
    def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 1 (offload weights only)."""
        clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches")
        mock_allocator_class = patch_cumem_allocator(mocker)
        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
@@ -129,7 +128,6 @@ class TestDiffusionWorkerSleep:
        # Verify sleep was called with correct tags
        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
        clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
        assert bool(result) is True
        # Verify buffers were NOT saved (level 1 doesn't save buffers)
        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
@@ -48,7 +48,6 @@ def _mock_request(tag: str):
        request_id=tag,
        prompt=f"prompt_{tag}",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        diffusion_kv_requests=None,
    )
@@ -1,300 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Correctness contracts for request-gated VAE decoder fast paths."""
 from __future__ import annotations
 import pytest
 import torch
 import torch.nn.functional as F
 from torch import nn
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.vae_optimizations.gate import (
    VaeFastPathGate,
    register_vae_fast_path_gate,
    use_pipeline_vae_fast_path,
    use_vae_fast_path,
 )
 pytestmark = [pytest.mark.diffusion]
 class _PipelineWithVae(nn.Module):
    def __init__(self, vae: nn.Module) -> None:
        super().__init__()
        self.vae = vae
 class _DistributedExecutorConfig:
    def __init__(self, parallel_mode: str) -> None:
        self.parallel_mode = parallel_mode
@pytest.mark.core_model
@pytest.mark.cpu
 def test_vae_fast_path_gate_restores_nested_state() -> None:
    vae = nn.Identity()
    gate = VaeFastPathGate()
    register_vae_fast_path_gate(vae, gate)
    with use_vae_fast_path(vae, True):
        assert gate.enabled
        with use_vae_fast_path(vae, False):
            assert not gate.enabled
        assert gate.enabled
    assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
 def test_pipeline_vae_fast_path_controls_discovered_vae() -> None:
    vae = nn.Identity()
    pipeline = _PipelineWithVae(vae)
    gate = VaeFastPathGate()
    register_vae_fast_path_gate(vae, gate)
    with use_pipeline_vae_fast_path(pipeline, True):
        assert gate.enabled
    assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
 def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None:
    from vllm_omni.diffusion.vae_optimizations import image
    gate = VaeFastPathGate()
    gate.enabled = True
    norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate)
    inputs = torch.randn(1, 4, 3, 3)
    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
    monkeypatch.setattr(
        image,
        "group_norm_silu_4d",
        lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"),
    )
    actual = norm(inputs)
    expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps))
    torch.testing.assert_close(actual, expected)
 def _small_image_vae() -> nn.Module:
    from diffusers.models.autoencoders import AutoencoderKL
    return AutoencoderKL(
        in_channels=3,
        out_channels=3,
        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"),
        block_out_channels=(32, 32),
        layers_per_block=1,
        latent_channels=4,
        norm_num_groups=8,
        sample_size=16,
    )
 def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module:
    if vae_cls is None:
        from diffusers.models.autoencoders import AutoencoderKLWan
        vae_cls = AutoencoderKLWan
    return vae_cls(
        base_dim=8,
        decoder_base_dim=8,
        z_dim=4,
        dim_mult=[1, 2],
        num_res_blocks=1,
        attn_scales=[],
        temperal_downsample=[False],
        latents_mean=[0.0] * 4,
        latents_std=[1.0] * 4,
        scale_factor_temporal=1,
        scale_factor_spatial=2,
    )
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
 def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None:
    from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
    from vllm_omni.diffusion.vae_optimizations.image import (
        FusedGroupNormSiLU,
        FusedUpsample2xConv2d,
        maybe_optimize_image_vae,
    )
    torch.manual_seed(0)
    vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval()
    latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16)
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents).sample
    maybe_optimize_image_vae(vae)
    maybe_optimize_image_vae(vae)
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules())
    assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules())
    assert torch.equal(vae.decode(latents).sample, reference)
    with use_vae_fast_path(vae, True):
        fast = vae.decode(latents).sample
    torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0)
    assert any(
        buffer is not None
        for module in vae.modules()
        for name, buffer in module._buffers.items()
        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
    )
    clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae))
    assert all(
        buffer is None
        for module in vae.modules()
        for name, buffer in module._buffers.items()
        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
    )
    with use_vae_fast_path(vae, True):
        torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0)
    assert torch.equal(vae.decode(latents).sample, reference)
    assert vae.decoder.conv_in.weight.is_contiguous()
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None:
    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d
    torch.manual_seed(0)
    x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last)
    weight = torch.randn(64, device="cuda", dtype=dtype)
    bias = torch.randn(64, device="cuda", dtype=dtype)
    actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5)
    expected = F.silu(F.group_norm(x, 8, weight, bias))
    assert actual is not None
    atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2)
    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
    "activation_dtype,affine_dtype,atol,rtol",
    [
        (torch.float32, torch.float32, 1e-5, 1e-5),
        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),
    ],
 )
 def test_wan_rmsnorm_silu_matches_reference(
    activation_dtype: torch.dtype,
    affine_dtype: torch.dtype,
    atol: float,
    rtol: float,
 ) -> None:
    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu
    torch.manual_seed(0)
    x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous(
        memory_format=torch.channels_last_3d
    )
    gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype)
    bias = torch.randn_like(gamma)
    normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype)
    expected = F.silu(normalized * 96**0.5 * gamma + bias)
    actual = wan_rmsnorm_silu(x, gamma, bias)
    assert actual is not None
    assert actual.stride() == x.stride()
    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
 def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None:
    from vllm_omni.diffusion.vae_optimizations.wan import (
        FusedWanRMSNormSiLU,
        maybe_optimize_wan_vae,
    )
    torch.manual_seed(0)
    vae = _small_wan_vae().cuda().eval()
    latents = torch.randn(1, 4, 1, 8, 8, device="cuda")
    parameter_names = {name for name, _ in vae.named_parameters()}
    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
    reference = vae.decode(latents).sample
    maybe_optimize_wan_vae(vae)
    maybe_optimize_wan_vae(vae)
    assert {name for name, _ in vae.named_parameters()} == parameter_names
    vae.load_state_dict(state_dict, strict=True)
    assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules())
    assert torch.equal(vae.decode(latents).sample, reference)
    with use_vae_fast_path(vae, True):
        fast = vae.decode(latents).sample
    torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4)
    assert torch.equal(vae.decode(latents).sample, reference)
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"])
 def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None:
    import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations
    monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True)
    vae = _small_wan_vae()
    vae.distributed_executor = _DistributedExecutorConfig(parallel_mode)
    wan_optimizations.maybe_optimize_wan_vae(vae)
    assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules())
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
    ("precision", "decode_dtype"),
    [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
 )
 def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None:
    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
    from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
    torch.manual_seed(0)
    vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval()
    configure_wan_decode_precision(vae, precision)
    assert next(vae.encoder.parameters()).dtype == torch.float32
    assert next(vae.quant_conv.parameters()).dtype == torch.float32
    assert next(vae.decoder.parameters()).dtype == decode_dtype
    assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype
    _, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype))
    assert grid_spec.output_dtype == decode_dtype
    encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16))
    decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32))
    assert encoded.latent_dist.parameters.dtype == torch.float32
    assert decoded.sample.dtype == decode_dtype
@@ -228,6 +228,7 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
@pytest.mark.advanced_model
@pytest.mark.skip(reason="#issue 5475")
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):
@@ -3,7 +3,6 @@
 import importlib.util
 from pathlib import Path
 from types import SimpleNamespace
 import pytest
@@ -44,11 +43,3 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
    with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
        generate_examples.load_model_display_names()
 def test_only_general_serving_examples_get_generated_pages():
    general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
    model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))
    assert generate_examples.is_general_example(general_example)
    assert not generate_examples.is_general_example(model_example)
@@ -1,16 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Online serving smoke coverage for the dense LingBot-Video checkpoint."""
+"""
 Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
 The LingBot dense PR supports text-to-video only, so this module exercises the
 native ``/v1/videos`` path with a small 9-frame request.
 """
 import json
 import os
 import pytest
 from tests.helpers.assertions import assert_video_first_frame_matches
 from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -41,32 +43,11 @@ def _get_diffusion_feature_cases(model: str):
    ]
-@pytest.mark.slow
+@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
    openai_client.send_images_generations_http_request(
        {
            "json": {
                "model": omni_server.model,
                "prompt": PROMPT,
                "negative_prompt": NEGATIVE_PROMPT,
                "size": "320x192",
                "n": 1,
                "response_format": "b64_json",
                "num_inference_steps": 2,
                "guidance_scale": 3.0,
                "flow_shift": 3.0,
                "seed": 42,
            }
        }
    )
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
 def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
    request_config = {
        "model": omni_server.model,
        "form_data": {
@@ -84,18 +65,3 @@ def test_video_generation_modes_001(omni_server: OmniServer, openai_client: Open
        },
    }
    openai_client.send_video_diffusion_request(request_config)
    synthetic_image = generate_synthetic_image(320, 192, force_regenerate=True, seed=42)
    request_config["form_data"]["prompt"] = "the red block moves slowly while the camera remains fixed"
    # Keep TI2V dimensions model-scoped until /v1/videos forwards raw reference images.
    request_config["form_data"].pop("height")
    request_config["form_data"].pop("width")
    request_config["form_data"]["extra_params"] = json.dumps({"size": "320x192"}, separators=(",", ":"))
    request_config["image_reference"] = f"data:image/jpeg;base64,{synthetic_image['base64']}"
    responses = openai_client.send_video_diffusion_request(request_config)
    assert responses[0].videos
    assert_video_first_frame_matches(
        responses[0].videos[0],
        synthetic_image["np_array"],
        max_mean_absolute_error=0.25,
    )
@@ -4,10 +4,9 @@
 """
 L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
-This file remains dense-only. Basic T2I/T2V/TI2V serving for both dense and MoE
+This file remains dense-only. Basic single-GPU MoE serving is covered by
-checkpoints is covered by ``test_lingbot_video.py`` and
+``test_lingbot_video_moe.py``; T2I, I2V, TI2V, and multi-GPU feature rows
-``test_lingbot_video_moe.py``. Distributed feature rows belong in follow-up
+belong in follow-up PRs.
 PRs.
 """
 import json
@@ -3,14 +3,11 @@
 """Single-GPU serving smoke for ``robbyant/lingbot-video-moe-30b-a3b``."""
 import json
 import os
 import pytest
 from tests.helpers.assertions import assert_video_first_frame_matches
 from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -41,32 +38,10 @@ def _get_server_cases(model: str):
    ]
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
 def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
    openai_client.send_images_generations_http_request(
        {
            "json": {
                "model": omni_server.model,
                "prompt": PROMPT,
                "negative_prompt": NEGATIVE_PROMPT,
                "size": "320x192",
                "n": 1,
                "response_format": "b64_json",
                "num_inference_steps": 2,
                "guidance_scale": 3.0,
                "flow_shift": 3.0,
                "seed": 42,
            }
        }
    )
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
-def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_text_to_video_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
    request_config = {
        "model": omni_server.model,
        "form_data": {
@@ -82,18 +57,3 @@ def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: Open
        },
    }
    openai_client.send_video_diffusion_request(request_config)
    synthetic_image = generate_synthetic_image(320, 192, force_regenerate=True, seed=42)
    request_config["form_data"]["prompt"] = "the red block moves slowly while the camera remains fixed"
    # Keep TI2V dimensions model-scoped until /v1/videos forwards raw reference images.
    request_config["form_data"].pop("height")
    request_config["form_data"].pop("width")
    request_config["form_data"]["extra_params"] = json.dumps({"size": "320x192"}, separators=(",", ":"))
    request_config["image_reference"] = f"data:image/jpeg;base64,{synthetic_image['base64']}"
    responses = openai_client.send_video_diffusion_request(request_config)
    assert responses[0].videos
    assert_video_first_frame_matches(
        responses[0].videos[0],
        synthetic_image["np_array"],
        max_mean_absolute_error=0.25,
    )
@@ -223,17 +223,6 @@ def test_client(mock_async_diffusion):
    return TestClient(app)
@pytest.fixture
 def lingbot_test_client(test_client):
    test_client.app.state.stage_configs = [
        SimpleNamespace(
            stage_type="diffusion",
            engine_args={"model_class_name": "LingBotVideoPipeline"},
        )
    ]
    return test_client
@pytest.fixture
 def async_omni_test_client():
    """Create test client with mocked AsyncOmni engine."""
@@ -3602,63 +3602,6 @@ class TestGLMTTSServing:
        load_tokenizer.assert_called_once()
@pytest.fixture
 def ming_flash_omni_tts_server(mocker: MockerFixture):
    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
    mock_engine_client = mocker.MagicMock()
    mock_engine_client.errored = False
    mock_engine_client.model_config = mocker.MagicMock(
        model="inclusionAI/Ming-omni-tts-0.5B",
    )
    mock_engine_client.default_sampling_params_list = [
        SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
    ]
    mock_engine_client.tts_batch_max_items = 32
    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
    mock_engine_client.stage_configs = [
        SimpleNamespace(
            engine_args=SimpleNamespace(model_stage="ming_tts"),
            tts_args={},
        )
    ]
    mock_models = mocker.MagicMock()
    mock_models.is_base_model.return_value = True
    return OmniOpenAIServingSpeech(
        engine_client=mock_engine_client,
        models=mock_models,
        request_logger=mocker.MagicMock(),
    )
 class TestMingFlashOmniTTSServing:
    def test_validate_ming_flash_omni_tts_rejects_ref_text(self, ming_flash_omni_tts_server):
        request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript")
        error = ming_flash_omni_tts_server._validate_tts_request(request)
        assert error is not None
        assert "ref_text" in error
    def test_validate_ming_flash_omni_tts_rejects_ref_audio(self, ming_flash_omni_tts_server):
        request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
        error = ming_flash_omni_tts_server._validate_tts_request(request)
        assert error is not None
        assert "ref_audio" in error
    def test_ming_flash_omni_tts_adapter_builds_prompt(self, ming_flash_omni_tts_server, mocker: MockerFixture):
        ming_flash_omni_tts_server._build_ming_flash_omni_prompt = mocker.MagicMock(
            return_value={
                "prompt_token_ids": [1, 2, 3],
                "additional_information": {"voice": ["test"]},
            }
        )
        request = OpenAICreateSpeechRequest(input="Hello", voice="test")
        asyncio.run(ming_flash_omni_tts_server._prepare_speech_generation(request))
        ming_flash_omni_tts_server._build_ming_flash_omni_prompt.assert_called_once()
 class TestTTSAsyncOffloading:
    """Tests for event-loop-safe offloading of blocking TTS operations."""
@@ -63,6 +63,12 @@ def test_resolve_unknown_returns_none():
    assert resolve_adapter(None) is None
 def test_ming_flash_omni_not_migrated():
    """ming_flash_omni is intentionally excluded from the adapter migration in
    this PR; it stays on the legacy inline dispatch in serving_speech.py."""
    assert resolve_adapter("ming_flash_omni_tts") is None
 def test_voxcpm2_resolves():
    """VoxCPM2 (the served ``latent_generator`` model) resolves cleanly.
@@ -28,7 +28,6 @@ from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
    VideoGenerationRequest,
    VideoGenerationStatus,
    VideoParams,
    VideoResponse,
 )
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
@@ -441,75 +440,6 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
    assert input_image.size == (96, 64)
 def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture):
    image_bytes = _make_test_image_bytes((48, 48))
    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
        data={
            "prompt": "A bear playing with yarn.",
            "extra_params": json.dumps({"width": 96, "height": 64}),
        },
        files={"input_reference": ("input.png", image_bytes, "image/png")},
    )
    assert response.status_code == 200
    video_id = response.json()["id"]
    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
    engine = test_client.app.state.openai_serving_video._engine_client
    input_image = engine.captured_prompt["multi_modal_data"]["image"]
    assert isinstance(input_image, Image.Image)
    assert input_image.size == (48, 48)
    sampling_params = engine.captured_sampling_params_list[0]
    assert sampling_params.extra_args["width"] == 96
    assert sampling_params.extra_args["height"] == 64
@pytest.mark.parametrize(
    ("generation_request", "expected_num_frames", "expected_duration"),
    [
        (
            VideoGenerationRequest(prompt="top-level frames", seconds="5", num_frames=9),
            9,
            5.0,
        ),
        (
            VideoGenerationRequest(prompt="nested frames", video_params=VideoParams(num_frames=9)),
            9,
            None,
        ),
        (
            VideoGenerationRequest(prompt="seconds only", seconds="5"),
            120,
            5.0,
        ),
    ],
 )
 def test_video_generation_bridges_request_fields(generation_request, expected_num_frames, expected_duration):
    engine = FakeAsyncOmni()
    handler = OmniOpenAIServingVideo.for_diffusion(
        diffusion_engine=engine,
        model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
    )
    asyncio.run(handler._run_and_extract(generation_request, "field-bridge"))
    sampling = engine.captured_sampling_params_list[0]
    # Top-level ``seconds`` bridges into extra_args["duration"]; num_frames is
    # passed through (or derived as seconds x fps when omitted). No private
    # provenance channel is injected.
    assert "_vllm_request_context" not in sampling.extra_args
    assert sampling.num_frames == expected_num_frames
    if expected_duration is None:
        assert "duration" not in sampling.extra_args
    else:
        assert sampling.extra_args["duration"] == expected_duration
 def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
@@ -3,7 +3,6 @@ import re
 from types import SimpleNamespace
 import pytest
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from tests.helpers.mark import hardware_test
@@ -22,13 +21,10 @@ async def _noop(*args, **kw):
    pass
-def get_fake_add_request(submitted_request_ids, submitted_lora_requests=None):
+def get_fake_add_request(submitted_request_ids):
    async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
        lora_request = kwargs.get("lora_request")
        del prompt, sampling_params_list, final_stage_id, kwargs
        submitted_request_ids.append(request_id)
        if submitted_lora_requests is not None:
            submitted_lora_requests.append(lora_request)
    return fake_add_request_async
@@ -96,37 +92,6 @@ def test_generate_submits_randomized_id_to_engine():
    asyncio.run(run())
@pytest.mark.cpu
 def test_generate_forwards_lora_request_to_engine():
    """Ensure the lora_request passed to generate() reaches add_request_async.
    Regression test for https://github.com/vllm-project/vllm-omni/issues/5369:
    AsyncOmni.generate() accepted lora_request but silently dropped it when
    submitting the request, so generation always used the base model.
    """
    async def run():
        submitted_ids = []
        submitted_loras = []
        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras))
        lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
        async for _ in omni.generate(
            prompt={"prompt": "test"},
            request_id="lora-req",
            sampling_params_list=[SimpleNamespace()],
            output_modalities=["text"],
            lora_request=lora,
        ):
            pass
        assert len(submitted_ids) == 1
        assert len(submitted_loras) == 1
        assert submitted_loras[0] is lora
    asyncio.run(run())
@pytest.mark.cpu
@pytest.mark.parametrize(
    "req_ids,cancel_prefix,expected_cancel_count",
@@ -389,24 +389,6 @@ def test_serve_cli_accepts_request_batch_max_wait_ms():
    assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0
@pytest.mark.parametrize("bad_wait", ["nan", "inf", "-inf", "-1"])
 def test_serve_cli_rejects_invalid_request_batch_max_wait_ms(bad_wait: str):
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)
    with pytest.raises(SystemExit):
        parser.parse_args(
            [
                "serve",
                "Qwen/Qwen-Image",
                "--omni",
                "--request-batch-max-wait-ms",
                bad_wait,
            ]
        )
 def test_serve_cli_accepts_additional_config():
    """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config."""
    parser = TrackingArgumentParser()
@@ -76,17 +76,6 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
    assert args.cfg_parallel_size == 4
 def test_serve_parser_accepts_wan_decode_precision() -> None:
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    OmniServeCommand().subparser_init(subparsers)
    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
    assert args.vae_decode_precision == "bf16"
    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
 def _make_headless_args(**kwargs) -> TrackingNamespace:
    defaults = {
        "model": "fake-model",
@@ -12,7 +12,6 @@ from typing import TYPE_CHECKING, Any, Literal
 if TYPE_CHECKING:
    from tests.helpers.runtime import DiffusionResponse
 import av
 import numpy as np
 import soundfile as sf
 from PIL import Image
@@ -20,7 +19,6 @@ from PIL import Image
 from tests.helpers.media import (
    convert_audio_bytes_to_text,
    cosine_similarity_text,
    decode_b64_image,
    preprocess_text,
 )
@@ -114,34 +112,6 @@ def assert_image_diffusion_response(
                    )
 def assert_images_generations_response(
    response: dict[str, Any],
    request_config: dict[str, Any],
    run_level: str | None = None,
 ) -> None:
    """Validate a successful ``/v1/images/generations`` JSON response."""
    del run_level
    request_body = request_config.get("json") or {}
    data = response.get("data")
    assert isinstance(data, list), "Image generation response is missing data[]"
    expected_count = int(request_body.get("n", 1))
    assert len(data) == expected_count, f"Expected {expected_count} images, got {len(data)}"
    width = height = None
    size = request_body.get("size")
    if isinstance(size, str) and "x" in size:
        width_value, height_value = size.lower().split("x", 1)
        width, height = int(width_value), int(height_value)
    for item in data:
        assert isinstance(item, dict), "Image generation data entries must be objects"
        b64_json = item.get("b64_json")
        assert isinstance(b64_json, str) and b64_json, "Image generation response is missing b64_json"
        image = decode_b64_image(b64_json)
        assert_image_valid(image, width=width, height=height)
 def assert_video_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
@@ -188,29 +158,6 @@ def assert_video_diffusion_response(
        )
 def assert_video_first_frame_matches(
    video: Path | bytes | BytesIO,
    expected: np.ndarray,
    *,
    max_mean_absolute_error: float,
 ) -> None:
    """Assert that an encoded video's first frame matches a reference image."""
    if isinstance(video, Path):
        source: str | BytesIO = str(video)
    else:
        video_bytes = video if isinstance(video, bytes) else video.getvalue()
        source = BytesIO(video_bytes)
    with av.open(source) as container:
        first_frame = next(container.decode(video=0)).to_ndarray(format="rgb24")
    assert first_frame.shape == expected.shape
    mean_absolute_error = float(np.abs(first_frame.astype(np.float32) - expected).mean() / 255.0)
    assert mean_absolute_error < max_mean_absolute_error, (
        f"Expected first-frame MAE < {max_mean_absolute_error}, got {mean_absolute_error:.6f}."
    )
 def assert_audio_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
@@ -41,7 +41,6 @@ from tests.helpers.assertions import (
    assert_audio_speech_response,
    assert_diffusion_response,
    assert_http_error,
    assert_images_generations_response,
    assert_omni_response,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
@@ -1505,11 +1504,6 @@ class OpenAIClientHandler:
            err_code=cfg.get("err_code"),
            err_message=cfg.get("err_message"),
        )
        if cfg.get("err_code") is None:
            assert resp.success, resp.error_message
            payload = resp.json_body
            assert isinstance(payload, dict)
            assert_images_generations_response(payload, cfg, run_level=self.run_level)
        return [resp]
    def send_images_edits_http_request(
@@ -112,90 +112,6 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
    assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
 def test_lingbot_extra_registry_declares_request_params() -> None:
    assert get_extra_body_params("LingBotVideoPipeline") == frozenset(
        {
            "batch_cfg",
            "duration",
            "flow_shift",
            "negative_prompt",
            "null_cond_clone_zero",
            "offload_vae_during_denoise",
            "output_type",
            "refiner_sigma_tail_steps",
            "resolution",
            "ratio",
            "shift",
            "t_thresh",
        }
    )
    assert get_extra_output_params("LingBotVideoPipeline") == frozenset()
    assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
 def test_lingbot_text_to_image_prompt_builder_preserves_empty_negative_prompt() -> None:
    assert build_text_to_image_prompt(
        "LingBotVideoPipeline",
        prompt="a red fox in fresh snow",
        negative_prompt="",
        height=192,
        width=320,
    ) == {
        "prompt": "a red fox in fresh snow",
        "modalities": ["image"],
        "negative_prompt": "",
    }
@pytest.mark.core_model
@pytest.mark.cpu
 def test_lingbot_image_to_video_prompt_builder() -> None:
    image = Image.new("RGB", (320, 192), "red")
    result = build_image_to_video_prompt(
        "LingBotVideoPipeline",
        prompt="the fox looks toward the camera",
        negative_prompt=None,
        media_inputs={"image": image},
        height=192,
        width=320,
        num_frames=9,
    )
    assert result == {
        "prompt": "the fox looks toward the camera",
        "modalities": ["video"],
        "multi_modal_data": {"image": image},
        "num_frames": 9,
    }
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
    "media_inputs",
    [
        {},
        {"image": "input.png"},
        {"image": [Image.new("RGB", (16, 16))]},
        {"image": Image.new("RGB", (16, 16)), "video": b"video"},
    ],
    ids=["missing", "path", "multiple", "extra-modality"],
 )
 def test_lingbot_image_to_video_prompt_builder_rejects_invalid_media(
    media_inputs: dict[str, object],
 ) -> None:
    with pytest.raises(ValueError, match="exactly one PIL image"):
        build_image_to_video_prompt(
            "LingBotVideoPipeline",
            prompt="move",
            negative_prompt=None,
            media_inputs=media_inputs,
        )
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
@@ -31,8 +31,8 @@ SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
 ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py")
 # Budgets. These may only go DOWN. See the module docstring.
-MAX_MODEL_TYPE_BRANCHES = 27
+MAX_MODEL_TYPE_BRANCHES = 29
-MAX_LEGACY_DETECTORS = 0
+MAX_LEGACY_DETECTORS = 1
 _GUIDANCE = """
 Move the per-model behaviour into its adapter under
@@ -123,29 +123,18 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
    return found
-def _check(
+def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
    label: str,
    path: Path,
    actual: int,
    budget: int,
    sample: list[str],
    budget_name: str,
    notices: list[str],
 ) -> list[str]:
    if actual > budget:
        listing = "\n".join(f"      {line}" for line in sample)
        return [
            f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
        ]
    if actual < budget:
-        # Progress is never a build failure. Removing branches is the goal, so a
+        return [
-        # PR that does it must not hand-edit this file to stay green -- nor break
+            f"{path}: {actual} {label}, down from a budget of {budget}. "
-        # `main` for everyone else when it merges from a stale base.
+            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
-        notices.append(
+            f"so the ratchet holds the new ground.",
-            f"{path}: {actual} {label}, below the budget of {budget}. "
+        ]
            f"Lower {budget_name} to {actual} in "
            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
        )
    return []
@@ -163,7 +152,6 @@ def main(argv: list[str] | None = None) -> int:
            print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
            return 1
    notices: list[str] = []
    branches = _model_type_branches(ast.parse(serving_path.read_text()))
    detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
@@ -174,7 +162,6 @@ def main(argv: list[str] | None = None) -> int:
        MAX_MODEL_TYPE_BRANCHES,
        [f"line {lineno}: {src}" for lineno, src in branches],
        "MAX_MODEL_TYPE_BRANCHES",
        notices,
    )
    errors += _check(
        "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -183,11 +170,8 @@ def main(argv: list[str] | None = None) -> int:
        MAX_LEGACY_DETECTORS,
        [f"line {lineno}: {name}" for lineno, name in detectors],
        "MAX_LEGACY_DETECTORS",
        notices,
    )
    for notice in notices:
        print(f"check_tts_adapter: {notice}", file=sys.stderr)
    for error in errors:
        print(f"check_tts_adapter: {error}", file=sys.stderr)
    return 1 if errors else 0
@@ -37,7 +37,6 @@ from vllm_omni.config.stage_config import (
    build_stage_runtime_overrides,
    load_deploy_config,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = {
    StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
@@ -550,7 +549,6 @@ class _DiffusionConfigProjection:
    cache_backend: str = "none"
    cache_config: Any = field(default_factory=dict)
    enable_cache_dit_summary: bool = False
    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
    enable_prompt_embed_cache: bool = False
    prompt_embed_cache_size: int = Field(default=32, ge=1)
    enable_session_state_manager: bool = False
@@ -573,7 +571,6 @@ class _DiffusionConfigProjection:
    fa_deterministic: bool = False
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None
    mask_strategy_file_path: str | None = None
    skip_time_steps: int = 15
    VSA_sparsity: float = 0.0
@@ -626,7 +623,6 @@ class _DiffusionConfigProjection:
            build_attention_config,
            parse_kv_cache_skip_selector,
        )
        from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
        from vllm_omni.quantization import build_quant_config
        if self.tf_model_config is None:
@@ -688,7 +684,6 @@ class _DiffusionConfigProjection:
                f"got {type(self.diffusion_attention_config)!r}"
            )
        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
        self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
        self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
@@ -401,7 +401,6 @@ class StageDeployConfig:
    step_execution: bool | None = None
    vae_use_slicing: bool | None = None
    vae_use_tiling: bool | None = None
    vae_decode_precision: str | None = None
    boundary_ratio: float | None = None
    flow_shift: float | None = None
    diffusion_kv_cache_dtype: str | None = None
@@ -20,17 +20,6 @@ class AttentionBackend(ABC):
    # avoid a slower masked-attention plan when tail padding is not semantic.
    supports_prefix_kv_slicing: bool = False
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        """Whether packed attention never reads attn_mask on this platform.
        When True, models that pack a [real, pad] two-document layout and
        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
        constructing the padding mask entirely. Backends whose mask-free
        behavior is platform-dependent must check current_omni_platform.
        """
        return False
    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
    # Platform resolution rejects an explicit selection outside this set, so a
    # hardware-specific backend fails before the model is built.
@@ -132,17 +121,6 @@ class AttentionMetadata:
    #     the packed cu_seqlens tensors.
    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
    #     backend that advertises supports_prefix_kv_slicing.
    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
    #     Requires the [real, pad] two-document packing contract; see
    #     FlashAttentionImpl._forward_varlen_packed_npu.
    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
    #     with this set (>1), q/k/v are divided by the factor before the op,
    #     the kernel scale_value is multiplied by its square, and the output
    #     is scaled back (exact for power-of-two factors). Absent means no
    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
    # Piecewise attention metadata (mixed causal/full masks).
    # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from functools import partial
 import torch
@@ -17,7 +16,6 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
    piecewise_attn,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
@@ -26,15 +24,6 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supports_piecewise_spans: bool = True
    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
        # before ever reading attn_mask. NPU: forward_fa_npu honors the
        # npu_attn_varlen opt-in and rebuilds the padding mask itself if the
        # packed contract fails. XPU reads attn_mask, so models must keep
        # constructing it there.
        return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
    @classmethod
    def supports_attention_mask(cls) -> bool:
        return True
@@ -426,36 +415,7 @@ class FlashAttentionImpl(AttentionImpl):
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
        # model marks extra["npu_attn_varlen"] and carries packed metadata, so
        # the padding document is excluded without reading or materializing
        # the quadratic full_qk mask.
        #   - default: packed varlen via mindiesd attention_forward_varlen;
        #   - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
        #     slicing via mindiesd attention_forward (so the env-selected op
        #     type is honored; the TND varlen op cannot be).
        extra = attn_metadata.extra if attn_metadata else {}
        if extra.get("npu_attn_varlen", False):
            if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
            else:
                out = self._forward_varlen_packed_npu(query, key, value, extra)
            if out is not None:
                return out
        attention_mask = attn_metadata.attn_mask if attn_metadata else None
        if attention_mask is None and extra.get("npu_attn_varlen", False):
            # Models skip mask construction when this opt-in is set
            # (FlashAttentionBackend.supports_packed_mask_free). The packed
            # paths above declined (contract mismatch), so rebuild the padding
            # mask here to keep the masked fallback correct.
            used = extra.get("valid_kv_length")
            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
                raise ValueError(
                    "npu_attn_varlen packed metadata is unusable and no attn_mask "
                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
                    "refusing to run unmasked attention over padding rows."
                )
            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
        # NPU aclnnFlashAttentionScore requires mask shape to be one of:
        # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -473,161 +433,3 @@ class FlashAttentionImpl(AttentionImpl):
            op_type="fused_attn_score",
            layout=layout,
        )
    def _resolve_packed_seq_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        extra: dict,
    ) -> tuple[list[int], list[int]] | None:
        """Resolve packed document boundaries (cumulative end offsets per doc).
        Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
        or None unless the metadata matches the exact contract MiniMax-H3
        produces:
          - q/k/v are [1, T, N, D] (BSND, single packed batch);
          - cu_seqlens describe a [real, pad] two-document packing, with the
            padding document as a strict suffix;
          - max_seqlen_q/k are Python ints equal to the real document length
            (so boundaries are derived without any device sync).
        """
        if self.causal:
            return None
        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
        if any(k not in extra for k in packed_keys):
            return None
        cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
        used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
        if query.shape[0] != 1 or key.shape[0] != 1:
            return None
        if not isinstance(used_q, int) or not isinstance(used_k, int):
            return None
        total_q, total_k = query.shape[1], key.shape[1]
        # .shape is host-side metadata: counting documents never syncs.
        if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
            return None
        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
            return None
        if used_q == total_q and used_k == total_k:
            # No padding document: one full-length document.
            return [total_q], [total_k]
        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
            # [real, pad] packing: the real document must be the longer one
            # (consistent with the max_seqlen naming).
            return [used_q, total_q], [used_k, total_k]
        return None
    def _forward_varlen_packed_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
        Returns None (caller falls back to the mask path) when the packed
        contract does not hold; see _resolve_packed_seq_npu.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        seq_q, seq_k = resolved
        try:
            from mindiesd import attention_forward_varlen
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        q = query.squeeze(0)  # [T, N, D] == TND
        k = key.squeeze(0)
        v = value.squeeze(0)
        # attention_forward_varlen wants cu_seqlens as host lists of cumulative
        # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
        out = attention_forward_varlen(
            q,
            k,
            v,
            [0, *seq_q],
            [0, *seq_k],
            softmax_scale=self.softmax_scale,
        )
        return out.unsqueeze(0)
    def _forward_prefix_kv_slice_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        extra: dict,
    ) -> torch.Tensor | None:
        """Mask-free attention by slicing K/V to the valid prefix (zero-copy
        views), then mindiesd attention_forward without a mask.
        Same numerical contract as the varlen path: the padding document is a
        strict suffix, so dropping it from K/V is identical to masking it out.
        Query keeps full length; outputs on padding rows are never consumed
        downstream. Returns None when the packed contract does not hold.
        Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
        cannot honor the env-selected op type, while attention_forward can.
        The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
        bf16 activations with large outliers overflow 65504 into ±inf and the
        affected rows turn NaN. Models may opt into exact power-of-two input
        pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
        are divided by the factor, the kernel scale_value is multiplied by its
        square, and the output is scaled back. Absent or invalid factor means
        no pre-scaling.
        """
        resolved = self._resolve_packed_seq_npu(query, key, extra)
        if resolved is None:
            return None
        _, seq_k = resolved
        used_k = seq_k[0]  # real document length (first cumulative end)
        try:
            from mindiesd import attention_forward
        except ImportError:
            raise ImportError(
                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
                "Please install MindIE-SD to enable NPU attention support. "
                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
            )
        # mindiesd always takes BSND input regardless of `layout`; the arg
        # selects the op-internal layout and laser only supports BNSD, so do
        # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
        layout = "BNSD"
        key = key[:, :used_k]
        value = value[:, :used_k]
        input_scale = extra.get("laser_input_scale")
        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
        if preserve_input_range:
            query = query / input_scale
            key = key / input_scale
            value = value / input_scale
        scale = self.softmax_scale
        if preserve_input_range:
            scale *= input_scale**2
        def _postprocess(out: torch.Tensor) -> torch.Tensor:
            return out * input_scale if preserve_input_range else out
        return _postprocess(
            attention_forward(
                query,
                key,
                value,
                attn_mask=None,
                opt_mode="manual",
                op_type="fused_attn_score",  # MINDIE_SD_FA_TYPE env overrides this
                layout=layout,
                scale=scale,
            )
        )
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
    other case — warmup denoise steps, exempt layers, a layer that does not declare
    ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection — delegates to FlashAttention,
+    segments too short to pay for block selection or not aligned to the kernel's
-    so a model can select this backend unconditionally. MindIE-SD handles an
+    block size — delegates to FlashAttention, so a model can select this backend
-    irregular video tail internally, retaining it outside the sparse blocks.
+    unconditionally.
    """
    def __init__(
@@ -292,6 +292,25 @@ class RainFusionAttentionImpl(AttentionImpl):
                _MIN_VIDEO_BLOCKS,
            )
            return None
        if video_len % _BLOCK_SIZE != 0:
            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
            # builds its block mask from the two segments pooled apart. The counts
            # only agree when the video segment ends on a block boundary; otherwise
            # a prefix block is forced on out of range, and the block straddling the
            # seam holds real prefix rows that selection may drop. Padding the
            # sequence to realign it is not an option because rf_v2 ignores
            # attn_mask, so pad keys would take a share of every softmax denominator.
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
                "multiple of %d to enable sparsity.",
                tuple(latent_shape),
                video_len,
                _BLOCK_SIZE,
                _BLOCK_SIZE,
            )
            return None
        logger.info_once(
            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
@@ -73,34 +73,6 @@ def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) ->
    )
 def enable_minimax_h3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
    """Enable TeaCache for the calibrated MiniMax-H3 FL2VA partition."""
    partition = getattr(pipeline, "partition", None)
    if partition == "ref2va":
        raise ValueError(
            "TeaCache only supports the MiniMax-H3 FL2VA partition; Ref2VA TeaCache has not been calibrated"
        )
    if partition not in {"fl2va", "combined"}:
        raise ValueError(f"Unsupported MiniMax-H3 partition for TeaCache: {partition!r}")
    teacache_config = TeaCacheConfig(
        transformer_type="MiniMaxH3DiTModel",
        rel_l1_thresh=config.rel_l1_thresh,
        coefficients=config.coefficients,
    )
    apply_teacache_hook(pipeline.transformer, teacache_config)
    if partition == "combined":
        logger.warning(
            "TeaCache is enabled only for MiniMax-H3 FL2VA; "
            "Ref2VA requests on this combined pipeline will run without TeaCache"
        )
    logger.info(
        f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
        "transformer_class=MiniMaxH3DiTModel, partition=FL2VA"
    )
 def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
    """
    Enable TeaCache for Flux2 Klein model.
@@ -124,7 +96,6 @@ CUSTOM_TEACACHE_ENABLERS = {
    "BagelPipeline": enable_bagel_teacache,
    "Flux2KleinPipeline": enable_flux2_klein_teacache,
    "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
    "MiniMaxH3Pipeline": enable_minimax_h3_teacache,
    "SenseNovaU1Pipeline": enable_sensenova_u1_teacache,
 }
@@ -174,6 +145,7 @@ class TeaCacheBackend(CacheBackend):
            # Create TeaCacheConfig from DiffusionCacheConfig with transformer_type
            # Access parameters via attribute access: config.rel_l1_thresh
            # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig
            try:
                teacache_config = TeaCacheConfig(
                    transformer_type=transformer_type,
@@ -83,19 +83,6 @@ _MODEL_COEFFICIENTS = {
    ],
    # LongCat Image transformer coefficients
    "LongCatImageTransformer2DModel": [652.5980, -424.1615, 84.5526, -4.5923, 0.1694],
    # MiniMax-H3 FL2VA coefficients.
    "MiniMaxH3DiTModel": [
        2.283704065852778e03,
        -7.775977277886368e02,
        9.408414741359490e01,
        -4.232669906169421e00,
        2.173782527946167e-01,
    ],
 }
 _DEFAULT_REL_L1_THRESH = 0.2
 _MODEL_DEFAULT_REL_L1_THRESH = {
    "MiniMaxH3DiTModel": 0.17,
 }
@@ -110,8 +97,10 @@ class TeaCacheConfig:
    Args:
        rel_l1_thresh: Threshold for accumulated relative L1 distance. When below threshold,
-            cached residual is reused. If None, uses the model-specific default or 0.2
+            cached residual is reused. Values in [0.1, 0.3] work best:
-            when no model-specific default is registered.
+            - 0.2: ~1.5x speedup with minimal quality loss
            - 0.4: ~1.8x speedup with slight quality loss
            - 0.6: ~2.0x speedup with noticeable quality loss
        coefficients: Polynomial coefficients for rescaling L1 distance. If None, uses
            model-specific defaults based on transformer_type.
        transformer_type: Transformer class name (e.g., "QwenImageTransformer2DModel").
@@ -119,18 +108,14 @@ class TeaCacheConfig:
            Defaults to "QwenImageTransformer2DModel".
    """
-    rel_l1_thresh: float | None = None
+    rel_l1_thresh: float = 0.2
    coefficients: list[float] | None = None
    transformer_type: str = "QwenImageTransformer2DModel"
    def __post_init__(self) -> None:
        """Validate and set default coefficients."""
-        threshold = self.rel_l1_thresh
+        if self.rel_l1_thresh <= 0:
-        if threshold is None:
+            raise ValueError(f"rel_l1_thresh must be positive, got {self.rel_l1_thresh}")
            threshold = _MODEL_DEFAULT_REL_L1_THRESH.get(self.transformer_type, _DEFAULT_REL_L1_THRESH)
        if threshold <= 0:
            raise ValueError(f"rel_l1_thresh must be positive, got {threshold}")
        self.rel_l1_thresh = threshold
        if self.coefficients is None:
            # Use model-specific coefficients, explicitly check if the type exists or not
@@ -1249,152 +1249,6 @@ def extract_sensenova_u1_context(
    )
 def extract_minimax_h3_context(
    module: nn.Module,
    **kwargs: Any,
 ) -> CacheContext:
    """Extract cache context for MiniMaxH3DiTModel.
    Mirrors the packed multimodal forward in ``MiniMaxH3DiTModel.forward``:
    preprocessing via ``_embed``, cacheable ``blocks`` loop, and
    ``final_layer`` postprocessing with row selection and update masks.
    """
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        _BF16_DTYPE,
        _FORWARD_SUPPORTED_KWARGS,
        MINIMAX_H3_ADALN_MODALITY_NUM,
        _modulate_scale_shift,
        _required_kwarg,
    )
    if not hasattr(module, "blocks") or len(module.blocks) == 0:
        raise ValueError("Module must have blocks")
    unexpected = sorted(set(kwargs) - _FORWARD_SUPPORTED_KWARGS)
    if unexpected:
        raise TypeError(
            "MiniMaxH3DiTModel.forward received unexpected kwargs: "
            f"{unexpected}; supported kwargs: "
            f"{sorted(_FORWARD_SUPPORTED_KWARGS)}"
        )
    x = _required_kwarg(kwargs, "x")
    audio_x = _required_kwarg(kwargs, "audio_x")
    img_position_ids = _required_kwarg(kwargs, "img_position_ids")
    unique_timesteps = _required_kwarg(kwargs, "unique_timesteps")
    inverse_indices = _required_kwarg(kwargs, "inverse_indices").view(-1).to(torch.long)
    update_mask = _required_kwarg(kwargs, "update_mask")
    token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long)
    skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False))
    update_audio_mask = kwargs.get("update_audio_mask")
    text_selected = _required_kwarg(kwargs, "prompt_embeds")
    img_pos = module._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info")
    audio_pos = module._pos_ids(_required_kwarg(kwargs, "audio_pos_info"), "audio_pos_info")
    text_pos = module._pos_ids(
        _required_kwarg(kwargs, "text_pos_info"),
        "text_pos_info",
    )
    infer_out_pos = module._pos_ids(
        _required_kwarg(kwargs, "img_pos_for_infer_output_info"),
        "img_pos_for_infer_output_info",
    )
    psp = _required_kwarg(kwargs, "packed_seq_params")
    cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
    max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
    refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
    refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
    refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
    video_layout = kwargs.get("video_token_layout")
    if x.dim() != 3 or x.shape[0] != 1:
        raise ValueError(f"x must be [1, S, C], got {list(x.shape)}")
    seq_len = int(x.shape[1])
    if token_tags.shape[0] != seq_len:
        raise ValueError(f"token_tags must cover the full packed sequence ({seq_len}), got {token_tags.shape[0]}.")
    if inverse_indices.shape[0] != seq_len:
        raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
    device = x.device
    rope_freqs = module.rope(img_position_ids).to(device)
    decoder_input, t_emb = module._embed(
        x=x,
        audio_x=audio_x,
        text_embeddings_selected=text_selected,
        unique_timesteps=unique_timesteps.view(-1).to(device),
        img_pos=img_pos.to(device),
        audio_pos=audio_pos.to(device),
        text_pos=text_pos.to(device),
        refiner_cu_seqlens=refiner_cu.to(device),
        refiner_max_seqlen=refiner_max,
        seq_len=seq_len,
        device=device,
    )
    combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)
    inverse_indices = inverse_indices.to(device)
    cu_seqlens = cu_seqlens.to(device)
    shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb)
    modulated_hidden = module.blocks[0].norm1(decoder_input)
    modulated_input = _modulate_scale_shift(
        modulated_hidden,
        shift_msa,
        scale_msa,
        combined_indices,
        dtype=_BF16_DTYPE,
    )
    def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
        hidden, block_rope, block_combined = module.sp_prepare(
            decoder_input,
            rope_freqs,
            combined_indices,
        )
        for block in module.blocks:
            hidden = block(
                hidden,
                t_emb=t_emb,
                combined_indices=block_combined,
                rope_freqs=block_rope,
                cu_seqlens=cu_seqlens,
                max_seqlen=max_seqlen,
                packed_total=seq_len,
                video_layout=video_layout,
            )
        hidden = module.sp_gather(hidden)
        return (hidden,)
    def postprocess(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        video_logits, audio_logits = module.final_layer(
            hidden,
            t_emb=t_emb,
            inverse_indices=inverse_indices,
        )
        video_logits = video_logits.index_select(0, infer_out_pos.to(device))
        audio_logits = audio_logits.index_select(0, audio_pos.to(device))
        if not skip_mask_out_condition:
            mask = update_mask.view(-1).to(device)
            if mask.shape[0] != video_logits.shape[0]:
                raise ValueError(f"update_mask length mismatch: {mask.shape[0]} != {video_logits.shape[0]}")
            video_logits = video_logits * mask.unsqueeze(-1)
            if update_audio_mask is not None:
                audio_logits = audio_logits * update_audio_mask.view(-1).unsqueeze(-1)
        return video_logits, audio_logits
    return CacheContext(
        modulated_input=modulated_input,
        hidden_states=decoder_input,
        encoder_hidden_states=None,
        temb=t_emb,
        run_transformer_blocks=run_transformer_blocks,
        postprocess=postprocess,
    )
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1411,7 +1265,6 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
    "LongCatImageTransformer2DModel": extract_longcat_context,
    "FluxTransformer2DModel": extract_flux_context,
    "SenseNovaU1ForCausalLM": extract_sensenova_u1_context,
    "MiniMaxH3DiTModel": extract_minimax_h3_context,
    # Future models:
    # "CogVideoXTransformer3DModel": extract_cogvideox_context,
 }
@@ -231,9 +231,7 @@ class TeaCacheHook(ModelHook):
        state.accumulated_rel_l1_distance += abs(rescaled_distance)
        # Decision: below threshold = cache, above = compute
-        rel_l1_thresh = self.config.rel_l1_thresh
+        if state.accumulated_rel_l1_distance < self.config.rel_l1_thresh:
        assert rel_l1_thresh is not None
        if state.accumulated_rel_l1_distance < rel_l1_thresh:
            return False  # Use cache
        else:
            state.accumulated_rel_l1_distance = 0.0  # Reset accumulator
@@ -23,10 +23,6 @@ from vllm.model_executor.layers.quantization.base_config import (
 )
 from vllm.transformers_utils.repo_utils import get_model_path
 from vllm_omni.diffusion.diffusion_kv.config import (
    DiffusionKVCacheMode,
    parse_diffusion_kv_cache_mode,
 )
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -416,14 +412,14 @@ class DiffusionCacheConfig:
        >>> # Access via attribute
        >>> print(config.rel_l1_thresh)  # 0.3 (from dict)
        >>> print(config.Fn_compute_blocks)  # 8 (default)
-        >>> # Empty dict defers model-specific defaults to the TeaCache backend
+        >>> # Empty dict uses all defaults
        >>> default_config = DiffusionCacheConfig.from_dict({})
-        >>> print(default_config.rel_l1_thresh)  # None
+        >>> print(config.rel_l1_thresh)  # 0.2 (default)
    """
    # TeaCache parameters [tea_cache only]
-    # None defers to the model-specific TeaCache default (0.2 fallback).
+    # Default: 0.2 provides ~1.5x speedup with minimal quality loss (optimal balance)
-    rel_l1_thresh: float | None = None
+    rel_l1_thresh: float = 0.2
    coefficients: list[float] | None = None  # Uses model-specific defaults if None
    # MagCache parameters [mag_cache only]
@@ -662,9 +658,6 @@ class OmniDiffusionConfig:
    # the string form only).
    engine_backend: str | type = "default"
    # Local Diffusion KV ownership and cache-layout mode.
    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
    # Optional override for the diffusion model runner class (import path).
    # Precedence in the worker: this override > the runner declared by the
    # selected engine class (``default_diffusion_model_runner_cls``) > the
@@ -711,9 +704,6 @@ class OmniDiffusionConfig:
    # VAE memory optimization parameters
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    # Optional Wan decode-only precision. When set, the VAE encoder remains
    # FP32 while decoder-side modules use the selected precision.
    vae_decode_precision: str | None = None
    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
@@ -926,11 +916,6 @@ class OmniDiffusionConfig:
        )
    def __post_init__(self):
        if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"):
            raise ValueError(
                "vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, "
                f"got {self.vae_decode_precision!r}"
            )
        if self.diffusion_compile_granularity not in {"regional", "full"}:
            raise ValueError(
                "diffusion_compile_granularity must be 'regional' or 'full', "
@@ -938,28 +923,10 @@ 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}")
        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
        if self.omni_kv_config is None:
            self.omni_kv_config = {}
        elif isinstance(self.omni_kv_config, Mapping):
            self.omni_kv_config = dict(self.omni_kv_config)
        else:
            raise TypeError("omni_kv_config must be a mapping")
        if self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and self.omni_kv_config.get(
            "need_recv_cache", False
        ):
            raise ValueError(
                "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
                "disable need_recv_cache until connector-aware admission is implemented"
            )
        self.master_port = self._resolve_master_port()
        self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
-        if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:
+        if self.request_batch_max_wait_ms < 0:
-            raise ValueError(
+            raise ValueError(f"request_batch_max_wait_ms must be non-negative, got {self.request_batch_max_wait_ms}.")
                f"request_batch_max_wait_ms must be a finite non-negative number, got {self.request_batch_max_wait_ms}."
            )
        if isinstance(self.profiler_config, dict):
            from vllm.config import ProfilerConfig
@@ -310,16 +310,14 @@ class DiffusionEngine:
        diffusion_engine_start_time = time.perf_counter()
        preprocess_time = 0.0
-        has_preprocessor = getattr(self, "pre_process_func", None) is not None
+        if self.pre_process_func is not None:
-        preprocess_start_time = time.perf_counter() if has_preprocessor else None
+            preprocess_start_time = time.perf_counter()
-        request = self._prepare_request_for_admission(request)
+            request = self.pre_process_func(request)
        if preprocess_start_time is not None:
            preprocess_time = time.perf_counter() - preprocess_start_time
            logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)
        exec_start_time = time.perf_counter()
-        request_id = self._add_prepared_request(request)
+        generator = self.async_add_req_and_stream_response(request)
        generator = self.get_output_stream(request_id)
        async for output in generator:
            exec_total_time = time.perf_counter() - exec_start_time
            # Async mode: wait for background D2H/SHM to complete.
@@ -434,7 +432,8 @@ class DiffusionEngine:
                    # Only RPC / abort work pending; loop back to drain it.
                    continue
-                self._wait_for_admission_if_needed_locked()
+                if self.supports_request_batch or self.dp_concurrent:
                    self._wait_for_request_batch_admission_locked()
                sched_output = self.scheduler.schedule()
@@ -468,39 +467,53 @@ class DiffusionEngine:
        # Engine is stopping: fail any RPCs still queued so callers don't hang.
        self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))
-    def _wait_for_admission_if_needed_locked(self) -> None:
+    def _wait_for_request_batch_admission_locked(self) -> None:
-        """Apply scheduler admission policy while holding the engine condition.
+        """Wait for compatible requests to accumulate before scheduling a wave.
        Caller must hold ``self._cv``.
        """
-        start = time.monotonic()
+        if self.step_execution or (not self.supports_request_batch and not self.dp_concurrent):
        decision = self.scheduler.get_admission_wait_decision(
            now=start,
            dp_concurrent=self.dp_concurrent,
        )
        if not decision.should_wait:
            return
        max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0
        if max_wait_s == 0:
            return
        max_batch = self.scheduler.max_num_running_reqs
        waiting = self.scheduler.num_waiting_requests()
        running = self.scheduler.num_running_requests()
        if running > 0:
            return
        start = time.monotonic()
        deadline = start + max_wait_s
        last_waiting = -1
        stable_since = start
        # For dp_concurrent, use a longer stable window so all dp_size
        # HTTP requests have time to land before scheduling.
        if self.dp_concurrent:
            stable_window_s = min(0.3, max_wait_s / 2.0)
        else:
            stable_window_s = min(0.05, max_wait_s / 5.0)
        while not self.stop_event.is_set():
            waiting = self.scheduler.num_waiting_requests()
            now = time.monotonic()
            if waiting >= max_batch:
                break
            if waiting > 0 and (now - stable_since) >= stable_window_s:
                break
            if now >= deadline:
                break
            if waiting > last_waiting:
                stable_since = now
                last_waiting = waiting
-            if self.scheduler.should_end_admission_wait(
+            remaining = deadline - now
-                decision,
+            self._cv.wait(timeout=min(remaining, 0.002))
                now=now,
                stable_since=stable_since,
            ):
                break
            remaining = decision.deadline - now if decision.deadline is not None else 0.002
            self._cv.wait(timeout=min(max(remaining, 0.0), 0.002))
        waited_ms = (time.monotonic() - start) * 1000.0
        final_waiting = self.scheduler.num_waiting_requests()
@@ -508,7 +521,7 @@ class DiffusionEngine:
            logger.info(
                "[RequestBatch] admission wait done waiting=%d max_batch=%d waited_ms=%.1f",
                final_waiting,
-                decision.max_batch,
+                max_batch,
                waited_ms,
            )
@@ -676,17 +689,7 @@ class DiffusionEngine:
        engine.run_startup_warmup()
        return engine
-    def _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+    def add_request(self, request: OmniDiffusionRequest) -> str:
        """Run model-owned preprocessing once, before entering Engine locks."""
        pre_process_func = getattr(self, "pre_process_func", None)
        if pre_process_func is None:
            return request
        return pre_process_func(request)
    def _add_prepared_request(self, request: OmniDiffusionRequest) -> str:
        """Admit a request whose model-owned preprocessing is complete."""
        with self._cv:
            if self._closed:
                raise RuntimeError("DiffusionEngine is closed.")
@@ -697,10 +700,6 @@ class DiffusionEngine:
        return request_id
    def add_request(self, request: OmniDiffusionRequest) -> str:
        request = self._prepare_request_for_admission(request)
        return self._add_prepared_request(request)
    async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
        with self._cv:
            queue = self._out_streams.get(request_id)
@@ -743,7 +742,6 @@ class DiffusionEngine:
        return final_output
    def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
        request = self._prepare_request_for_admission(request)
        with self._rpc_lock:
            if self._closed:
                raise RuntimeError("DiffusionEngine is closed.")
@@ -888,7 +886,8 @@ class DiffusionEngine:
            ),
        )
        logger.info("dummy run to warm up the model")
-        output = self.add_req_and_wait_for_response(req)
+        request = self.pre_process_func(req) if self.pre_process_func is not None else req
        output = self.add_req_and_wait_for_response(request)
        if output.error:
            raise RuntimeError(f"Dummy run failed: {output.error}")
@@ -1,2 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -1,39 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 from enum import Enum
 class DiffusionKVCacheMode(str, Enum):
    """Migration mode for diffusion KV ownership."""
    DENSE_LEGACY = "dense_legacy"
    # Reserved for the experimental Worker-owned paging design. Production
    # configuration rejects it until that engine is migrated to the common
    # Scheduler-owned path.
    PAGED_WORKER_LOCAL = "paged_worker_local"
    PAGED_SCHEDULER = "paged_scheduler"
 def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode:
    """Parse a selectable Diffusion KV ownership mode.
    ``PAGED_WORKER_LOCAL`` remains in the enum to document the migration
    topology, but it is not a valid production selection until that runtime is
    wired. Keep this validation shared by structured and runtime configs so an
    unsupported mode fails before managed-process startup.
    """
    try:
        mode = DiffusionKVCacheMode(value)
    except (TypeError, ValueError) as exc:
        supported = ", ".join(mode.value for mode in DiffusionKVCacheMode)
        raise ValueError(f"Unsupported Diffusion KV diffusion_kv_mode {value!r}; expected one of: {supported}") from exc
    if mode is DiffusionKVCacheMode.PAGED_WORKER_LOCAL:
        raise ValueError(
            "Diffusion KV diffusion_kv_mode 'paged_worker_local' is reserved but not implemented; "
            "use 'dense_legacy' or 'paged_scheduler'"
        )
    return mode
@@ -1,116 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 from collections.abc import Sequence
 from dataclasses import dataclass
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.request import RequestStatus
@dataclass(frozen=True)
 class DiffusionKVContext:
    """An independently managed K/V context outside the primary sequence.
    ``context_id`` identifies the logical context within one execution
    sequence. ``cache_role`` binds it to a logical attention cache role exposed
    by the Worker. Physical cache geometry remains native ``KVCacheSpec`` /
    ``KVCacheConfig`` state and is deliberately absent here.
    """
    context_id: str
    cache_role: str
    num_tokens: int
    block_hashes: tuple[BlockHash, ...] = ()
    def __post_init__(self) -> None:
        if not self.context_id:
            raise ValueError("context_id must be non-empty")
        if not self.cache_role:
            raise ValueError("cache_role must be non-empty")
        if self.num_tokens <= 0:
            raise ValueError(f"num_tokens must be positive, got {self.num_tokens}")
        object.__setattr__(self, "block_hashes", tuple(self.block_hashes))
 class DiffusionKVRequest:
    """Scheduler-owned KV state for one diffusion execution sequence.
    The primary sequence follows an ordered ``[prefix | target]`` policy, such
    as one Hunyuan CFG row. ``prefix_len`` is the contiguous reusable prefix;
    ``target_len`` is overwritten by every denoise step; ``num_tokens`` is the
    complete first-step allocation boundary.
    Independent cross/joint-attention K/V does not belong to that token axis
    and is described by ``kv_contexts``. This object also exposes the minimal
    mutable Request surface consumed by native ``KVCacheManager``.
    An empty ``block_hashes`` sequence means the prefix has no canonical cache
    identity yet. Such a request may use native request-local page allocation,
    but consumers must not publish it through ``KVCacheManager.cache_blocks``.
    Prefix publication becomes valid only after model preprocessing supplies
    one canonical hash for every cacheable full block.
    """
    def __init__(
        self,
        request_id: str,
        *,
        sequence_id: int,
        prefix_len: int,
        target_len: int,
        seq_len: int,
        block_hashes: Sequence[BlockHash] = (),
        kv_contexts: Sequence[DiffusionKVContext] = (),
    ) -> None:
        if not request_id:
            raise ValueError("request_id must be non-empty")
        if sequence_id < 0:
            raise ValueError(f"sequence_id must be non-negative, got {sequence_id}")
        if prefix_len < 0:
            raise ValueError(f"prefix_len must be non-negative, got {prefix_len}")
        if target_len <= 0:
            raise ValueError(f"target_len must be positive, got {target_len}")
        if seq_len <= 0:
            raise ValueError(f"seq_len must be positive, got {seq_len}")
        if prefix_len + target_len > seq_len:
            raise ValueError(
                "prefix_len + target_len must not exceed seq_len: "
                f"prefix_len={prefix_len}, target_len={target_len}, seq_len={seq_len}"
            )
        contexts = tuple(kv_contexts)
        if any(not isinstance(context, DiffusionKVContext) for context in contexts):
            raise TypeError("kv_contexts must contain only DiffusionKVContext values")
        context_ids = [context.context_id for context in contexts]
        if len(context_ids) != len(set(context_ids)):
            raise ValueError(f"kv_contexts must have unique context_id values, got {context_ids}")
        # Diffusion semantics consumed by the Scheduler-side facade.
        self.sequence_id = sequence_id
        self.prefix_len = prefix_len
        self.target_len = target_len
        self.kv_contexts = contexts
        # Native vLLM Request surface. Keep this list intentionally small and
        # cover it with real-KVCacheManager conformance tests.
        self.request_id = request_id
        self.num_tokens = seq_len
        self.num_prompt_tokens = prefix_len
        self.num_computed_tokens = 0
        self.block_hashes = list(block_hashes)
        self.skip_reading_prefix_cache = not self.block_hashes
        self.status = RequestStatus.WAITING
        self.num_preemptions = 0
        self.num_in_flight_tokens = 0
        # vLLM 0.26 uses this optional boundary when publishing cached blocks;
        # zero means that no sparse-retention boundary is active.
        self.shared_prefix_boundary = 0
    @property
    def seq_len(self) -> int:
        """Complete first-step sequence length and allocation boundary."""
        return self.num_tokens
@@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
            split_dims=(3, 4),
            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
            tile_spec=tile_spec,
-            output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype),
+            output_dtype=self.dtype,
        )
        return tiletask_list, grid_spec
		`@@ -1,2 +0,0 @@`
			`# SPDX-License-Identifier: Apache-2.0`
			`# SPDX-FileCopyrightText: Copyright contributors to the vLLM project`