Merge branch 'main' into perf/minimax-h3-lossless-hot-path

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: Wang  fuyin <wfy2003324@163.com>

.buildkite/cuda/test-nightly.yml

@@ -460,6 +460,21 @@ steps:
             exit $$EXIT2
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3"
+        key: nightly-diffusion-x2v-performance-minimax-h3
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_4
+
   - group: ":card_index_dividers: Diffusion Test"
     key: nightly-diffusion-test-group
     depends_on: upload-nightly-pipeline

.claude/skills/add-diffusion-model/SKILL.md

@@ -309,7 +309,7 @@ For Omni or custom models, create:
 
 Required updates:
 1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
-2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
+2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table)
 3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
 4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
 5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)

benchmarks/diffusion/backends.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import json
 import mimetypes
 import os
 import time
@@ -13,6 +14,17 @@ from tqdm import tqdm
 DEFAULT_EDITS_BOT_TASK = "think"
 
 
+def _is_minimax_h3_mixed_reference(model: str) -> bool:
+    """MiniMax-H3 Ref2VA consumes reference videos as file paths.
+
+    Its reference preparer ffprobes / ffmpegs the reference file, so single
+    ``video_reference`` uploads are sent as a base64 data URL: the server
+    persists them to disk and keeps ``source_path``, handing the pipeline a
+    real file path instead of a decoded frame list.
+    """
+    return "MiniMax-H3" in (model or "")
+
+
 @dataclass
 class RequestFuncInput:
     prompt: str
@@ -355,13 +367,21 @@ async def async_request_v1_videos(
         )
     elif input.video_paths and len(input.video_paths) > 0:
         video_path = input.video_paths[0]
-        image_file = open(video_path, "rb")
-        form.add_field(
-            "input_reference",
-            image_file,
-            filename=os.path.basename(video_path),
-            content_type=_guess_mime_type(video_path),
-        )
+        if _is_minimax_h3_mixed_reference(input.model):
+            with open(video_path, "rb") as reference_file:
+                video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
+            form.add_field(
+                "video_reference",
+                json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
+            )
+        else:
+            image_file = open(video_path, "rb")
+            form.add_field(
+                "input_reference",
+                image_file,
+                filename=os.path.basename(video_path),
+                content_type=_guess_mime_type(video_path),
+            )
 
     job_id = None
     job_status = None

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -858,7 +858,12 @@ class RandomDataset(BaseDataset):
             height = int(self.args.height or profile.get("height") or 720)
             num_frames = int(self.args.num_frames or profile.get("num_frames") or 16)
             fps = float(self.args.fps or profile.get("fps") or 8)
-            writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+            # cv2 typically writes MPEG-4 Part 2 (``mp4v``), which pipelines such
+            # as MiniMax-H3 Ref2VA reject because they require H.264/H.265
+            # reference videos. Write the synthetic clip first, then re-encode
+            # it to H.264 when ffmpeg is available; otherwise keep ``mp4v``.
+            raw_path = video_path + ".raw.mp4"
+            writer = cv2.VideoWriter(raw_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
             if not writer.isOpened():
                 raise RuntimeError("cv2.VideoWriter failed to open")
             for frame_idx in range(num_frames):
@@ -871,6 +876,10 @@ class RandomDataset(BaseDataset):
                 frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220)
                 writer.write(frame)
             writer.release()
+            if _reencode_video_to_h264(raw_path, video_path):
+                os.unlink(raw_path)
+            else:
+                os.replace(raw_path, video_path)
             return [video_path]
         except Exception as e:
             raise RuntimeError(
@@ -878,6 +887,40 @@ class RandomDataset(BaseDataset):
             ) from e
 
 
+def _reencode_video_to_h264(source: str, destination: str) -> bool:
+    """Re-encode a video to H.264/MP4 using ffmpeg when available.
+
+    Returns False when ffmpeg is missing or the re-encode fails so callers can
+    fall back to the original container/codec.
+    """
+    import shutil
+    import subprocess
+
+    if shutil.which("ffmpeg") is None:
+        return False
+    try:
+        result = subprocess.run(
+            [
+                "ffmpeg",
+                "-y",
+                "-loglevel",
+                "error",
+                "-i",
+                source,
+                "-c:v",
+                "libx264",
+                "-pix_fmt",
+                "yuv420p",
+                destination,
+            ],
+            capture_output=True,
+            text=True,
+        )
+        return result.returncode == 0 and os.path.exists(destination)
+    except Exception:
+        return False
+
+
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
 

docs/.nav.yml

@@ -68,12 +68,12 @@ nav:
         - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
         - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
         - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
-      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload.md
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md
       - Execution Modes: user_guide/diffusion/execution_modes.md
     - Quantization:
-      - Overview: user_guide/quantization/overview.md
+      - Quantization: user_guide/quantization/overview.md
       - Online Quantization: user_guide/quantization/online.md
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
@@ -106,7 +106,6 @@ nav:
     - Feature Design:
       - Runtime and Stage Execution:
         - Disaggregated Inference: design/feature/disaggregated_inference.md
-        - Ray-based Execution: design/feature/ray_based_execution.md
         - Async Chunk: design/feature/async_chunk.md
         - Async Diffusion Output: design/feature/async_diffusion_output.md
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
@@ -130,10 +129,11 @@ nav:
           - VAE Patch Parallelism: design/feature/vae_parallel.md
         - Attention Backends:
           - Skip-Softmax: design/feature/skip_softmax.md
-        - Quantization: user_guide/quantization/overview.md
+        - Quantization: design/feature/quantization.md
         - Cache-DiT: design/feature/cache_dit.md
         - TeaCache: design/feature/teacache.md
         - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
+        - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
     - Infrastructure and Performance:
       - Prometheus Metrics: design/metrics.md
       - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -1,6 +1,6 @@
 **Scope**
 
-For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
 
 **Test Case Design**
 

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -200,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par
 | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
 | request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
 | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
-| baseline        | object      | No       | see below       | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.
+| baseline        | object      | No       | see below       | Optional thresholds. **Hardware-nested only**: top-level keys must be `[hardware-resource]` pytest markers from `pyproject.toml` (`get_hardware_mark_list()`). Under each hardware key, metric names are free-form; values are scalars or sweep-aligned lists. Example: `{"H100": {"throughput_qps": […], "custom_stage_ms": 1.0}, "A3": {…}}`. Runtime filename aliases (`_RUNTIME_DEVICE_ALIASES`) are separate and do not expand this allowlist. |

docs/contributing/model/adding_diffusion_model.md

@@ -824,7 +824,7 @@ omni = Omni(model="your-model",
 
 ### CPU Offload
 
-See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md)
+See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md)
 
 vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
 

docs/design/feature/distributed_layerwise_offload.md

@@ -0,0 +1,156 @@
+# Distributed Layerwise Offload
+
+This document describes distributed layerwise offload (DLO) for diffusion
+models. DLO keeps only a small number of DiT blocks on the accelerator and
+streams the remaining blocks from host memory. The distributed backend can
+either shard those host-side weights across an existing parallel group or keep
+the standard loader's rank-local weights and avoid an additional collective.
+
+For user-facing commands, see the
+[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
+and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
+
+## Status
+
+DLO is implemented for multi-device diffusion execution. The default
+AllGather path is the primary path for DP and SP deployments. The
+`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
+for standard-loader sharding, workstation bring-up, and systems where an
+additional DLO collective is undesirable, but it does not reduce host weight
+storage across ranks.
+
+The compatibility matrix below describes the current implementation. The
+unit-level guards are covered, but not every parallelism combination has a
+full model-and-hardware end-to-end test.
+
+## Design
+
+### DLO consumes the existing parallel topology
+
+DLO does not create a new DP, TP, or SP topology. It reads the configured
+`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
+standard distributed groups have been initialized.
+
+The DLO weight-sharding group is selected as follows:
+
+1. Use the existing DP group when `data_parallel_size > 1`.
+2. When DP is one and SP is greater than one, use the SP group.
+3. Otherwise, run rank-locally without a DLO process group.
+
+TP is deliberately not used as DLO's AllGather group. HSDP has its own
+parameter-sharding lifecycle and is not allowed to be sharded a second time by
+DLO's AllGather path.
+
+### AllGather path
+
+With the default `dlo_use_allgather=True`, each rank stores approximately
+`1 / group_size` of each streamable block in pinned host memory. The next
+block's shard is copied to a device buffer and reconstructed with
+`all_gather_into_tensor` on a communication stream while the current block is
+executing.
+
+```text
+Compute:    [Block N]             [Block N+1]          [Block N+2]
+H2D:                      [shard N+1]           [shard N+2]
+AllGather:                [full N+1]             [full N+2]
+Buffers:    [current slot]       [prefetch slot]       [current slot]
+```
+
+![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif)
+
+The backend uses two shared device buffers, so accelerator weight residency is
+bounded by the largest streamed blocks rather than the complete model.
+
+When DP is greater than one, the engine can process one request per DP rank in
+the same denoising wave. Because AllGather is a collective, all participating
+requests must take the same execution path at every denoising step.
+
+### Rank-local path without DLO AllGather
+
+With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
+one. The regular model loader remains responsible for preparing each rank's
+weights, including TP-local tensors or HSDP-managed parameters. DLO then
+streams those rank-local tensors block by block using H2D copies only.
+
+This mode means:
+
+- DP still provides independent replicas, but DLO does not shard weights
+  across DP ranks.
+- SP still performs its normal activation/attention collectives, but DLO does
+  not shard weights across SP ranks.
+- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
+  DLO's additional weight AllGather is disabled.
+- Pure DP deployments keep a full host-side model copy per rank, subject to
+  shared OS page-cache behavior.
+- The scheduler does not require a synchronized DP request wave for DLO.
+
+This path is intentionally not implemented by reusing the AllGather mmap
+loader. It relies on the standard loader so model-specific TP/HSDP transforms
+remain intact.
+
+## Parallelism compatibility
+
+| Parallelism | DLO + AllGather | DLO without AllGather |
+|---|---|---|
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
+| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
+| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
+
+### Combined dimensions
+
+- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
+  than one; SP continues to use its own sequence-parallel group. If DP is one,
+  the SP group becomes DLO's sharding group in AllGather mode.
+- **DP + TP/SP without AllGather:** standard model loading defines the
+  rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
+  collective.
+- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
+  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
+  materialization and synchronization.
+- **HSDP + DP or TP:** rejected independently by the diffusion parallel
+  configuration.
+
+## Request and loading constraints
+
+AllGather DP multi-concurrency requires:
+
+- explicit `num_inference_steps`;
+- the same `num_inference_steps` for all requests in a wave; and
+- identical request arguments that affect the collective execution path.
+
+The no-AllGather path does not impose these DLO-specific synchronized-wave
+requirements.
+
+The mmap loader is used only by the supported DLO+AllGather path when the
+model has a compatible checkpoint layout. Online quantization is incompatible
+with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
+quantization. Models that do not meet the mmap requirements use the regular
+loader path.
+
+## Validation coverage
+
+Current source-level validation includes:
+
+- HSDP + DLO + AllGather rejection;
+- HSDP + DLO without AllGather acceptance at configuration level;
+- TP rejection in the DLO+AllGather mmap path;
+- resident-layer requests requiring no-AllGather;
+- DP request-wave validation for denoising-step compatibility;
+- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
+  tests.
+
+The highest-value missing coverage is end-to-end numerical comparison against
+ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
+on the target CUDA/NCCL or CANN/HCCL hardware.
+
+## Recommendations
+
+- Use **DP + DLO AllGather** for the supported throughput and host-memory
+  scaling path.
+- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
+  not the goal.
+- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
+  the expectation of higher host-memory use and lower validation confidence.
+- Prefer **HSDP alone** for production HSDP deployments until the combined
+  HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/design/feature/quantization.md

@@ -0,0 +1,249 @@
+# Quantization
+
+This document describes the quantization architecture in vLLM-Omni and the
+extension points for contributors. For user-facing configuration, supported
+hardware, and method-specific instructions, see the
+[quantization user guide](../../user_guide/quantization/overview.md).
+
+## Goals and scope
+
+The design has four goals:
+
+1. Provide one `quantization_config` entry point for diffusion-only,
+   multi-stage omni/TTS, and multi-stage diffusion models.
+2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where
+   a vLLM backend already provides the required behavior.
+3. Allow different components of one pipeline to use different quantization
+   configurations.
+4. Preserve the distinction between online quantization of a BF16/FP16
+   checkpoint and loading a pre-quantized checkpoint with serialized scales.
+
+This document covers configuration resolution, component routing, model
+integration, backend extension, and validation. Kernel implementations and
+user-facing support matrices are outside its scope.
+
+## Architecture
+
+The quantization path is split across the following layers:
+
+| Layer | Responsibility | Source |
+| --- | --- | --- |
+| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` |
+| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` |
+| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` |
+| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` |
+| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` |
+| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` |
+
+The runtime contract remains vLLM's `QuantizationConfig`:
+
+```text
+user/CLI or model config
+          |
+          v
+OmniDiffusionConfig.quantization_config
+          |
+          +--> build_quant_config()
+          |        |
+          |        +--> Omni override
+          |        +--> vLLM registry backend
+          |        +--> ComponentQuantizationConfig
+          |
+          +--> checkpoint metadata reconciliation
+          |
+          +--> diffusion worker / vLLM config
+          |
+          +--> model layers call get_quant_method(layer, prefix)
+```
+
+The model layer receives either a method implementation or `None`. Returning
+`None` leaves that layer unquantized.
+
+## Configuration resolution
+
+`build_quant_config()` accepts the following forms:
+
+| Input | Resolution |
+| --- | --- |
+| `None` or `"none"` | Disable quantization. |
+| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. |
+| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. |
+| Per-component dictionary | Construct a `ComponentQuantizationConfig`. |
+| Existing `QuantizationConfig` | Pass through without rebuilding it. |
+
+Method aliases are normalized case-insensitively, with `-` and `_` treated as
+equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
+ModelOpt checkpoint metadata is detected from its method, producer metadata, or
+`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
+
+The factory uses Omni-specific overrides for methods whose diffusion behavior
+or platform dispatch is not provided by the vLLM registry. The current override
+set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
+AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through
+vLLM's registry when their configuration is compatible.
+
+## Per-component routing
+
+Multi-stage models do not necessarily want one quantization policy for every
+module. A component map expresses that policy without requiring each pipeline
+to construct separate vLLM configurations:
+
+```python
+quantization_config = {
+    "transformer": {"method": "fp8"},
+    "vae": None,
+    "default": None,
+}
+```
+
+`ComponentQuantizationConfig` applies the following rules:
+
+1. Component keys are interpreted as runtime layer-prefixes.
+2. The longest matching prefix wins.
+3. `None` disables quantization for the matching component.
+4. `default`, when present, handles prefixes with no explicit match.
+5. A matching child config delegates to its normal
+   `get_quant_method(layer, prefix)` implementation.
+
+Prefix matching happens after the model's weight-name mapping. New model
+integrations must verify the actual runtime prefixes, especially when a
+`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route
+a layer to the default configuration, so component-routing tests should cover
+both positive matches and fall-through behavior.
+
+## Online and pre-quantized checkpoints
+
+The same method name can represent two different loading paths:
+
+- **Online**: the source checkpoint contains BF16/FP16 weights; the method
+  creates quantized weights or scales while the model is loaded.
+- **Pre-quantized**: the checkpoint contains quantized weights and the metadata
+  required by the method; the loader consumes those serialized tensors.
+
+Method configuration owns the mode flag, for example
+`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The
+configuration's `get_quant_method()` selects the corresponding linear method.
+
+For pipelines with transformer-specific `config.json` files, such as cascade
+models, `resolve_quant_config_from_disk()` reconciles the active configuration
+with each transformer's checkpoint metadata. Its invariants are:
+
+- If no active configuration exists, valid checkpoint metadata can be used for
+  auto-detection.
+- A method mismatch is an error rather than an implicit conversion.
+- Serialized-checkpoint flags rebuild an online configuration into the
+  matching offline configuration.
+- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active
+  configuration when they differ.
+- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path.
+
+This prevents a model from interpreting serialized weights with an online
+loader or silently applying the wrong method to one transformer in a
+multi-transformer pipeline.
+
+## Model and layer integration
+
+Diffusion model constructors receive the resolved configuration and pass it to
+quantizable vLLM layers. Each method configuration is responsible for:
+
+- identifying supported layers in `get_quant_method()`;
+- creating the correct parameter shapes and loaders;
+- processing weights after loading when conversion is required; and
+- applying the platform-specific quantized operation at runtime.
+
+Model code must keep precision-sensitive or unsupported modules unquantized.
+The shared `safe_quant_config()` helper is used by model integrations for
+normalization and modulation layers where a generic FP8 configuration would be
+unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
+are not implicitly quantized just because the main transformer is quantized;
+they require explicit component routing and method support.
+
+Weight quantization and quantized KV cache are separate paths. The former is
+resolved through `quantization_config`; runtime attention/KV-cache quantization
+is configured by the diffusion attention configuration and must be validated
+independently.
+
+## Platform and parallelism boundaries
+
+Platform checks belong in the method implementation rather than in the common
+factory. This keeps configuration parsing portable while allowing CUDA, NPU,
+and XPU methods to select different kernels or reject unsupported execution at
+model construction time. The user guide is the source of truth for the current
+hardware support matrix.
+
+Quantization methods must also respect the weight layout produced by the
+configured parallelism strategy:
+
+- Tensor-parallel layers may receive partitioned weight shapes and loaders.
+- Sequence-parallel execution changes activation movement but does not by
+  itself change the quantization configuration.
+- Data-parallel replicas build their own method state; quantization metadata
+  must be available consistently on every rank.
+- `WeightsMapper` transformations must be applied before matching ignored-layer
+  lists or component prefixes.
+
+Distributed layerwise offload has an additional boundary: its sharded mmap
+AllGather path is not compatible with online quantization because the path
+expects already materialized checkpoint weights. Use the rank-local
+`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
+combining these features. See the [DLO design document](distributed_layerwise_offload.md)
+for the full compatibility matrix.
+
+## Adding a quantization backend
+
+Contributors should follow this sequence:
+
+1. Check whether an existing vLLM backend can provide the required weight
+   format and kernel behavior. Prefer reuse over a duplicate Omni backend.
+2. If Omni-specific behavior is required, implement a `QuantizationConfig` and
+   its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`.
+3. Implement `get_name()`, supported activation dtypes, minimum capability,
+   checkpoint parsing, ignored-layer handling, and `get_quant_method()`.
+4. Add a lazy factory override in `factory.py`. Keep optional platform
+   dependencies out of package import paths and normalize aliases to one
+   canonical method name.
+5. If the method is pre-quantized, define the checkpoint metadata and the
+   serialized loading path before adding an online path.
+6. Add model-specific component routing only when the model contains stages
+   that must use different policies.
+7. Document user-visible configuration in the quantization user guide and keep
+   implementation rationale here.
+
+The public factory registration hook,
+`register_quantization_override(method, builder)`, is intended for integrations
+that need to register an Omni-compatible builder without changing the vLLM
+registry.
+
+## Validation
+
+Validation should cover the complete path, not only whether a config object can
+be instantiated:
+
+| Area | Required checks |
+| --- | --- |
+| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps |
+| Method config | Online/offline selection, ignored layers, platform guards, and weight loading |
+| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata |
+| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts |
+| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts |
+| Performance | Peak device memory, generation latency, and any backend-specific throughput metric |
+
+The unit tests live under `tests/diffusion/quantization/`. The trajectory
+similarity tool in
+`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py`
+provides a repeatable quality and memory comparison for diffusion models.
+Method support and feature combinations remain model- and hardware-dependent;
+passing factory tests alone does not establish production support.
+
+## Design invariants
+
+- `quantization_config` is the single runtime source of truth after config
+  normalization.
+- A checkpoint method mismatch fails loudly instead of risking weight
+  corruption.
+- Components without an explicit supported quantization path remain in their
+  original precision.
+- Optional backend dependencies are imported lazily.
+- Method implementations own platform-specific capability checks and kernels.
+- User-facing support claims require model, hardware, numerical, and memory
+  validation in addition to configuration-level tests.

docs/design/feature/ray_based_execution.md

@@ -1,59 +0,0 @@
-# Distributed utils
-
-This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
-## 1. Installation
-```bash
-pip install "ray[default]"
-```
-## 2. Ray Utils
-
-The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
-*   **Multi-node deployment**: Running pipeline stages across different physical machines.
-*   **Resource management**: Efficient GPU/CPU allocation.
-
-### 2.1 Basic Usage
-
-To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
-
-**Command Line Example:**
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B \
-  --omni \
-  --port 8091 \
-  --worker-backend ray \
-  --ray-address auto
-```
-
-### 2.2 Cluster Setup
-
-**Step 1: Start Head Node**
-Run this on your primary machine:
-```bash
-ray start --head --port=6399
-```
-
-**Step 2: Connect Worker Nodes**
-Run this on each worker machine:
-```bash
-ray start --address=<HEAD_NODE_IP>:6399
-```
-
-> **Tip**: For a complete cluster setup script, refer to the vLLM example:
-> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
-
-### 2.3 Distributed Connector Support
-
-When running on Ray, the system automatically adapts its communication strategy:
-
-*   **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
-*   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
-*   **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`.
-
-### 2.4 Internal Helpers
-
-*   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
-
-## 3. Troubleshooting
-
-*   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
-*   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.

docs/design/index.md

@@ -13,7 +13,6 @@ while the source paths remain stable so existing links continue to work.
 ### Runtime and stage execution
 
 - [Disaggregated Inference](feature/disaggregated_inference.md)
-- [Ray-based Execution](feature/ray_based_execution.md)
 - [Async Chunk](feature/async_chunk.md)
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
@@ -48,11 +47,12 @@ while the source paths remain stable so existing links continue to work.
 
 #### Quantization
 
-- [Quantization Overview](../user_guide/quantization/overview.md)
+- [Quantization](feature/quantization.md)
 
 - [Cache-DiT](feature/cache_dit.md)
 - [TeaCache](feature/teacache.md)
 - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
+- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
 
 ## Infrastructure and Performance
 

docs/design/module/dit_module.md

@@ -539,7 +539,7 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ
 
 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
 
-A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code.
+A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code.
 
 For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
 

docs/user_guide/diffusion/attention_backends.md

@@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 
 | Value | Notes |
 |---|---|
-| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
+| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
 | `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
@@ -23,6 +23,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
 | `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
 | `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
+| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
 
 
 ## Configuration
@@ -69,8 +70,9 @@ vllm-omni serve <model> \
     --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
 ```
 
-A backend that needs configuration exposes it as a typed field on the spec. Today the only
-one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)):
+A backend that needs configuration exposes it as a typed field on the spec:
+`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
+`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
 
 ```bash
 --diffusion-attention-config.default.backend TRTLLM_ATTN \
@@ -103,7 +105,7 @@ A plain dict is also accepted and normalized to `AttentionConfig`.
 
 Auto-route preference, in order:
 
-1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model is mask-free** (see below)
+1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
 2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
 3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
 4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
@@ -111,7 +113,7 @@ Auto-route preference, in order:
 
 `TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
 
-`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell because it is the only backend that can enable Skip-Softmax, not because its dense kernel is faster — dense, the two are comparable. It cannot honor an attention mask, so the auto-default only applies to pipelines verified mask-free (the Wan family); mask-using pipelines keep `CUDNN_ATTN`. `TRTLLM_ATTN` stays available everywhere via explicit `--diffusion-attention-backend TRTLLM_ATTN` (which raises clearly if a mask is received).
+`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
 
 The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
 
@@ -204,11 +206,80 @@ AttentionConfig(
 )
 ```
 
-
 **Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
 availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
 (B300); `int8` QK kernels are compiled for **SM100 only**.
 
+## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
+
+`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
+128-token block of keys, ranks them per query block, and computes attention against only the
+top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
+spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
+
+Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
+first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
+nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
+exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
+delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
+
+Enable it through the typed `block_sparse` block on the attention spec, shared by every
+block-sparse backend:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
+| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
+| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
+      "block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
+```
+
+Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="RAINFUSION_ATTN",
+        block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
+    ),
+)
+```
+
+**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
+cheapest way to recover structure, because the early high-noise steps decide global layout while the
+later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
+`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
+
+Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
+incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
+blocks — use Ulysses SP (`ring_degree=1`).
+
+### Which geometries run sparse
+
+Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
+is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
+speedup, so pick an aligned geometry when you want one.
+
+Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
+spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
+to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
+
+"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
+resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+which you can buy back with `start_step`.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.

docs/user_guide/diffusion/cpu_offload.md

@@ -184,10 +184,15 @@ class Flux2Transformer2DModel(nn.Module):
 ### How It Works
 
 Distributed layerwise offloading extends single-GPU layerwise offloading to
-multi-device deployments.  Each DP rank stores only **1/dp_size** of the model
-weights in host memory; full layer weights are reconstructed at runtime via
-**AllGather** on a dedicated communication stream, overlapped with computation
-via a fixed double-buffer scheme.
+multi-device deployments. In the default mode, each rank stores only
+**1/dp_size** of the model weights in host memory; full layer weights are
+reconstructed at runtime via **AllGather** on a dedicated communication stream,
+overlapped with computation via a fixed double-buffer scheme. The
+`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
+weights independently and does not shard weights across ranks.
+
+For the implementation design and compatibility matrix, see
+[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
 
 **Key features:**
 - **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
@@ -219,7 +224,7 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4
 
-# Without AllGather (each rank loads full weights, no sharding)
+# Without DLO AllGather (each rank streams standard-loader rank-local weights)
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4 \
@@ -247,9 +252,9 @@ m = Omni(
 | Flag | Description | Default |
 |------|-------------|---------|
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
+| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
 | `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
-| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
+| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
 
 ### How model weights are loaded (mmap path)
 
@@ -303,9 +308,15 @@ request while AllGather synchronizes only weight shards (request-independent).
 
 ### Limitations
 - Online quantization (FP8) is incompatible with mmap loading — falls back to
-  regular `load_weights()` automatically
-- Tensor Parallel is not supported (DLO uses DP-based sharding)
-- HSDP + AllGather is rejected (would double-shard weights)
+  regular loading only when the model is outside the DLO+AllGather mmap path;
+  online quantization with DLO+AllGather is rejected. Use
+  `--dlo-no-use-allgather` or disable online quantization.
+- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
+  weight-loader callbacks are bypassed. The no-AllGather path retains the
+  standard TP loader and is experimental.
+- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
+  with no-AllGather is accepted at configuration level, but has limited
+  end-to-end validation.
 - `num_inference_steps=None` is not allowed in DP multi-concurrency mode
 
 **Module Discovery**
@@ -338,7 +349,7 @@ OffloadBackend (base class)
 ├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 │                          (single-GPU, full weights on host)
 └── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
-                                         (multi-GPU, 1/dp_size sharded weights + AllGather)
+                                         (multi-GPU, sharded AllGather or rank-local mode)
 ```
 
 Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -369,4 +380,10 @@ reasoner/generator components inside the model forward pass.
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
 - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
-- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required.  See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
+- Distributed Layerwise Offloading reuses the layerwise block topology, but the
+  model and parallelism combination still needs validation. AllGather mode
+  additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
+  retains the standard loader's rank-local tensors. See the
+  [DLO design](../../design/feature/distributed_layerwise_offload.md) and
+  [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
+  for usage examples.

docs/user_guide/diffusion_features.md

@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
 
 | Method | Description | Best For |
 |--------|-------------|----------|
-| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
+| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
 | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
 | **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
 
@@ -104,8 +104,8 @@ The following tables show which models support each feature:
 
 > Notes:
 
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
-> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
 
@@ -274,9 +274,9 @@ Measured on NVIDIA H800:
 
 **Memory Optimization:**
 
-- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
+- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
 - **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
-- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
+- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 
 **Extensions:**
 

docs/user_guide/quantization/overview.md

@@ -5,6 +5,9 @@ path. The same configuration entrypoint is used across diffusion-only models,
 multi-stage omni/TTS models, and multi-stage diffusion models, but each model
 type has a different quantization scope.
 
+For the internal architecture and backend extension points, see the
+[quantization design document](../../design/feature/quantization.md).
+
 ## Quantization Modes
 
 | Mode | Guide | Description | Methods |

mkdocs.yml

@@ -75,6 +75,8 @@ plugins:
       redirect_maps:
         user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
         user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
+        user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
+        user_guide/quantization/index.md: user_guide/quantization/overview.md
         design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
         design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
   - glightbox

pyproject.toml

@@ -250,13 +250,16 @@ markers = [
     "npu: Tests that run on NPU/Ascend (auto-added)",
     "musa: Tests that run on MUSA/Moore Threads (auto-added)",
     # specified computation resources marks (auto-added)
-    "H100: Tests that require H100 GPU",
-    "L4: Tests that require L4 GPU",
-    "MI325: Tests that require MI325 GPU (AMD/ROCm)",
-    "B60: Tests that require Intel Arc Pro B60 XPU",
-    "S5000: Tests that require S5000 GPU (Moore Threads/MUSA)",
-    "A2: Tests that require A2 NPU",
-    "A3: Tests that require A3 NPU",
+    # Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list().
+    "H100: [hardware-resource] Tests that require H100 GPU",
+    "H200: [hardware-resource] Tests that require H200 GPU",
+    "L4: [hardware-resource] Tests that require L4 GPU",
+    "B200: [hardware-resource] Tests that require B200 GPU",
+    "MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)",
+    "B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU",
+    "S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)",
+    "A2: [hardware-resource] Tests that require A2 NPU",
+    "A3: [hardware-resource] Tests that require A3 NPU",
     "distributed_cuda: Tests that require multi cards on CUDA platform",
     "distributed_rocm: Tests that require multi cards on ROCm platform",
     "distributed_xpu: Tests that require multi cards on XPU platform",

recipes/MiniMaxAI/MiniMax-H3-MUSA.md

@@ -12,19 +12,17 @@
 - Maintainer: Community
 
 This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments.
-MiniMax H3 has two independently served checkpoint partitions:
-
-- `FL2VA` serves text-to-video+audio (`t2va`) and
-  first-frame-to-video+audio (`fl2va`).
-- `Ref2VA` serves image/audio or video-reference generation (`ref2va`).
+The modular H3 pipeline can load both task-specific DiTs, but this MUSA recipe
+uses a **single task partition**: `FL2VA` for `t2va`/`fl2va`, or `Ref2VA` for
+`ref2va`. This avoids loading both DiTs.
 
 ## Prerequisites
 
 ### Checkpoint
 
-Only download the partition needed by the server. Downloading the whole model
-repository also retrieves duplicate shared components and the original
-top-level pipeline layout.
+Download the partition required by the selected startup task. The commands below
+run a single-task server; do not use the modular default (combined) mode on MUSA
+until its two-DiT memory and performance profile is validated.
 
 For T2VA and FL2VA:
 
@@ -66,14 +64,13 @@ soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg.
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA, or to `Ref2VA` for Ref2VA.
+Use the partition path and set `--task-type` explicitly: `fl2va` for T2VA/FL2VA
+or `ref2va` for Ref2VA.
 
-The validated Ref2VA configuration uses tensor parallelism across four MTT
-S5000 GPUs and offloads inactive model components to CPU:
+The validated Ref2VA configuration uses tensor parallelism across four MTT S5000 GPUs and offloads inactive model components to CPU:
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL="${MODEL_ROOT}/Ref2VA"
 export PORT=8091
 
 MUSA_VISIBLE_DEVICES=0,1,2,3 \
@@ -85,6 +82,7 @@ vllm serve "${MODEL}" \
   --port "${PORT}" \
   --trust-remote-code \
   --num-gpus 4 \
+  --task-type ref2va \
   --tensor-parallel-size 4 \
   --vae-patch-parallel-size 4 \
   --enable-cpu-offload \
@@ -133,7 +131,8 @@ Keep `--vae-use-tiling` enabled for this serving profile.
 - Keep Ring Attention at degree 1. The current Ring path does not preserve
   MiniMax H3 packed padding boundaries.
 - H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
-- Each server process loads only one checkpoint partition.
+- This MUSA recipe uses an explicitly selected single task partition; combined
+  serving loads both DiTs.
 - H3 currently executes one generation request per diffusion batch.
 - FP8 quantization has not been enabled for MiniMax H3.
 - MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -26,12 +26,12 @@ environments. Differences from the GPU path:
 
 ### Checkpoint
 
-Same as the GPU recipe — Hugging Face access approval is required:
+Same as the GPU recipe — Hugging Face access approval is required. Authenticate
+once; `vllm serve` downloads the required nested components automatically:
 
 ```bash
 hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+export MODEL=MiniMaxAI/MiniMax-H3
 ```
 
 ### Environment
@@ -50,7 +50,8 @@ uv pip install -e .
 ```
 
 Install the **mindie-sd** operator library to enable Ascend-optimized fused
-operators (`adalayernorm`, etc.):
+operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse
+attention kernel:
 
 ```bash
 git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
@@ -72,24 +73,21 @@ pip install mindiesd-*.whl
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
+Pass the repository ID directly. The pipeline loads the two nested DiTs while
+sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
 
-### Multi-NPU: 768P validated configuration
+### Multi-NPU: 768P combined-service configuration
 
-Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
-parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
-VAE patch parallelism degree 8, and layerwise offload:
+Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
+degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
 
 ```bash
 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 export PORT=9098
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
 export PYTHONDONTWRITEBYTECODE=1
-export HF_HUB_OFFLINE=1
-export TRANSFORMERS_OFFLINE=1
 
 vllm serve "${MODEL}" \
   --omni \
@@ -111,12 +109,31 @@ Do not add `--enforce-eager`. The first request includes regional
 compilation; warm the server once before measuring steady-state latency.
 H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 
-To serve Ref2VA, stop the FL2VA server and restart with:
+The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
+partition restart is required. Layerwise offload applies to both DiTs.
+
+### Optional optimizations
+
+Two independent optimizations may be enabled on top of the configuration
+above. Both are validated for T2VA only.
+
+**RainFusion block-sparse attention** — switch the attention backend, keeping
+every other flag unchanged:
 
 ```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
+  --diffusion-attention-backend RAINFUSION_ATTN
 ```
 
+**INT8 online quantization** — add one flag to the server command above:
+
+```bash
+  --quantization int8
+```
+
+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.
+
 ## HTTP API examples
 
 The request format is identical to the GPU recipe; see
@@ -145,7 +162,8 @@ throughput guarantee.
 
 ## Known limitations
 
-- Each server process loads only one checkpoint partition.
+- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
+  both task-specific DiTs, and loads one copy of every shared component.
 - H3 currently executes one generation request per diffusion batch.
 - The first regional-compile request is a warmup and should not be included
   in steady-state performance measurements.
@@ -154,9 +172,16 @@ throughput guarantee.
   the same HTTP request contract.
 - VAE patch parallelism requires size 1 or the full DiT group size and
   supports the H3 native `tile` mode only.
+- RainFusion block-sparse attention and INT8 quantization are validated for
+  T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA.
+- Online quantization cannot be combined with distributed layerwise offload
+  while AllGather is enabled; pass `--dlo-no-use-allgather` in that case.
 
 ## Additional resources
 
 - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention)
+  — RainFusion knobs and tuning
+- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
 - [Supported models](../../docs/models/supported_models.md)
 - [Video API](../../docs/serving/videos_api.md)

recipes/MiniMaxAI/MiniMax-H3.md

@@ -12,27 +12,34 @@
 - Maintainer: Community
 
 MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
-checkpoint has two independently served partitions:
+checkpoint has two task-specific DiT partitions:
 
 - `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
   (`fl2va`)
 - `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported
   image/video/audio combinations (`ref2va`; audio-only is rejected)
 
+One vLLM-Omni diffusion stage can load both DiTs while instantiating the
+tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
+once. Requests select the DiT with `extra_params.task`.
+
 The generated MP4 contains H.264 video and synchronized stereo audio.
 
 ## Prerequisites
 
-The checkpoint requires Hugging Face access approval. Download it and point
-`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
-subdirectories:
+The checkpoint requires Hugging Face access approval. Authenticate once;
+`vllm serve` downloads the required components automatically:
 
 ```bash
 hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+export MODEL=MiniMaxAI/MiniMax-H3
 ```
 
+The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
+`Ref2VA/transformer/**`. It does not download or load the diffusers-format
+`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
+duplicate Ref2VA copies of shared components.
+
 Install vLLM-Omni from the checkout containing MiniMax H3 support. The
 two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
 FlashAttention-4. Install the optional dependency only for the four-GPU
@@ -41,6 +48,13 @@ B300/GB200 `FLASH_ATTN` profile:
 ```bash
 uv venv
 source .venv/bin/activate
+uv pip install -e .
+```
+
+To keep FA4 available as an explicit option on Blackwell, install the
+FlashAttention-4 extra:
+
+```bash
 uv pip install -e '.[fa4]'
 ```
 
@@ -49,16 +63,17 @@ reference-video preparation and MP4 output.
 
 ## Start a server
 
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
+Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery
+and shared components, and loads the second DiT from `Ref2VA/transformer`.
 
 ### Memory and storage requirements
 
 Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
 H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
-BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
-therefore needs roughly **270 GiB** of model storage, although the examples
-serve only one partition at a time.
+BF16 safetensors (about **135 GiB** on disk). Keeping both partitions
+locally therefore needs roughly **270 GiB** of model storage. A combined
+service downloads both; `--task-type fl2va` or `--task-type ref2va` downloads
+only the selected partition.
 
 CPU offload and distributed layerwise offload reduce GPU residency; they do
 not make the model weights disappear. With `--dlo-no-use-allgather`, each
@@ -79,7 +94,7 @@ This matches the accuracy-qualified reference path and prevents the Qwen3-VL
 encoder and DiT from being resident on the GPU at the same time.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0 \
@@ -96,8 +111,8 @@ vllm serve "${MODEL}" \
 ```
 
 Use a GPU with enough memory for the active H3 component and enough system RAM
-for the offloaded components. CPU offload reduces GPU memory pressure but adds
-PCIe/NVLink transfer latency.
+for both offloaded DiTs plus the shared components. Model-level offload keeps
+the two DiTs mutually exclusive on GPU, but adds PCIe/NVLink transfer latency.
 
 ### Two 24/32 GB GPUs: TP2 distributed layerwise offload
 
@@ -109,7 +124,7 @@ copied to the GPUs once per denoise stage, reused by every sampling step, and
 released before VAE decode so the decoder can reuse their HBM.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1 \
@@ -120,6 +135,7 @@ vllm serve "${MODEL}" \
   --host 0.0.0.0 \
   --port "${PORT}" \
   --trust-remote-code \
+  --task-type fl2va \
   --num-gpus 2 \
   --tensor-parallel-size 2 \
   --usp 1 \
@@ -191,23 +207,24 @@ bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
 The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
 `PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
 
-### Four GPUs: measured best-practice throughput
+### Four GPUs: throughput-oriented combined service
 
-The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
+For a combined service on four high-memory GPUs, use:
 
 - no CPU or layerwise offload;
 - Ulysses sequence parallelism degree 4;
 - native tiled VAE patch parallelism degree 4;
 - regional `torch.compile` for the repeated DiT blocks;
-- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP
-  left at 1.
+- dense BF16 `TRTLLM_ATTN`, with Ring and TP left at 1.
+
+Both DiTs remain resident in this no-offload configuration. If they do not fit,
+use model-level CPU offload.
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
@@ -220,19 +237,76 @@ vllm serve "${MODEL}" \
   --ring 1 \
   --vae-patch-parallel-size 4 \
   --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
+  --vae-use-tiling
 ```
 
-On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
-`Using CuTe FlashAttention-4 on Blackwell` before recording measurements.
-
 Do not add `--enforce-eager` to this performance configuration. The first
 request includes regional compilation; warm the server once before measuring
 steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
 1. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
+### Attention Backends
+
+On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16
+`TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly,
+use:
+
+```bash
+--diffusion-attention-backend TRTLLM_ATTN
+```
+
+Stable measurements with the four-GPU profile above put dense `TRTLLM_ATTN`
+and FA4 within 2% of each other. `TRTLLM_ATTN` remains the datacenter Blackwell
+default and enables the optional optimizations below. Confirm the server log
+contains `Defaulting to diffusion attention backend TRTLLM_ATTN` before
+recording measurements when using the default selection.
+
+FA4 remains available by explicitly selecting the `FLASH_ATTN` backend:
+
+```bash
+--diffusion-attention-backend FLASH_ATTN
+```
+
+On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
+`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements.
+
+`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main
+DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse
+Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses
+`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V
+to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short
+attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled
+for it. The example enables the calibration-free Skip-Softmax path with
+`threshold=0.05`, after the normalized timestep reaches `0.97`:
+
+```bash
+--diffusion-attention-config '{
+  "default": {
+    "backend": "TRTLLM_ATTN",
+    "quant": {
+      "dtype_qk": "fp8_e4m3",
+      "q_block_size": 1,
+      "k_block_size": 16
+    },
+    "skip_softmax": {
+      "threshold": 0.05,
+      "disabled_until_timestep": 0.97
+    }
+  },
+  "per_role": {
+    "minimax_h3.token_refiner": {
+      "backend": "TRTLLM_ATTN"
+    }
+  }
+}'
+```
+
+For configuration details, see
+[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax)
+and
+[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization).
+
 ### Text encoder tensor parallelism
 
 The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
@@ -243,11 +317,10 @@ vLLM-style tensor-parallel layers and runs with distributed collectives over
 its own encoder process group):
 
 ```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
+export MODEL=MiniMaxAI/MiniMax-H3
 export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
@@ -261,8 +334,7 @@ vllm serve "${MODEL}" \
   --text-encoder-tp-size 4 \
   --vae-patch-parallel-size 4 \
   --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
+  --vae-use-tiling
 ```
 
 `N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
@@ -275,12 +347,9 @@ the reference path within bf16 rounding: every encoder rank all-reduces the
 row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
 replicated on every rank.
 
-To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
-command with:
-
-```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
-```
+No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
+`task=ref2va` routes to
+`Ref2VA/transformer`. T2VA uses the FL2VA DiT.
 
 ### Online FP8 quantization
 
@@ -336,7 +405,7 @@ export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
 
 ### 1. T2VA: text to video and audio
 
-Run this request against the `FL2VA` partition:
+Run this request against the combined service:
 
 ```bash
 curl -sS -X POST "${API_URL}" \
@@ -354,7 +423,7 @@ curl -sS -X POST "${API_URL}" \
 
 ### 2. FL2VA: first frame to video and audio
 
-Run this request against the `FL2VA` partition. When width and height are
+Run this request against the combined service. When width and height are
 omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
 edge.
 
@@ -393,10 +462,11 @@ curl -sS -X POST "${API_URL}" \
 
 ### 3. Ref2VA: image-only, image/audio, or mixed references
 
-Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits
-`audio_reference`; adding one or more audio references is optional. The typed
-fields accept one object or an ordered JSON list. `audio_reference` accepts an
-HTTP(S) URL or a `data:` URL.
+Run these requests against the combined service or a Ref2VA-only service.
+Image-only Ref2VA omits `audio_reference`; adding one or more audio references
+is optional. The typed fields accept one object or an ordered JSON list.
+`audio_reference` accepts an HTTP(S) URL or a `data:` URL. In one terminal,
+expose the local reference assets to the serving host:
 
 ```bash
 python -m http.server 8092 \
@@ -446,7 +516,7 @@ shorter, the reference soundtrack is truncated to the generated clip.
 
 ### 4. Ref2VA: video, separate audio, and mixed references
 
-Run this request against the `Ref2VA` partition. Repeat the
+Run this request against the combined service. Repeat the
 `input_references` multipart field once per source video. H3 consumes the
 videos in form order and preserves their original soundtracks during
 conditioning.
@@ -508,7 +578,7 @@ than one is requested.
 
 | Parameter | Recommended value | Notes |
 |-----------|-------------------|-------|
-| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
+| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
 | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
 | `fps` | `24` | H3 output FPS is fixed |
 | `num_inference_steps` | `50` | Matches the reference accuracy workloads |
@@ -550,7 +620,8 @@ reduction.
 
 ## Known limitations
 
-- Each server process loads only one checkpoint partition.
+- 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.
 - The first regional-compile request is a warmup and should not be included in
   steady-state performance measurements.

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -56,10 +56,14 @@ vllm serve /path/to/Cosmos3-Super \
   --data-parallel-size 4
 ```
 
-### Without AllGather (full weights per rank, no sharding)
+### Without DLO AllGather (rank-local weights, no DLO sharding)
 
-Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
-no AllGather synchronization, no concurrent request requirement, but N× host memory.
+Use `--dlo-no-use-allgather` when you want each rank to stream the standard
+loader's rank-local weights independently — no DLO AllGather synchronization,
+no concurrent request requirement, but no DLO weight sharding and higher host
+memory use. In pure DP, this is equivalent to independent full-weight
+layerwise-offload replicas; with TP or HSDP, each rank retains the standard
+loader's local tensor layout.
 
 ```bash
 vllm serve /path/to/Cosmos3-Nano \
@@ -154,9 +158,9 @@ wait
 | Flag | Description | Default |
 |------|-------------|---------|
 | `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
-| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
+| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` |
 | `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
-| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
+| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` |
 | `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
 
 ## Mode comparison
@@ -166,7 +170,7 @@ wait
 | DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
 | DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
 | DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
-| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
+| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU |
 | HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
 
 ## Memory expectations
@@ -185,7 +189,11 @@ wait
 
 - `--dlo-use-allgather` (default) requires all concurrent requests to have the same
   `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
-- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
-  `load_weights()` automatically
-- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
-- HSDP + AllGather is rejected (would double-shard weights)
+- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
+  use `--dlo-no-use-allgather` or disable online quantization
+- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
+  `--dlo-no-use-allgather` retains the standard TP loader and should be treated
+  as an experimental compatibility path
+- HSDP + AllGather is rejected (would double-shard weights). HSDP with
+  `--dlo-no-use-allgather` is accepted at configuration level, but is not yet a
+  broadly validated production combination

tests/config/test_omni_config.py

@@ -39,7 +39,7 @@ from vllm_omni.config.stage_config import (
     load_deploy_config,
     merge_pipeline_deploy,
 )
-from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -124,10 +124,16 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
         assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
         assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
         assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
-        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90)
-        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128)
+        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization")
+        assert omni_stage.cache_config.enable_prefix_caching == engine_args.get("enable_prefix_caching")
+        expected_disable_hybrid = engine_args.get("disable_hybrid_kv_cache_manager")
+        if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_GENERATION:
+            expected_disable_hybrid = True if expected_disable_hybrid is None else expected_disable_hybrid
+        assert omni_stage.cache_config.disable_hybrid_kv_cache_manager == expected_disable_hybrid
+        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs")
         assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens")
-        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True)
+        assert omni_stage.scheduler_config.enable_chunked_prefill == engine_args.get("enable_chunked_prefill")
+        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling")
         legacy_parallel_config = engine_args.get("parallel_config") or {}
         assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get(
             "tensor_parallel_size",
@@ -198,6 +204,64 @@ def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime
     assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size
 
 
+@pytest.mark.parametrize(
+    "cli_overrides",
+    [
+        {"enable_lora": True},
+        {"stage_0_enable_lora": True},
+    ],
+    ids=["global", "stage-scoped"],
+)
+def test_from_pipeline_config_rejects_explicit_unowned_engine_cli_fields(cli_overrides):
+    with pytest.raises(ValueError, match=r"no structured config owner: enable_lora"):
+        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+
+
+@pytest.mark.parametrize(
+    ("engine_extras", "unowned_field"),
+    [
+        ({"enable_lora": True}, "enable_lora"),
+        ({"parallel_config": {"cfg_parallel_size": 2}}, "parallel_config.cfg_parallel_size"),
+    ],
+    ids=["top-level", "nested-parallel-config"],
+)
+def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras, unowned_field):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, engine_extras=engine_extras)])
+
+    with pytest.raises(ValueError, match=rf"no structured config owner: {unowned_field}"):
+        VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
+
+
+@pytest.mark.parametrize(
+    ("cli_overrides", "stage_id"),
+    [
+        ({"kv_cache_dtype": "fp8"}, 0),
+        ({"stage_0_kv_cache_dtype": "fp8"}, 0),
+        ({"stage_1_kv_cache_dtype": "fp8"}, 1),
+    ],
+    ids=["global", "ar-stage", "generation-stage"],
+)
+def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
+    with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
+        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+
+
+def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
+    stage = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path="dreamzero_tp1_cfg2",
+        cli_overrides={"kv_cache_dtype": "fp8"},
+    ).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8"
+
+
+def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged():
+    assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True}
+
+
 def test_runtime_num_gpus_is_derived_from_parallel_world_size():
     omni_config = _from_pipeline_key("hunyuan_image3_dit")
     stage = omni_config.stage_by_id(0)
@@ -380,8 +444,10 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
     assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields
 
 
-def test_runtime_config_fields_match_rfc_runtime_scope():
+def test_runtime_config_fields_match_structured_runtime_scope():
     assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
+        "distributed_executor_backend",
+        "worker_cls",
         "devices",
         "num_replicas",
         "env",
@@ -392,8 +458,21 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
     }
 
 
-def test_sub_config_fields_match_rfc_scopes():
+def test_sub_config_fields_match_structured_scopes():
     assert {f.name for f in fields(OmniStageModelConfig)} == {
+        "model",
+        "model_arch",
+        "revision",
+        "tokenizer_revision",
+        "code_revision",
+        "seed",
+        "logits_processors",
+        "trust_remote_code",
+        "dtype",
+        "attention_backend",
+        "moe_backend",
+        "hf_overrides",
+        "limit_mm_per_prompt",
         "active_stream_window",
         "duplex_max_sessions",
         "enable_sleep_mode",
@@ -404,6 +483,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "task_type",
         "codec_frame_rate_hz",
         "enforce_eager",
+        "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",
         "compilation_config",
         "enable_multithread_weight_load",
@@ -416,12 +496,16 @@ def test_sub_config_fields_match_rfc_scopes():
         "tokenizer_subdir",
     }
     assert {f.name for f in fields(OmniStageLoadConfig)} == {
+        "tokenizer",
+        "download_dir",
+        "skip_tokenizer_init",
         "load_format",
         "tokenizer_mode",
         "config_format",
         "skip_mm_profiling",
     }
     assert {f.name for f in fields(OmniStageCacheConfig)} == {
+        "kv_cache_memory_bytes",
         "gpu_memory_utilization",
         "enable_prefix_caching",
         "disable_hybrid_kv_cache_manager",
@@ -435,6 +519,8 @@ def test_sub_config_fields_match_rfc_scopes():
         "async_scheduling",
     }
     assert {f.name for f in fields(OmniStageConnectorConfig)} == {
+        "async_chunk",
+        "omni_kv_config",
         "stage_connector",
         "output_connectors",
         "input_connectors",
@@ -827,7 +913,7 @@ def test_from_pipeline_config_matches_build_engine_args_dict_behavior_for_repres
     pipeline = _resolve_pipeline_or_skip("qwen3_tts")
     legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
     omega_stage = legacy_stage.to_omegaconf()
-    legacy_engine_args = build_engine_args_dict(
+    legacy_engine_args = build_legacy_engine_args_dict(
         omega_stage,
         model="/tmp/qwen3-tts",
         stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}},
@@ -922,6 +1008,8 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
 
 
 def test_diffusion_config_field_classification_covers_current_fields():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
     classified_fields = (
         omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
         | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
@@ -929,8 +1017,12 @@ def test_diffusion_config_field_classification_covers_current_fields():
     )
 
     assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
+    assert {f.name for f in fields(OmniDiffusionConfig)} <= (
+        classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+    )
     assert {
         "enable_prompt_embed_cache",
+        "enable_session_state_manager",
         "prompt_embed_cache_size",
         "diffusion_kv_cache_dtype",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS

tests/dfx/conftest.py

@@ -1,4 +1,3 @@
-import copy
 import json
 import os
 import re
@@ -10,7 +9,7 @@ from typing import Any
 
 import pytest
 
-from tests.helpers.mark import hardware_marks
+from tests.helpers.mark import get_hardware_mark_list, hardware_marks
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
 from vllm_omni.platforms import current_omni_platform
@@ -447,36 +446,16 @@ def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> li
     return unique
 
 
-def extract_mark_resource_label(mark_field: Any) -> str:
-    """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values.
-
-    Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``.
-
-    Prefer :func:`get_runtime_resource_label` for perf result filenames so labels
-    reflect the machine that actually ran the benchmark.
-    """
-    if isinstance(mark_field, list):
-        for item in mark_field:
-            if isinstance(item, dict) and "hardware_marks" in item:
-                return extract_mark_resource_label(item)
-        return "na"
-    if not isinstance(mark_field, dict):
-        return "na"
-    hw = mark_field.get("hardware_marks")
-    if not isinstance(hw, dict):
-        return "na"
-    res = hw.get("res")
-    if not isinstance(res, dict) or not res:
-        return "na"
-    labels = [_safe_filename_token(value) for value in res.values()]
-    return "-".join(labels) if labels else "na"
-
-
-_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = (
+# Device marketing-name tokens for ``get_runtime_resource_label`` only.
+# Longer / more specific tokens must appear before shorter substrings
+# (e.g. ``H200`` before ``H20``, ``910B4`` before ``910``).
+_RUNTIME_DEVICE_ALIASES: tuple[str, ...] = (
     "H100",
     "H800",
     "H200",
     "H20",
+    "B200",
+    "GB200",
     "L40S",
     "L40",
     "L4",
@@ -505,7 +484,7 @@ def _normalize_runtime_device_label(raw: str) -> str:
     if not raw or not str(raw).strip():
         return "na"
     upper = str(raw).upper()
-    for token in _KNOWN_RUNTIME_RESOURCE_TOKENS:
+    for token in _RUNTIME_DEVICE_ALIASES:
         if token.upper() in upper:
             return _safe_filename_token(token)
     compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw))
@@ -558,10 +537,108 @@ def resource_label_for_filename(resource_label: str | None) -> str:
     return token
 
 
-def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
-    """Return runtime hardware label for perf result filenames."""
-    del configs
-    return get_runtime_resource_label()
+def is_hardware_nested_baseline(baseline: dict[str, Any]) -> bool:
+    """True for ``{"H100": {"metric": ...}, "A3": {...}}`` baselines.
+
+    Detection is structural + hardware-label allowlist (not metric-name hints):
+
+    - every top-level value is a non-empty dict (per-hardware metric map);
+    - every top-level key is in :func:`tests.helpers.mark.get_hardware_mark_list`
+      (``[hardware-resource]`` markers from ``pyproject.toml``).
+
+    Metric names under each hardware bucket are unconstrained.
+    Runtime device aliases (:data:`_RUNTIME_DEVICE_ALIASES`) are unrelated.
+    """
+    if not baseline:
+        return False
+    if not all(isinstance(value, dict) and bool(value) for value in baseline.values()):
+        return False
+    allowed = {label.upper() for label in get_hardware_mark_list()}
+    return all((token := str(key).strip()) and token.upper() in allowed for key in baseline.keys())
+
+
+def resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+) -> Any:
+    """Pick the baseline threshold for one metric at this sweep step.
+
+    Accepts a sweep-aligned list/tuple (indexed by ``sweep_index``) or a scalar.
+    """
+    if baseline_raw is None:
+        return None
+    if isinstance(baseline_raw, dict):
+        raise TypeError(
+            "per-metric baseline dict keyed by concurrency/request-rate is not supported; "
+            "use a sweep-aligned list or a scalar under each hardware bucket"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        if sweep_index is None:
+            raise ValueError("list baseline requires sweep_index")
+        if not (0 <= sweep_index < len(baseline_raw)):
+            raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def resolve_baseline_for_sweep(
+    baseline_config: dict[str, Any] | None,
+    *,
+    sweep_index: int | None = None,
+) -> dict[str, Any]:
+    """Resolve sweep-dependent baselines while keeping every hardware bucket.
+
+    Expects hardware-nested input only::
+
+        {"H100": {"throughput_qps": [a, b, c]}, "A3": {"throughput_qps": [x, y, z]}}
+
+    with ``sweep_index=2`` (e.g. concurrency 32) becomes::
+
+        {"H100": {"throughput_qps": c}, "A3": {"throughput_qps": z}}
+
+    Per-metric values may be a sweep-aligned list or a scalar.
+    """
+    if not baseline_config:
+        return {}
+    if not is_hardware_nested_baseline(baseline_config):
+        allowed = sorted(get_hardware_mark_list())
+        unknown = [
+            str(key).strip()
+            for key in baseline_config.keys()
+            if str(key).strip().upper() not in {label.upper() for label in allowed}
+        ]
+        nested_shape = all(isinstance(value, dict) and bool(value) for value in baseline_config.values())
+        if nested_shape and unknown:
+            raise ValueError(
+                "baseline top-level keys must be [hardware-resource] pytest markers from "
+                "pyproject.toml (tool.pytest.ini_options.markers). "
+                f"Unknown hardware label(s): {unknown!r}. "
+                f"Known markers: {allowed!r}. "
+                "Add a marker such as '\"A100: [hardware-resource] Tests that require A100 GPU\"' "
+                "to pyproject.toml, then retry."
+            )
+        raise ValueError(
+            "baseline must be hardware-nested with [hardware-resource] marker names as "
+            "top-level keys, e.g. "
+            '{"H100": {"throughput_qps": [...]}, "A3": {"custom_metric": 1.0}}. '
+            f"got top-level keys={list(baseline_config.keys())!r}; "
+            f"known markers={allowed!r}. "
+            "If you need a new hardware bucket, add it to pyproject.toml markers with "
+            "the [hardware-resource] tag."
+        )
+
+    return {
+        hardware: {
+            metric_name: resolve_baseline_value(
+                baseline_raw,
+                sweep_index=sweep_index,
+            )
+            for metric_name, baseline_raw in metrics.items()
+        }
+        for hardware, metrics in baseline_config.items()
+        if isinstance(metrics, dict)
+    }
 
 
 def run_benchmark(
@@ -572,17 +649,18 @@ def run_benchmark(
     num_prompt: int,
     *,
     baseline_config: dict[str, Any] | None = None,
+    sweep_index: int | None = None,
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
     resource_label: str | None = None,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
-    After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the raw
-    ``baseline`` object from the test config (when ``baseline_config`` is provided),
-    including any hardware-nested map (e.g. ``{"H100": {...}}``). Downstream report
-    tooling can select a hardware bucket when needed. If the benchmark exits without
-    writing a result file, ``result_omni_template.json`` is used as a fallback.
+    After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the
+    sweep-resolved baseline (when ``baseline_config`` is provided). Hardware-nested
+    maps keep every hardware bucket, but each metric is reduced to the value for
+    this concurrency / request-rate step. If the benchmark exits without writing a
+    result file, ``result_omni_template.json`` is used as a fallback.
     """
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     ri = _safe_filename_token(random_input_len)
@@ -649,8 +727,11 @@ def run_benchmark(
             result = json.load(f)
 
     if baseline_config:
-        # Persist the JSON baseline as-is (no hardware / sweep resolution).
-        result["baseline"] = copy.deepcopy(baseline_config)
+        # Keep every hardware bucket; resolve list metrics to this sweep step.
+        result["baseline"] = resolve_baseline_for_sweep(
+            baseline_config,
+            sweep_index=sweep_index,
+        )
     else:
         result["baseline"] = {}
     if random_input_len is not None:

tests/dfx/perf/scripts/diffusion_result_template.json

@@ -21,10 +21,12 @@
       "num-input-images": 0,
       "enable-negative-prompt": false,
       "baseline": {
-        "throughput_qps": 0,
-        "latency_mean": 0,
-        "peak_memory_mb_max": 0,
-        "peak_memory_mb_mean": 0
+        "H100": {
+          "throughput_qps": 0,
+          "latency_mean": 0,
+          "peak_memory_mb_max": 0,
+          "peak_memory_mb_mean": 0
+        }
       }
     },
     "result": {

tests/dfx/perf/scripts/result_omni_template.json

@@ -46,9 +46,11 @@
   "median_audio_duration_s": 0,
   "p99_audio_duration_s": 0,
   "baseline": {
-    "mean_ttft_ms": 0,
-    "mean_audio_ttfp_ms": 0,
-    "mean_audio_rtf": 0
+    "H100": {
+      "mean_ttft_ms": 0,
+      "mean_audio_ttfp_ms": 0,
+      "mean_audio_rtf": 0
+    }
   },
   "random_input_len": 0,
   "random_output_len": 0

tests/dfx/perf/scripts/run_benchmark.py

@@ -237,7 +237,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
         break
 
     # QPS / request-rate sweep
-    for qps, num_prompt in zip(qps_list, num_prompt_list):
+    for sweep_index, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
             args=args,
@@ -246,6 +246,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
             baseline_config=params.get("baseline"),
+            sweep_index=sweep_index,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,
@@ -253,7 +254,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
         assert_result(result, num_prompt)
 
     # concurrency test
-    for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
+    for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
             args=args,
@@ -262,6 +263,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
             baseline_config=params.get("baseline"),
+            sweep_index=sweep_index,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -23,7 +23,6 @@ of all runs from cases in that file). Bulk load without ``--test-config-file`` u
 the same per-file aggregation; ``-m`` only selects which cases run.
 """
 
-import copy
 import json
 import os
 import socket
@@ -45,6 +44,7 @@ from tests.dfx.conftest import (
     get_runtime_resource_label,
     hardware_json_value,
     is_diffusion_perf_config,
+    resolve_baseline_for_sweep,
     resolve_pytest_marks,
     resource_label_for_filename,
 )
@@ -350,10 +350,49 @@ def _kill_process_tree(pid: int) -> None:
 
 
 def _resolve_offline_model(model: str) -> str:
-    """Under HF_HUB_OFFLINE, resolve a HF repo id to its local snapshot dir."""
+    """Resolve a model string to a local directory.
+
+    Local paths pass through unchanged. HF repo ids are resolved to their local
+    snapshot dir when ``HF_HUB_OFFLINE`` is set. ``repo_id/subfolder`` (e.g.
+    ``MiniMaxAI/MiniMax-H3/FL2VA``) is resolved to the matching partition
+    subfolder of the snapshot, downloading only that subfolder when online.
+    """
     import huggingface_hub
 
-    if not model or os.path.isdir(model) or not huggingface_hub.constants.HF_HUB_OFFLINE:
+    if not model or os.path.isdir(model):
+        return model
+
+    # MiniMax-H3 custom code uses relative imports across files, which breaks
+    # when the HF cache snapshot is symlinked (``get_class_from_dynamic_module``
+    # resolves ``realpath`` into the blobs dir). Mirror the accuracy test's env
+    # overrides so CI / local runs can point at a materialized partition. The
+    # repo-root key (used with ``--task-type fl2va|ref2va``) expects a
+    # materialized root containing the FL2VA/Ref2VA subfolders.
+    model_env_overrides = {
+        "MiniMaxAI/MiniMax-H3": "VLLM_TEST_MINIMAX_H3_MODEL",
+        "MiniMaxAI/MiniMax-H3/FL2VA": "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL",
+        "MiniMaxAI/MiniMax-H3/Ref2VA": "VLLM_TEST_MINIMAX_H3_REF2VA_MODEL",
+    }
+    env_name = model_env_overrides.get(model)
+    if env_name:
+        env_model = os.environ.get(env_name)
+        if env_model:
+            return env_model
+
+    parts = model.split("/")
+    if len(parts) >= 3:
+        repo_id = "/".join(parts[:2])
+        subfolder = "/".join(parts[2:])
+        from huggingface_hub import snapshot_download
+
+        snapshot_root = snapshot_download(
+            repo_id,
+            allow_patterns=[f"{subfolder}/**"],
+            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+        )
+        return str(Path(snapshot_root) / subfolder)
+
+    if not huggingface_hub.constants.HF_HUB_OFFLINE:
         return model
     from huggingface_hub import snapshot_download
 
@@ -387,6 +426,11 @@ class DiffusionServer:
     def _start_server(self) -> None:
         env = os.environ.copy()
         env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        # Give each run its own writable video-storage dir. The server default
+        # (/tmp/storage) may be missing or not writable on shared hosts / CI.
+        storage_path = Path(os.environ.get("VLLM_OMNI_STORAGE_PATH", str(BENCHMARK_RESULT_DIR / "storage")))
+        storage_path.mkdir(parents=True, exist_ok=True)
+        env["VLLM_OMNI_STORAGE_PATH"] = str(storage_path)
 
         cmd = [
             sys.executable,
@@ -766,8 +810,8 @@ def run_benchmark(
     completed = metrics.get("completed_requests", metrics.get("completed", 0))
     failed = metrics.get("failed_requests", metrics.get("failed", 0))
 
-    # Persist the JSON baseline as-is (no hardware / sweep resolution).
-    baseline = copy.deepcopy(params.get("baseline") or {})
+    # Persist sweep-resolved baseline from params (already narrowed in _build_run_params).
+    baseline = params.get("baseline") or {}
     metrics["baseline"] = baseline
 
     record: dict[str, Any] = {
@@ -862,6 +906,7 @@ def _build_run_params(
     params: dict[str, Any],
     *,
     num_prompts: int,
+    sweep_index: int | None = None,
     request_rate: Any | None = None,
     max_concurrency: Any | None = None,
 ) -> dict[str, Any]:
@@ -873,6 +918,12 @@ def _build_run_params(
         run_params["request-rate"] = request_rate
     if max_concurrency is not None:
         run_params["max-concurrency"] = max_concurrency
+    if "baseline" in params:
+        # Keep all hardware buckets; pick the metric value for this sweep step only.
+        run_params["baseline"] = resolve_baseline_for_sweep(
+            params.get("baseline"),
+            sweep_index=sweep_index,
+        )
     return run_params
 
 
@@ -895,19 +946,20 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
 
     sweep_runs: list[dict[str, Any]] = []
 
-    for request_rate, num_prompts in zip(request_rate_list, num_prompt_list):
+    for sweep_index, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
         sweep_runs.append(
             {
                 "params": _build_run_params(
                     params,
                     request_rate=request_rate,
                     num_prompts=num_prompts,
+                    sweep_index=sweep_index,
                 ),
                 "num_prompts": num_prompts,
             }
         )
 
-    for max_concurrency, num_prompts in zip(max_concurrency_list, num_prompt_list):
+    for sweep_index, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         sweep_runs.append(
             {
                 "params": _build_run_params(
@@ -915,6 +967,7 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
                     max_concurrency=max_concurrency,
                     num_prompts=num_prompts,
                     request_rate="inf",
+                    sweep_index=sweep_index,
                 ),
                 "num_prompts": num_prompts,
             }
@@ -927,6 +980,7 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
                 "params": _build_run_params(
                     params,
                     num_prompts=default_num_prompts,
+                    sweep_index=0,
                 ),
                 "num_prompts": default_num_prompts,
             }

tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json

@@ -25,8 +25,8 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "H100": {
-                        "throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
-                        "latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
+                        "throughput_qps": [0.0303, 0.0302, 0.0302],
+                        "latency_mean": [32.6, 62.2, 111.7],
                         "peak_memory_mb_max": 40500,
                         "peak_memory_mb_mean": 40500
                     }

tests/dfx/perf/tests/test_boogu_image_vllm_omni.json

@@ -25,8 +25,8 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "H100": {
-                        "throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
-                        "latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
+                        "throughput_qps": [0.0549, 0.0547, 0.0549],
+                        "latency_mean": [18.0, 34.3, 61.4],
                         "peak_memory_mb_max": 40500,
                         "peak_memory_mb_mean": 40500
                     }

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -134,7 +134,7 @@
         "max-concurrency": 16,
         "num-prompts": 96,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.1874,
             "latency_p99": 101.852,
             "latency_mean": 79.0385,

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -38,7 +38,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.21,
             "latency_p99": 4.7469,
             "peak_memory_mb_max": 101100,

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -38,7 +38,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.20,
             "latency_p99": 5.1025,
             "peak_memory_mb_max": 97402,

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -37,7 +37,7 @@
         "num-prompts": 10,
         "max-concurrency": 1,
         "baseline": {
-          "H100": {
+          "H200": {
             "throughput_qps": 0.213,
             "latency_mean": 4.6953,
             "latency_p99": 4.8601,

tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json

@@ -0,0 +1,295 @@
+[
+  {
+    "test_name": "test_minimax_h3_t2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 T2VA text-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0211,
+            "latency_mean": 38.3252,
+            "peak_memory_mb_mean": 54219.6
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_ti2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 FL2VA first-frame image-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "ti2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Continue the scene: the camera slowly pushes in on the subject as the background city lights drift and blur, with the actor singing softly and emotionally.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"fl2va\",\"duration\":8.7,\"audio_flow_shift\":3.0}"
+        },
+        "num-input-images": 1,
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0212,
+            "latency_mean": 38.1385,
+            "peak_memory_mb_mean": 54622.8
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_v2v_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 Ref2VA video-reference-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "ref2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "v2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Restyle the reference footage into a cinematic night scene with neon reflections, keeping the same camera motion and subject.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"ref2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0077,
+            "latency_mean": 105.7508,
+            "peak_memory_mb_mean": 55180.8
+          }
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_minimax_h3_t2v_dlo_4gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "MiniMax-H3 T2VA with 4-card distributed layerwise offload (DLO), 8 denoise steps (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "MiniMaxAI/MiniMax-H3",
+      "serve_args": {
+        "trust-remote-code": true,
+        "task-type": "fl2va",
+        "num-gpus": 4,
+        "usp": 4,
+        "ring": 1,
+        "use-hsdp": true,
+        "hsdp-shard-size": 4,
+        "text-encoder-tp-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-parallel-mode": "tile",
+        "vae-use-tiling": true,
+        "diffusion-attention-backend": "FLASH_ATTN",
+        "stage-init-timeout": 1800,
+        "init-timeout": 1800,
+        "enable-diffusion-pipeline-profiler": true,
+        "enable-distributed-layerwise-offload": true,
+        "dlo-no-use-allgather": true,
+        "dlo-resident-layers": 8,
+        "enforce-eager": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1344x768_frames209_steps8",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "warmup-requests": 1,
+        "warmup-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1344,
+            "height": 768,
+            "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.",
+            "num_inference_steps": 8,
+            "num_frames": 209,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "flow_shift": 12.0,
+          "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}"
+        },
+        "baseline": {
+          "H100": {
+            "throughput_qps": 0.0212,
+            "latency_mean": 38.2089,
+            "peak_memory_mb_mean": 36218.4
+          }
+        }
+      }
+    ]
+  }
+]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -17,7 +17,7 @@ def test_task_excluded_from_cli_args():
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "percentile-metrics": "audio_rtf,audio_ttfp",
-        "baseline": {"mean_audio_rtf": [0.5]},
+        "baseline": {"H100": {"mean_audio_rtf": [0.5]}},
     }
     exclude_keys = {
         "request_rate",
@@ -123,35 +123,6 @@ def test_resolve_pytest_marks_rejects_legacy_object_format():
         )
 
 
-def test_extract_mark_resource_label():
-    from tests.dfx.conftest import extract_mark_resource_label
-
-    assert extract_mark_resource_label(None) == "na"
-    assert (
-        extract_mark_resource_label(
-            [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "full_model",
-            ]
-        )
-        == "H100"
-    )
-    assert (
-        extract_mark_resource_label(
-            [
-                {
-                    "hardware_marks": {
-                        "res": {"cuda": "H100", "rocm": "MI325"},
-                        "num_cards": 2,
-                    }
-                },
-                "diffusion",
-            ]
-        )
-        == "H100-MI325"
-    )
-
-
 def test_resource_label_for_filename():
     from tests.dfx.conftest import resource_label_for_filename
 
@@ -169,16 +140,14 @@ def test_hardware_json_value():
     assert hardware_json_value(None) == ""
 
 
-def test_extract_configs_resource_label(monkeypatch):
-    from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label
+def test_get_runtime_resource_label(monkeypatch):
+    from tests.dfx.conftest import get_runtime_resource_label
 
     monkeypatch.setattr(
         "tests.dfx.conftest._read_runtime_device_name",
         lambda *, device_id=0: "NVIDIA H100 80GB HBM3",
     )
-    get_runtime_resource_label(refresh=True)
-    assert extract_configs_resource_label([]) == "H100"
-    get_runtime_resource_label(refresh=True)
+    assert get_runtime_resource_label(refresh=True) == "H100"
     monkeypatch.setattr(
         "tests.dfx.conftest._read_runtime_device_name",
         lambda *, device_id=0: "Ascend910B2",
@@ -336,3 +305,183 @@ def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_pat
         ("start", tts_server),
         ("stop", tts_server),
     ]
+
+
+def test_is_hardware_nested_baseline():
+    from tests.dfx.conftest import (
+        _RUNTIME_DEVICE_ALIASES,
+        is_hardware_nested_baseline,
+    )
+    from tests.helpers.mark import get_hardware_mark_list
+
+    hardware_marks = get_hardware_mark_list()
+    assert hardware_marks
+    assert {"H100", "L4", "A3", "MI325", "B200"} <= hardware_marks
+    assert "cuda" not in hardware_marks
+    assert "full_model" not in hardware_marks
+
+    # Runtime aliases are a full static list (independent of baseline allowlist).
+    assert "H100" in _RUNTIME_DEVICE_ALIASES
+    assert "A100" in _RUNTIME_DEVICE_ALIASES
+
+    assert is_hardware_nested_baseline(
+        {
+            "H100": {"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]},
+            "A3": {"mean_ttft_ms": [0.9, 1.8], "mean_e2el_ms": [9.0, 18.0]},
+        }
+    )
+    # Custom metric names under a known hardware label are allowed.
+    assert is_hardware_nested_baseline({"H100": {"custom_stage_ms": 12.0, "foo_bar": [1.0, 2.0]}})
+    assert is_hardware_nested_baseline({"A3": {"request_throughput": 0.5}})
+    # Flat metric maps / concurrency-keyed maps are not hardware-nested.
+    assert not is_hardware_nested_baseline({"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]})
+    assert not is_hardware_nested_baseline({"mean_ttft_ms": {"1": 1.0, "32": 2.0}})
+    assert not is_hardware_nested_baseline({})
+    # Alias-only labels (not [hardware-resource] markers) are rejected for baselines.
+    assert not is_hardware_nested_baseline({"A100": {"throughput_qps": 1.0}})
+    # Unknown top-level hardware label is rejected.
+    assert not is_hardware_nested_baseline({"UnknownGPU": {"throughput_qps": 1.0}})
+    # Empty per-hardware metric map is rejected.
+    assert not is_hardware_nested_baseline({"H100": {}})
+
+
+def test_resolve_baseline_for_sweep_keeps_all_hardware_for_one_concurrency():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    baseline = {
+        "H100": {
+            "mean_ttft_ms": [96.4, 140.8, 271.9, 362.3, 507.8],
+            "mean_e2el_ms": [18507.0, 28365.0, 31907.0, 48161.0, 72630.0],
+        },
+        "A3": {
+            "mean_ttft_ms": [90.0, 130.0, 250.0, 340.0, 480.0],
+            "mean_e2el_ms": [17000.0, 26000.0, 30000.0, 45000.0, 70000.0],
+        },
+    }
+    # max_concurrency=[1,4,8,16,32] -> index 4 is concurrency 32
+    got = resolve_baseline_for_sweep(baseline, sweep_index=4)
+    assert got == {
+        "H100": {"mean_ttft_ms": 507.8, "mean_e2el_ms": 72630.0},
+        "A3": {"mean_ttft_ms": 480.0, "mean_e2el_ms": 70000.0},
+    }
+    # First sweep step keeps both hardware buckets too.
+    got0 = resolve_baseline_for_sweep(baseline, sweep_index=0)
+    assert set(got0) == {"H100", "A3"}
+    assert got0["H100"]["mean_ttft_ms"] == 96.4
+    assert got0["A3"]["mean_ttft_ms"] == 90.0
+
+
+def test_resolve_baseline_for_sweep_rejects_flat_baseline():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    with pytest.raises(ValueError, match=r"hardware-nested.*pyproject\.toml"):
+        resolve_baseline_for_sweep(
+            {"throughput_qps": [0.4, 0.6], "latency_mean": [1.0, 2.0]},
+            sweep_index=1,
+        )
+
+
+def test_resolve_baseline_for_sweep_rejects_unknown_hardware_label():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    with pytest.raises(ValueError, match=r"Unknown hardware label\(s\): \['A100'\].*pyproject\.toml"):
+        resolve_baseline_for_sweep(
+            {"A100": {"throughput_qps": 1.0}},
+            sweep_index=0,
+        )
+
+
+def test_resolve_baseline_for_sweep_supports_list_and_scalar_under_hardware():
+    from tests.dfx.conftest import resolve_baseline_for_sweep
+
+    # Sweep-aligned lists under each hardware bucket (canonical form).
+    listed = {
+        "H100": {"throughput_qps": [0.4, 0.6, 0.8], "latency_mean": [1.0, 2.0, 3.0]},
+        "A3": {"throughput_qps": [0.5, 0.7, 0.9], "latency_mean": [0.9, 1.8, 2.7]},
+    }
+    assert resolve_baseline_for_sweep(listed, sweep_index=1) == {
+        "H100": {"throughput_qps": 0.6, "latency_mean": 2.0},
+        "A3": {"throughput_qps": 0.7, "latency_mean": 1.8},
+    }
+
+    # Scalars under hardware stay as-is (single-concurrency cases).
+    scalar = {"H100": {"throughput_qps": 0.5}}
+    assert resolve_baseline_for_sweep(scalar, sweep_index=0) == {"H100": {"throughput_qps": 0.5}}
+    # Custom metric names are preserved.
+    custom = {"H100": {"custom_stage_ms": [10.0, 20.0]}}
+    assert resolve_baseline_for_sweep(custom, sweep_index=1) == {"H100": {"custom_stage_ms": 20.0}}
+    assert resolve_baseline_for_sweep(None) == {}
+    assert resolve_baseline_for_sweep({}) == {}
+
+
+def test_resolve_baseline_value_errors():
+    from tests.dfx.conftest import resolve_baseline_value
+
+    with pytest.raises(ValueError, match="sweep_index"):
+        resolve_baseline_value([1.0, 2.0], sweep_index=None)
+    with pytest.raises(IndexError):
+        resolve_baseline_value([1.0], sweep_index=1)
+    with pytest.raises(TypeError, match="not supported"):
+        resolve_baseline_value({"1": 0.4}, sweep_index=0)
+
+
+def test_diffusion_build_run_params_resolves_baseline_per_sweep(tmp_path, monkeypatch):
+    """``_build_run_params`` / ``_iter_sweep_runs`` narrow list baselines per concurrency."""
+    import importlib
+    import sys
+
+    cfg = tmp_path / "mini_diffusion_perf.json"
+    cfg.write_text(
+        json.dumps(
+            [
+                {
+                    "test_name": "test_mini",
+                    "server_type": "vllm-omni",
+                    "mark": [
+                        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                        "diffusion",
+                        "full_model",
+                    ],
+                    "server_params": {"model": "m/mini"},
+                    "benchmark_params": [{"name": "p0", "num-prompts": 1, "max-concurrency": 1}],
+                }
+            ]
+        ),
+        encoding="utf-8",
+    )
+    monkeypatch.setattr(sys, "argv", ["pytest", "--test-config-file", str(cfg)])
+    sys.modules.pop("tests.dfx.perf.scripts.run_diffusion_benchmark", None)
+    from tests.dfx.perf.scripts import run_diffusion_benchmark as rdb
+
+    importlib.reload(rdb)
+
+    params = {
+        "name": "c_sweep",
+        "dataset": "random",
+        "task": "t2i",
+        "num-prompts": [8, 16, 32],
+        "max-concurrency": [1, 8, 32],
+        "baseline": {
+            "H100": {"throughput_qps": [0.1, 0.2, 0.3], "latency_mean": [10.0, 20.0, 30.0]},
+            "A3": {"throughput_qps": [0.11, 0.22, 0.33], "latency_mean": [9.0, 19.0, 29.0]},
+        },
+    }
+    run32 = rdb._build_run_params(
+        params,
+        num_prompts=32,
+        max_concurrency=32,
+        request_rate="inf",
+        sweep_index=2,
+    )
+    assert run32["max-concurrency"] == 32
+    assert run32["baseline"] == {
+        "H100": {"throughput_qps": 0.3, "latency_mean": 30.0},
+        "A3": {"throughput_qps": 0.33, "latency_mean": 29.0},
+    }
+
+    sweeps = rdb._iter_sweep_runs(params)
+    assert len(sweeps) == 3
+    assert sweeps[2]["params"]["max-concurrency"] == 32
+    assert sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"] == 0.3
+    assert "A3" in sweeps[2]["params"]["baseline"]
+    assert not isinstance(sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"], list)

tests/diffusion/attention/test_attention_config.py

@@ -90,6 +90,35 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
+    def test_block_sparse_defaults_applied_when_backend_selected(self):
+        spec = AttentionSpec(backend="RAINFUSION_ATTN")
+        assert spec.block_sparse.sparsity == 0.8
+        assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
+
+    def test_block_sparse_skip_layers_selector_expanded(self):
+        spec = AttentionSpec(
+            backend="RAINFUSION_ATTN",
+            block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"},
+        )
+        assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38}
+        assert spec.backend_kwargs() == {
+            "sparsity": 0.9,
+            "start_step": 12,
+            "skip_layers": [0, 1, 2, 38],
+        }
+
+    def test_block_sparse_rejected_on_dense_backend(self):
+        with pytest.raises(ValueError, match="block_sparse is only supported by"):
+            AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8})
+
+    @pytest.mark.parametrize(
+        "block_sparse",
+        [{"sparsity": 1.5}, {"start_step": -1}],
+    )
+    def test_block_sparse_invalid_values(self, block_sparse):
+        with pytest.raises(ValueError):
+            AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):

tests/diffusion/attention/test_rainfusion_plan.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
+
+The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
+1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
+first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+mask from the video and prefix segments separately, so the two only agree on the
+block count when the video segment ends on a boundary.
+"""
+
+import dataclasses
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
+    _BLOCK_SIZE,
+    RainFusionAttentionImpl,
+    RainFusionPlan,
+)
+from vllm_omni.platforms import current_omni_platform
+
+PREFIX_ROWS = 710  # 14 text rows + 696 audio rows
+ALIGNED_GRID = (62, 24, 40)  # 1280x768 -> 59520 video rows, 465 blocks
+MISALIGNED_GRID = (62, 24, 42)  # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
+
+
+def make_impl(**backend_kwargs):
+    return RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8, **backend_kwargs},
+    )
+
+
+def make_metadata(grid, prefix_len=PREFIX_ROWS, max_seqlen_q=None):
+    video_rows = grid[0] * grid[1] * grid[2]
+    return AttentionMetadata(
+        extra={"max_seqlen_q": prefix_len + video_rows if max_seqlen_q is None else max_seqlen_q},
+        video_layout=VideoTokenLayout(prefix_len=prefix_len, latent_grid=grid),
+    )
+
+
+def test_block_aligned_video_segment_runs_sparse():
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 59520
+    assert plan.latent_shape == list(ALIGNED_GRID)
+
+
+def test_misaligned_video_segment_stays_dense():
+    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    # pad keys would take a share of every real query's softmax denominator.
+    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+
+
+@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
+def test_prefix_length_does_not_affect_alignment(prefix_len):
+    # Only the video segment has to land on a block boundary; rf_v2 pools the
+    # prefix separately and keeps every one of its blocks.
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID, prefix_len=prefix_len))
+
+    assert plan is not None
+    assert plan.prefix_len == prefix_len
+
+
+def test_sparsity_zero_never_resolves_a_plan():
+    assert make_impl(sparsity=0.0)._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+
+
+def test_missing_video_layout_falls_back_to_dense():
+    assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 60230})) is None
+
+
+def test_missing_max_seqlen_falls_back_to_dense():
+    metadata = AttentionMetadata(
+        video_layout=VideoTokenLayout(prefix_len=PREFIX_ROWS, latent_grid=ALIGNED_GRID),
+    )
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+def test_video_segment_must_be_the_tail_of_packed_document_zero():
+    metadata = make_metadata(ALIGNED_GRID, max_seqlen_q=PREFIX_ROWS + 59520 + 128)
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
+def test_short_video_stays_dense(grid):
+    assert make_impl()._resolve_plan(make_metadata(grid)) is None
+
+
+@pytest.mark.parametrize("qkv_layout", ["BNSD", "BSH"])
+def test_explicitly_wrong_layout_is_rejected(qkv_layout):
+    with pytest.raises(ValueError, match="BSND"):
+        RainFusionAttentionImpl(
+            num_heads=8,
+            head_size=128,
+            softmax_scale=128**-0.5,
+            qkv_layout=qkv_layout,
+            backend_kwargs={"sparsity": 0.8},
+        )
+
+
+def test_undeclared_layout_stays_dense():
+    # An undeclared layout is not an error -- the layer keeps working, just
+    # densely, and the fallback sees the same absent layout plain FLASH_ATTN would.
+    impl = RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        backend_kwargs={"sparsity": 0.8},
+    )
+
+    assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+    assert impl.dense_fallback.qkv_layout is None
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_fully_populated_mask_reproduces_dense_attention():
+    """At sparsity=0 every key block is kept, so rf_v2 must match dense attention.
+
+    This is the guard on the geometry the backend hands the kernel: a wrong
+    prefix length or latent grid still produces plausible output under sparse
+    selection, but shows up here as a mismatch against SDPA.
+    """
+    torch.manual_seed(0)
+    prefix_len, grid = 384, (4, 16, 64)  # 4096 video rows = 32 blocks
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    heads, head_dim = 4, 128
+    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
+
+    impl = make_impl()
+    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
+    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
+    out = impl._forward_sparse_npu(q, k, v, plan)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).transpose(1, 2)
+    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
+    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"

tests/diffusion/attention/test_trtllm_attn.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn.functional as F
 
 from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
@@ -142,6 +143,9 @@ def test_denoise_progress_mixin(monkeypatch):
     _Pipe().record_denoise_step(4)
     assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25)
 
+    _Pipe().record_denoise_step(5, normalized_timestep=0.7)
+    assert ctx.denoise_step_idx == 5 and ctx.denoise_timestep == pytest.approx(0.7)
+
 
 @pytest.mark.parametrize(
     "bk, field",
@@ -187,6 +191,7 @@ def test_mask_free_allowlist_gates_trtllm_default():
 
     assert meta("WanPipeline").attention_mask_free is True
     assert meta("WanVACEPipeline").attention_mask_free is True
+    assert meta("MiniMaxH3Pipeline").attention_mask_free is True
     assert meta("FluxPipeline").attention_mask_free is False
     assert meta(None).attention_mask_free is False
 
@@ -204,6 +209,91 @@ def test_masked_layer_raises():
         _impl().forward_cuda(q, k, v, _Meta())
 
 
+def test_packed_metadata_trims_padding_mask(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    mask = torch.tensor([[True, True, True, True, True, True, False, False]])
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+    captured = {}
+
+    def fake_attention(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(kwargs["query"])
+
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+    monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
+    monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
+
+    out = _impl().forward_cuda(q, k, v, metadata)
+
+    assert out.shape == q.shape
+    assert captured["query"].shape[0] == 6
+    assert captured["key"].shape[0] == 6
+    assert captured["value"].shape[0] == 6
+    assert captured["batch_size"] == 1
+    assert captured["seq_lens"].tolist() == [6]
+    assert captured["cum_seq_lens_q"].tolist() == [0, 6]
+    assert captured["cum_seq_lens_kv"].tolist() == [0, 6]
+    assert captured["max_q_len"] == 6
+    assert captured["max_kv_len"] == 6
+    assert torch.count_nonzero(out[:, 6:]) == 0
+
+
+def test_packed_metadata_rejects_non_prefix_mask(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]),
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+
+    with pytest.raises(ValueError, match="prefix-valid"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
+def test_packed_metadata_must_be_complete():
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, s], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed TRTLLM attention metadata"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
+def test_packed_metadata_must_cover_inputs(monkeypatch):
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    cu_seqlens = torch.tensor([0, 6], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        }
+    )
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+
+    with pytest.raises(ValueError, match="must cover all Q/K/V tokens"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
 def test_propagate_calibration_045_schema(monkeypatch):
     import types
 
@@ -295,6 +385,30 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+@requires_trtllm_attn
+def test_bf16_packed_padding_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 256, 192, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
+    metadata = AttentionMetadata(
+        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl().forward_cuda(q, k, v, metadata).float()
+    ref_used = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    used_rel = (out[:, :used] - ref_used).abs().mean() / ref_used.abs().mean()
+    assert used_rel < 0.01, f"BF16 packed valid-token rel err {used_rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
 requires_sage = pytest.mark.skipif(
     not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
     reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
@@ -321,6 +435,47 @@ def test_sage_quant_matches_sdpa(dtype_qk):
     assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
 
 
+@requires_sage
+def test_sage_packed_non_aligned_length_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 208, 194, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
+    metadata = AttentionMetadata(
+        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}).forward_cuda(q, k, v, metadata)
+    assert torch.isfinite(out).all()
+    ref = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    rel = (out[:, :used].float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.2, f"SAGE packed valid-token rel err {rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
+@requires_sage
+def test_sage_short_sequence_uses_dense_kernel():
+    torch.manual_seed(0)
+    b, s, h, d = 1, 14, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+
+    impl = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16})
+    impl._sage_quantize_fn = lambda *args, **kwargs: pytest.fail("short sequence must not use SAGE")
+    out = impl.forward_cuda(q, k, v)
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert torch.isfinite(out).all()
+    assert rel < 0.01, f"SAGE short-sequence dense fallback rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)

tests/diffusion/cache/test_cache_backends.py

@@ -86,9 +86,35 @@ class TestCacheDiTBackend:
 
         # Verify cache-dit was enabled
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_cache_dit.enable_cache.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
+        """Combined pipelines enable and refresh Cache-DiT for every DiT."""
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
+        pipeline._dit_modules = ["transformer", "transformers_ref"]
+        pipeline.transformer = Mock()
+        pipeline.transformers_ref = Mock()
+        for transformer in (pipeline.transformer, pipeline.transformers_ref):
+            transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+                block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
+            )
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=20)
+
+        assert len(backend._refresh_funcs) == 2
+        assert mock_cache_dit.enable_cache.call_count == 2
+        assert mock_cache_dit.refresh_context.call_count == 2
+        assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
+            pipeline.transformer,
+            pipeline.transformers_ref,
+        }
+
     @patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_enable_without_model_adapter_uses_cache_dit_registry(self, mock_cache_dit, mock_logger):
@@ -106,7 +132,7 @@ class TestCacheDiTBackend:
         backend.enable(pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         assert mock_cache_dit.enable_cache.call_args.args[0] is pipeline.transformer
         mock_logger.info.assert_any_call(
             "Transformer %s does not declare _cache_dit_adapter_config; "
@@ -138,7 +164,7 @@ class TestCacheDiTBackend:
 
         assert isinstance(exc_info.value.__cause__, ValueError)
         assert backend.enabled is False
-        assert backend._refresh_func is None
+        assert backend._refresh_funcs == []
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
@@ -245,7 +271,7 @@ class TestCacheDiTBackend:
         backend.enable(mock_pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.model
@@ -285,7 +311,7 @@ class TestCacheDiTBackend:
         backend.enable(mock_pipeline)
 
         assert backend.enabled is True
-        assert backend._refresh_func is not None
+        assert backend._refresh_funcs
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.transformer

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -24,12 +25,336 @@ def test_pipeline_import_registry_and_component_discovery():
         "pipeline_minimax_h3",
         "MiniMaxH3Pipeline",
     )
+    assert _DIFFUSION_MODELS["MiniMaxH3ModularPipeline"] == _DIFFUSION_MODELS["MiniMaxH3Pipeline"]
     assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
-    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
+    assert (
+        _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3ModularPipeline"] == _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"]
+    )
+    assert MiniMaxH3Pipeline._dit_modules == ["transformer", "transformers_ref"]
     assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
     assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
 
 
+def _write_partition_index(path, *, partition, tasks):
+    path.mkdir(parents=True)
+    (path / "model_index.json").write_text(
+        json.dumps(
+            {
+                "_class_name": "MiniMaxH3Pipeline",
+                "_diffusers_version": "0.32.2",
+                "_minimax_h3": {
+                    "partition": partition,
+                    "tasks": tasks,
+                    "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+
+
+def test_modular_diffusers_index_is_resolved_generically(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+    from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+    from vllm_omni.entrypoints.utils import resolve_model_config_path
+
+    (tmp_path / "modular_model_index.json").write_text(
+        json.dumps(
+            {
+                "_class_name": "MiniMaxH3ModularPipeline",
+                "_diffusers_version": "0.36.0.dev0",
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    assert is_diffusion_model(str(tmp_path))
+    assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
+    assert resolve_model_config_path(str(tmp_path)) is None
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+    assert config.model_class_name == "MiniMaxH3ModularPipeline"
+    assert config.supports_multimodal_inputs
+    assert config.max_multimodal_image_inputs == 9
+    assert config.supports_mixed_reference_inputs
+
+
+@pytest.mark.parametrize(
+    ("task_type", "partition"),
+    [
+        (None, "combined"),
+        ("auto", "combined"),
+        ("t2va", "fl2va"),
+        ("fl2va", "fl2va"),
+        ("ref2va", "ref2va"),
+    ],
+)
+def test_startup_task_selects_weight_partition(task_type, partition):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    assert _minimax_h3_partition_for_task(task_type) == partition
+
+
+def test_startup_auto_task_uses_explicit_local_partition(tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    for directory, partition, tasks in (
+        ("FL2VA", "fl2va", ["t2va", "fl2va"]),
+        ("Ref2VA", "ref2va", ["ref2va"]),
+    ):
+        model_path = tmp_path / directory
+        _write_partition_index(model_path, partition=partition, tasks=tasks)
+        assert _minimax_h3_partition_for_task(None, str(model_path)) == partition
+        assert _minimax_h3_partition_for_task("auto", str(model_path)) == partition
+        assert _minimax_h3_partition_for_task("combined", str(model_path)) == "combined"
+
+
+def test_startup_task_rejects_unsupported_value():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task
+
+    with pytest.raises(ValueError, match="task_type must be one of"):
+        _minimax_h3_partition_for_task("unsupported")
+
+
+def test_combined_task_inference_and_transformer_routing():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    pipeline.transformer = torch.nn.Identity()
+    pipeline.transformers_ref = torch.nn.Linear(1, 1)
+    assert pipeline._resolve_task(None, {}) == "t2va"
+    assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
+    assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
+    assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
+    pipeline.partition = "ref2va"
+    pipeline.supported_tasks = frozenset({"ref2va"})
+    assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
+
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    assert pipeline._transformer_for_task("fl2va") is pipeline.transformer
+    assert pipeline._transformer_for_task("ref2va") is pipeline.transformers_ref
+
+
+@pytest.mark.parametrize(
+    (
+        "task_type",
+        "expected_partition",
+        "expected_dits",
+        "component_partition",
+        "source_partitions",
+        "expected_tasks",
+    ),
+    [
+        (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}),
+        ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}),
+        ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}),
+    ],
+)
+def test_pipeline_loads_task_selected_dits_and_shared_components_once(
+    monkeypatch,
+    tmp_path,
+    task_type,
+    expected_partition,
+    expected_dits,
+    component_partition,
+    source_partitions,
+    expected_tasks,
+):
+    from vllm_omni.diffusion.data import (
+        DiffusionParallelConfig,
+        OmniDiffusionConfig,
+    )
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    _write_partition_index(
+        tmp_path / "FL2VA",
+        partition="fl2va",
+        tasks=["t2va", "fl2va"],
+    )
+    _write_partition_index(
+        tmp_path / "Ref2VA",
+        partition="ref2va",
+        tasks=["ref2va"],
+    )
+    for partition_name in ("FL2VA", "Ref2VA"):
+        for component in (
+            "transformer",
+            "tokenizer",
+            "processor",
+            "text_encoder",
+            "video_vae",
+            "audio_vae",
+        ):
+            (tmp_path / partition_name / component).mkdir()
+
+    created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+
+    class FakeModule(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+    class FakeDiT(FakeModule):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            created["dit"].append((args, kwargs))
+
+    def component_factory(name):
+        def create(path, *args, **kwargs):
+            created[name].append(str(path))
+            return FakeModule()
+
+        return create
+
+    tokenizer_calls = []
+    processor_calls = []
+    tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
+    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        tokenizer_calls.append((args, kwargs)) or object()
+    )
+    processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
+    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        processor_calls.append((args, kwargs)) or object()
+    )
+    monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3Qwen3VLEncoder",
+        component_factory("text_encoder"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3VideoVAE",
+        component_factory("video_vae"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "MiniMaxH3AudioVAE",
+        component_factory("audio_vae"),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "Qwen2TokenizerFast",
+        tokenizer_cls,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "Qwen3VLProcessor",
+        processor_cls,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "_dit_rank_world",
+        lambda: (None, 0, 1),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_local_device",
+        lambda: torch.device("cpu"),
+    )
+    download_calls = []
+
+    def fake_download(**kwargs):
+        download_calls.append(kwargs)
+        return str(tmp_path)
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "download_weights_from_hf_specific",
+        fake_download,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="MiniMaxAI/MiniMax-H3",
+        revision=None,
+        task_type=task_type,
+        quantization_config=None,
+        parallel_config=DiffusionParallelConfig(
+            cfg_parallel_size=1,
+            text_encoder_tp_size=1,
+        ),
+    )
+    pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config)
+
+    assert pipeline.partition == expected_partition
+    assert pipeline.supported_tasks == expected_tasks
+    assert len(created["dit"]) == expected_dits
+    component_path = tmp_path / component_partition
+    assert created["text_encoder"] == [str(component_path / "text_encoder")]
+    assert created["video_vae"] == [str(component_path / "video_vae")]
+    assert created["audio_vae"] == [str(component_path / "audio_vae")]
+    assert len(tokenizer_calls) == 1
+    assert len(processor_calls) == 1
+    expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"]
+    assert pipeline._dit_modules == expected_dit_modules
+    assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2)
+    if expected_partition == "ref2va":
+        assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
+    assert [source.model_or_path for source in pipeline.weights_sources] == [
+        str(tmp_path / partition_name) for partition_name in source_partitions
+    ]
+    expected_patterns = (
+        pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
+        if expected_partition == "combined"
+        else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
+    )
+    assert download_calls == [
+        {
+            "model_name_or_path": "MiniMaxAI/MiniMax-H3",
+            "cache_dir": None,
+            "allow_patterns": expected_patterns,
+            "revision": None,
+            "require_all": True,
+        }
+    ]
+
+
+def test_combined_weight_loader_routes_each_contiguous_partition():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    class FakeTransformer(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.loaded = []
+            self.post_load_calls = 0
+
+        def load_weights(self, weights):
+            self.loaded = [name for name, _ in weights]
+            return set(self.loaded)
+
+        def post_load_weights(self):
+            self.post_load_calls += 1
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = FakeTransformer()
+    pipeline.transformers_ref = FakeTransformer()
+    pipeline.text_encoder = torch.nn.Identity()
+    pipeline.video_vae = torch.nn.Identity()
+    pipeline.audio_vae = torch.nn.Identity()
+    loaded = pipeline.load_weights(
+        iter(
+            [
+                ("transformer.a", torch.ones(1)),
+                ("transformer.b", torch.ones(1)),
+                ("transformers_ref.a", torch.ones(1)),
+            ]
+        )
+    )
+
+    assert pipeline.transformer.loaded == ["a", "b"]
+    assert pipeline.transformers_ref.loaded == ["a"]
+    assert pipeline.transformer.post_load_calls == 1
+    assert pipeline.transformers_ref.post_load_calls == 1
+    assert loaded == {"transformer.a", "transformer.b", "transformers_ref.a"}
+
+
 def test_dlo_offload_plan_includes_token_refiner():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -451,6 +776,21 @@ def test_distributed_layerwise_resident_blocks_release_on_failure():
     controller.offload_resident_layers.assert_called_once_with()
 
 
+def test_distributed_layerwise_resident_blocks_can_be_skipped():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    controller = Mock()
+    pipeline._dlo_residency_controller = controller
+
+    with pipeline._resident_dit_layers_on_device(enabled=False):
+        pass
+
+    controller.load_resident_layers.assert_not_called()
+    controller.offload_resident_layers.assert_not_called()
+
+
 def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,

tests/diffusion/models/minimax_h3/test_minimax_h3_scheduler_finite_gate.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_scheduler_default_skips_large_finite_scans_and_preserves_results(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        scheduling_minimax_h3_euler_ancestral as scheduler,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", raising=False)
+    original_isfinite = torch.isfinite
+    scanned_numels = []
+
+    def record_isfinite(tensor):
+        scanned_numels.append(tensor.numel())
+        return original_isfinite(tensor)
+
+    monkeypatch.setattr(torch, "isfinite", record_isfinite)
+    state = torch.linspace(-1.0, 1.0, 8)
+    velocity = torch.linspace(0.25, -0.25, 8)
+    timestep = torch.tensor(0.5)
+    default_x0 = scheduler.minimax_h3_rf_v_to_x0(state, velocity, timestep)
+    default = scheduler.minimax_h3_euler_eta0_step(
+        state,
+        default_x0,
+        sigma_curr=0.5,
+        sigma_next=0.25,
+    )
+
+    assert scanned_numels == [1]
+
+    checked_x0 = scheduler.minimax_h3_rf_v_to_x0(
+        state,
+        velocity,
+        timestep,
+        validate_finite_tensors=True,
+    )
+    checked = scheduler.minimax_h3_euler_eta0_step(
+        state,
+        checked_x0,
+        sigma_curr=0.5,
+        sigma_next=0.25,
+        validate_finite_tensors=True,
+    )
+    torch.testing.assert_close(default_x0, checked_x0, rtol=0, atol=0)
+    torch.testing.assert_close(default, checked, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")])
+def test_scheduler_finite_gate_rejects_nonfinite_inputs(monkeypatch, invalid_value):
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        scheduling_minimax_h3_euler_ancestral as scheduler,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", "1")
+    invalid = torch.ones(8)
+    invalid[0] = invalid_value
+
+    with pytest.raises(ValueError, match="must be finite"):
+        scheduler.minimax_h3_rf_v_to_x0(
+            invalid,
+            torch.ones_like(invalid),
+            torch.tensor(0.5),
+        )
+    with pytest.raises(ValueError, match="must be finite"):
+        scheduler.minimax_h3_euler_eta0_step(
+            invalid,
+            torch.ones_like(invalid),
+            sigma_curr=0.5,
+            sigma_next=0.25,
+        )

tests/diffusion/quantization/test_int8_config.py

@@ -147,6 +147,130 @@ def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyP
     assert isinstance(method, UnquantizedLinearMethod)
 
 
+def test_fused_layer_can_be_ignored_by_its_prefix(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
+    """A layer opts out of quantization under the prefix its module was built with.
+
+    MiniMax H3 stores qkv as one checkpoint tensor, so the fused name is the
+    only handle a user has on it.
+    """
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+    config = build_quant_config("int8", ignored_layers=["blocks.0.attn.qkv_proj"])
+    layer = mocker.Mock(spec=LinearBase)
+
+    assert isinstance(config.get_quant_method(layer, "blocks.0.attn.qkv_proj"), UnquantizedLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.1.attn.qkv_proj"), NPUInt8OnlineLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.0.mlp.fc1"), NPUInt8OnlineLinearMethod)
+
+
+class TestNPUQuantMatmulShapeFallback:
+    """npu_quant_matmul rejects wide outputs, so those layers stay unquantized.
+
+    The limit applies to the per-rank shard, which is what the kernel sees. The
+    same layer can therefore be quantized at a higher TP degree.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _mock_tp(self, mocker):
+        # The fallback delegates to UnquantizedLinearMethod, which reads the TP
+        # group; stand one in so the shape gating can be tested without a
+        # distributed init.
+        mock_group = mocker.Mock()
+        mock_group.rank_in_group = 0
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+        mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    @pytest.fixture
+    def quant_config(self):
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config
+
+        return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+
+    def create_weights(self, method, layer, output_partition_sizes):
+        # Narrow input dim: the fallback path allocates for real, and only the
+        # output dimension is under test.
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=output_partition_sizes,
+            input_size=8,
+            output_size=sum(output_partition_sizes),
+            params_dtype=torch.bfloat16,
+            weight_loader=lambda param, loaded_weight, *args, **kwargs: None,
+        )
+
+    @pytest.mark.parametrize(
+        "method_name",
+        ["NPUInt8LinearMethod", "NPUInt8OnlineLinearMethod"],
+    )
+    def test_over_wide_layer_falls_back_to_unquantized(self, quant_config, method_name):
+        import vllm_omni.quantization.int8_config as int8_config
+
+        method = getattr(int8_config, method_name)(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [int8_config.NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1])
+
+        assert isinstance(layer.quant_method, UnquantizedLinearMethod)
+        assert layer.weight.dtype == torch.bfloat16
+
+    def test_layer_within_the_limit_keeps_int8(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [NPU_QUANT_MATMUL_MAX_OUT_FEATURES])
+
+        assert layer.quant_method is method
+
+    def test_tensor_parallel_sharding_brings_a_wide_layer_back(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        global_out_features = 4 * NPU_QUANT_MATMUL_MAX_OUT_FEATURES
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [global_out_features // 4])
+
+        assert layer.quant_method is method
+
+
+class TestOffloadAfterQuant:
+    def test_only_methods_advertising_the_capability_are_asked(self):
+        from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config, NPUInt8OnlineLinearMethod
+
+        class MetaDeviceOnlyMethod:
+            """Stands in for upstream online FP8: lazy, but no offload hook."""
+
+            uses_meta_device = True
+
+        model = torch.nn.Sequential(Module(), Module())
+        lazy_int8 = NPUInt8OnlineLinearMethod(
+            DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+        )
+        model[0].quant_method = lazy_int8
+        model[1].quant_method = MetaDeviceOnlyMethod()
+
+        marked = DiffusersPipelineLoader._request_offload_after_quant(model)
+
+        assert marked == 1
+        assert lazy_int8._offload_after_quant
+
+
 class TestInt8LinearMethod:
     @pytest.fixture
     def mock_quant_config(self, mocker):

tests/diffusion/test_diffusion_config_propagation.py

@@ -131,6 +131,11 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
+def test_task_type_roundtrip():
+    od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task")
+    assert od.task_type == "model-defined-task"
+
+
 def test_additional_config_roundtrip():
     additional_config = {"torchair_graph_config": {"enabled": True}}
     od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)

tests/engine/test_arg_utils.py

@@ -128,6 +128,24 @@ def test_qwen3_tts_codec_frame_rate_patching():
     assert omni_config.codec_frame_rate_hz == 12.3
 
 
+def test_qwen3_tts_startup_task_type_is_validated():
+    vllm_config = EngineArgs().create_model_config()
+
+    config = OmniModelConfig.from_vllm_model_config(
+        vllm_config,
+        model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
+        task_type="Base",
+    )
+    assert config.task_type == "Base"
+
+    with pytest.raises(ValueError, match="Qwen3-TTS --task-type must be one of"):
+        OmniModelConfig.from_vllm_model_config(
+            vllm_config,
+            model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
+            task_type="fl2va",
+        )
+
+
 def test_from_cli_args_picks_up_stage_configs_path():
     """from_cli_args should pick up stage_configs_path from namespace."""
     ns = argparse.Namespace(

tests/engine/test_stage_engine_args.py

@@ -0,0 +1,510 @@
+from __future__ import annotations
+
+import copy
+import sys
+import types
+from dataclasses import fields, replace
+from pathlib import Path
+
+import pytest
+from pydantic.fields import FieldInfo
+
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
+from vllm_omni.config.stage_config import (
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+    StageType,
+    build_stage_runtime_overrides,
+    load_deploy_config,
+    merge_pipeline_deploy,
+)
+from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
+from vllm_omni.engine import stage_init_utils
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.stage_init_utils import (
+    build_engine_args_dict,
+    build_engine_args_dict_from_omni_stage_config,
+    build_legacy_engine_args_dict,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
+
+
+def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]:
+    backend_fields = fields(config_cls)
+    backend_field_names = {backend_field.name for backend_field in backend_fields}
+    backend_config = config_cls(
+        **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
+    )
+    effective_values: dict[str, object] = {}
+    for backend_field in backend_fields:
+        value = getattr(backend_config, backend_field.name)
+        if isinstance(value, FieldInfo):
+            value = value.get_default(call_default_factory=True)
+        effective_values[backend_field.name] = value
+    return effective_values
+
+
+_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
+_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
+
+
+@pytest.fixture(autouse=True)
+def _stable_engine_arg_environment(monkeypatch, tmp_path):
+    from vllm_omni import platforms
+
+    monkeypatch.delenv("VLLM_USE_FLASHINFER_MOE_FP16", raising=False)
+    xcodec_model = tmp_path / "xcodec-model"
+    xcodec_model.mkdir()
+    monkeypatch.setenv("XCODEC1_PATH", str(xcodec_model))
+    platform = platforms.current_omni_platform
+    monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
+    monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
+    monkeypatch.setattr(platform, "is_rocm", lambda: False)
+    monkeypatch.setattr(
+        OmniDiffusionConfig,
+        "_resolve_master_port",
+        lambda config: config.master_port or 29500,
+    )
+
+    def _resolve_test_worker(engine_args):
+        worker_type = engine_args.get("worker_type")
+        if worker_type is not None and engine_args.get("worker_cls") in (None, "auto"):
+            engine_args["worker_cls"] = f"test.{worker_type}.Worker"
+
+    monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker)
+
+
+def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
+    parent_model = tmp_path / "parent-model"
+    stage_model = tmp_path / "stage-model"
+    diffusion_model = tmp_path / "diffusion-model"
+    for root in (parent_model, stage_model, diffusion_model):
+        root.mkdir()
+    (stage_model / "ar-model").mkdir()
+    (stage_model / "ar-tokenizer").mkdir()
+
+    pipeline = PipelineConfig(
+        model_type="typed-engine-args",
+        model_arch="PipelineModel",
+        stages=(
+            StagePipelineConfig(
+                stage_id=0,
+                model_stage="thinker",
+                execution_type=StageExecutionType.LLM_AR,
+                requires_multimodal_data=True,
+                hf_config_name="thinker_config",
+                engine_output_type="hidden_states",
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_subdir="ar-model",
+                tokenizer_subdir="ar-tokenizer",
+                retains_state_across_chunks=True,
+                async_chunk_process_next_stage_input_func="operator.add",
+            ),
+            StagePipelineConfig(
+                stage_id=1,
+                model_stage="talker",
+                execution_type=StageExecutionType.LLM_GENERATION,
+                input_sources=(0,),
+                engine_output_type="audio_tokens",
+            ),
+            StagePipelineConfig(
+                stage_id=2,
+                model_stage="dit",
+                execution_type=StageExecutionType.DIFFUSION,
+                input_sources=(1,),
+                final_output=True,
+                final_output_type="image",
+                model_arch="TypedDiffusionPipeline",
+            ),
+        ),
+    )
+    deploy = DeployConfig(
+        async_chunk=True,
+        trust_remote_code=True,
+        dtype="float16",
+        distributed_executor_backend="mp",
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                tensor_parallel_size=2,
+                max_num_seqs=8,
+                default_sampling_params={"extra_args": {"voice": "test"}},
+                engine_extras={
+                    "model": str(stage_model),
+                    "attention_backend": "FLASHINFER",
+                    "hf_overrides": {"rope_scaling": {"factor": 2}},
+                    "limit_mm_per_prompt": {"audio": 1},
+                    "logits_processors": ["test.custom.LogitsProcessor"],
+                    "kv_cache_memory_bytes": 1024,
+                    "omni_kv_config": {"need_send_cache": True},
+                    "max_cudagraph_capture_size": 0,
+                    "worker_cls": "test.custom.Worker",
+                },
+            ),
+            StageDeployConfig(
+                stage_id=1,
+                max_num_seqs=4,
+            ),
+            StageDeployConfig(
+                stage_id=2,
+                tensor_parallel_size=2,
+                vae_parallel_mode="spatial_shard_height",
+                diffusion_attention_config={
+                    "default": {"backend": "FLASH_ATTN"},
+                    "per_role": {"cross": {"backend": "TORCH_SDPA"}},
+                },
+                engine_extras={
+                    "model": str(diffusion_model),
+                    "engine_backend": "test.diffusion.Engine",
+                    "enable_session_state_manager": True,
+                },
+            ),
+        ],
+    )
+    return pipeline, deploy, str(parent_model)
+
+
+def _legacy_and_typed_stages(
+    pipeline: PipelineConfig,
+    deploy: DeployConfig,
+    model: str,
+    cli_overrides: dict[str, object] | None = None,
+):
+    resolved_cli_overrides = {"model": model, **(cli_overrides or {})}
+    legacy_stage_configs = merge_pipeline_deploy(
+        pipeline,
+        copy.deepcopy(deploy),
+    )
+    for stage in legacy_stage_configs:
+        stage.runtime_overrides = build_stage_runtime_overrides(
+            stage.stage_id,
+            resolved_cli_overrides,
+        )
+    legacy_stages = [stage.to_omegaconf() for stage in legacy_stage_configs]
+    omni_config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=copy.deepcopy(deploy),
+        cli_overrides=resolved_cli_overrides,
+    )
+    return legacy_stages, omni_config
+
+
+def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+    connector_spec = {"name": "SharedMemoryConnector", "extra": {"mode": "test"}}
+    cli_tokenizer = "/external/tokenizer"
+
+    typed_args_by_stage = {}
+    for stage_id in (0, 1):
+        legacy_args = build_legacy_engine_args_dict(
+            legacy_stages[stage_id],
+            model,
+            stage_connector_spec=connector_spec,
+            cli_tokenizer=cli_tokenizer,
+        )
+        typed_args = build_engine_args_dict_from_omni_stage_config(
+            omni_config.stage_by_id(stage_id),
+            model,
+            stage_connector_spec=connector_spec,
+            cli_tokenizer=cli_tokenizer,
+        )
+        typed_args_by_stage[stage_id] = typed_args
+
+        assert {name: typed_args[name] for name in legacy_args} == legacy_args
+
+    thinker_args = typed_args_by_stage[0]
+    assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
+    assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
+    assert thinker_args["stage_id"] == 0
+    assert thinker_args["stage_connector_spec"] == connector_spec
+    assert thinker_args["has_sampling_extra_args"] is True
+    assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
+    assert thinker_args["worker_cls"] == "test.custom.Worker"
+    assert thinker_args["attention_backend"] == "FLASHINFER"
+    assert thinker_args["trust_remote_code"] is True
+    assert thinker_args["dtype"] == "float16"
+    assert thinker_args["distributed_executor_backend"] == "mp"
+    assert thinker_args["tensor_parallel_size"] == 2
+    assert thinker_args["hf_overrides"] == {"rope_scaling": {"factor": 2}}
+    assert thinker_args["limit_mm_per_prompt"] == {"audio": 1}
+    assert thinker_args["logits_processors"] == ["test.custom.LogitsProcessor"]
+    assert thinker_args["retains_state_across_chunks"] is True
+    assert thinker_args["kv_cache_memory_bytes"] == 1024
+    assert thinker_args["max_cudagraph_capture_size"] == 0
+
+    talker_args = typed_args_by_stage[1]
+    assert talker_args["model"] == model
+    assert talker_args["tokenizer"] == cli_tokenizer
+    assert talker_args["worker_cls"] == "test.generation.Worker"
+    assert talker_args["disable_hybrid_kv_cache_manager"] is True
+    assert talker_args["enable_prefix_caching"] is False
+
+
+def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(2),
+        model,
+    )
+
+    assert typed_args["stage_id"] == legacy_args["stage_id"] == 2
+    assert typed_args["model"] == legacy_args["model"] == str(tmp_path / "diffusion-model")
+    assert omni_config.stage_by_id(2).diffusion_config.model == str(tmp_path / "diffusion-model")
+    assert typed_args["model_stage"] == legacy_args["model_stage"] == "dit"
+    assert typed_args["model_arch"] == legacy_args["model_arch"] == "TypedDiffusionPipeline"
+    assert typed_args["model_class_name"] == legacy_args["model_class_name"] == "TypedDiffusionPipeline"
+    assert typed_args["engine_backend"] == legacy_args["engine_backend"] == "test.diffusion.Engine"
+    assert typed_args["enable_session_state_manager"] is legacy_args["enable_session_state_manager"] is True
+    assert typed_args["parallel_config"]["tensor_parallel_size"] == 2
+    assert typed_args["parallel_config"]["vae_parallel_mode"] == "spatial_shard_height"
+    assert isinstance(typed_args["diffusion_attention_config"], AttentionConfig)
+    assert typed_args["diffusion_attention_config"].default.backend == "FLASH_ATTN"
+    assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA"
+
+
+def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    deploy.enable_prefix_caching = False
+    deploy.enable_chunked_prefill = False
+    thinker_deploy = deploy.stages[0]
+    thinker_deploy.gpu_memory_utilization = 0.75
+    thinker_deploy.disable_hybrid_kv_cache_manager = False
+    thinker_deploy.async_scheduling = False
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(0),
+        model,
+    )
+    expected = {
+        "gpu_memory_utilization": 0.75,
+        "enable_prefix_caching": False,
+        "disable_hybrid_kv_cache_manager": False,
+        "max_num_seqs": 8,
+        "enable_chunked_prefill": False,
+        "async_scheduling": False,
+    }
+
+    assert {name: legacy_args[name] for name in expected} == expected
+    assert {name: typed_args[name] for name in expected} == expected
+
+
+@pytest.mark.parametrize("stage_scoped", [False, True], ids=["global", "stage-scoped"])
+def test_typed_engine_args_preserve_inherited_model_and_load_cli_fields(tmp_path, stage_scoped):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    expected = {
+        "revision": "model-rev",
+        "tokenizer_revision": "tokenizer-rev",
+        "code_revision": "code-rev",
+        "seed": 123,
+        "download_dir": str(tmp_path / "downloads"),
+    }
+    cli_overrides = {f"stage_0_{name}" if stage_scoped else name: value for name, value in expected.items()}
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model,
+        cli_overrides,
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(0),
+        model,
+    )
+
+    assert {name: legacy_args[name] for name in expected} == expected
+    assert {name: typed_args[name] for name in expected} == expected
+
+
+def test_typed_engine_args_preserve_deploy_subdirectory_precedence(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    stage_model = tmp_path / "stage-model"
+    (stage_model / "override-model").mkdir()
+    (stage_model / "override-tokenizer").mkdir()
+    deploy.stages[0].engine_extras.update(
+        {
+            "model_subdir": "override-model",
+            "tokenizer_subdir": "override-tokenizer",
+        }
+    )
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    typed_stage = omni_config.stage_by_id(0)
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
+
+    assert typed_stage.model_config.model_subdir == "override-model"
+    assert typed_stage.model_config.tokenizer_subdir == "override-tokenizer"
+    assert typed_args["model"] == legacy_args["model"] == str(stage_model / "override-model")
+    assert typed_args["tokenizer"] == legacy_args["tokenizer"] == str(stage_model / "override-tokenizer")
+
+
+@pytest.mark.parametrize(
+    ("cli_overrides", "expected"),
+    [
+        ({}, {"need_recv_cache": True}),
+        (
+            {"stage_0_omni_kv_config": {"need_send_cache": False}},
+            {"need_send_cache": False},
+        ),
+    ],
+    ids=["topology-over-deploy", "cli-over-topology"],
+)
+def test_typed_engine_args_preserve_legacy_omni_kv_precedence(tmp_path, cli_overrides, expected):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    thinker = replace(
+        pipeline.stages[0],
+        omni_kv_config={"need_recv_cache": True},
+    )
+    pipeline = replace(pipeline, stages=(thinker, *pipeline.stages[1:]))
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model,
+        cli_overrides,
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
+    typed_stage = omni_config.stage_by_id(0)
+    typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
+
+    assert typed_stage.connector_config.omni_kv_config == expected
+    assert typed_args["omni_kv_config"] == legacy_args["omni_kv_config"] == expected
+
+
+def test_build_engine_args_dict_preserves_legacy_api(tmp_path):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, _ = _legacy_and_typed_stages(pipeline, deploy, model)
+    connector_spec = {"name": "SharedMemoryConnector", "extra": {}}
+
+    compatibility_args = build_engine_args_dict(
+        copy.deepcopy(legacy_stages[0]),
+        model,
+        stage_connector_spec=connector_spec,
+        cli_tokenizer="/external/tokenizer",
+    )
+    explicitly_legacy_args = build_legacy_engine_args_dict(
+        copy.deepcopy(legacy_stages[0]),
+        model,
+        stage_connector_spec=connector_spec,
+        cli_tokenizer="/external/tokenizer",
+    )
+
+    assert compatibility_args == explicitly_legacy_args
+
+
+def test_legacy_engine_args_drop_none_tp_without_mutating_stage_config():
+    stage_config = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"tensor_parallel_size": None},
+        default_sampling_params={},
+    )
+
+    engine_args = build_legacy_engine_args_dict(
+        stage_config,
+        model="test-model",
+    )
+
+    assert "tensor_parallel_size" not in engine_args
+    assert stage_config.engine_args == {"tensor_parallel_size": None}
+
+
+def test_typed_engine_args_own_rocm_attention_default(monkeypatch, tmp_path):
+    from vllm_omni import platforms
+
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+    monkeypatch.setattr(platforms.current_omni_platform, "is_rocm", lambda: True)
+    aiter_module = types.ModuleType("vllm._aiter_ops")
+    setattr(
+        aiter_module,
+        "rocm_aiter_ops",
+        types.SimpleNamespace(is_enabled=lambda: False),
+    )
+    monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_module)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(1),
+        model,
+    )
+
+    assert "attention_backend" not in legacy_args
+    assert typed_args["attention_backend"] == "TRITON_ATTN"
+
+
+def test_typed_ming_image_engine_args_defer_diffusion_batch_default():
+    pipeline = resolve_pipeline_config("ming_flash_omni_image")
+    assert pipeline is not None
+    deploy = load_deploy_config(_DEPLOY_DIR / "ming_flash_omni_image.yaml")
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model="/tmp",
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model="/tmp")
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(1),
+        model="/tmp",
+    )
+    typed_backend_args = {name: value for name, value in typed_args.items() if name in _DIFFUSION_BACKEND_FIELDS}
+
+    assert "max_num_seqs" not in legacy_args
+    assert "max_num_seqs" not in typed_args
+    assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1
+
+
+@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
+def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
+    pipeline = resolve_pipeline_config(model_type)
+    if pipeline is None:
+        pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
+
+    deploy = (
+        load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name)
+        if pipeline.default_deploy_config_name is not None
+        else DeployConfig()
+    )
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model="/tmp",
+    )
+
+    for legacy_stage in legacy_stages:
+        stage_id = legacy_stage.stage_id
+        legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp")
+        typed_args = build_engine_args_dict_from_omni_stage_config(
+            omni_config.stage_by_id(stage_id),
+            model="/tmp",
+        )
+        if legacy_stage.stage_type == StageType.DIFFUSION:
+            backend_fields = _DIFFUSION_BACKEND_FIELDS
+            backend_config_cls = OmniDiffusionConfig
+        else:
+            backend_fields = _LLM_BACKEND_FIELDS
+            backend_config_cls = OmniEngineArgs
+
+        missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
+        assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
+
+        legacy_effective_args = _effective_backend_values(backend_config_cls, legacy_args)
+        typed_effective_args = _effective_backend_values(backend_config_cls, typed_args)
+        assert typed_effective_args == legacy_effective_args, (
+            f"{model_type} stage {stage_id} changed effective backend arguments"
+        )

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -48,6 +48,52 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
     assert mux_calls[0]["audio"] is None
 
 
+def test_float_frames_are_converted_without_stacking_full_video(monkeypatch):
+    frame = np.array(
+        [
+            [[0.0, 0.5, 1.0, 0.2], [1.5, -0.5, 0.5, 0.8]],
+        ],
+        dtype=np.float32,
+    )
+    original = frame.copy()
+
+    def fail_stack(*args, **kwargs):
+        raise AssertionError("float video conversion must not stack all frames")
+
+    monkeypatch.setattr(video_api_utils.np, "stack", fail_stack)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame])
+
+    expected = np.array([[[128, 191, 255], [255, 64, 191]]], dtype=np.uint8)
+    assert frames.flags.c_contiguous
+    np.testing.assert_array_equal(frames, np.array([expected, expected]))
+    np.testing.assert_array_equal(frame, original)
+
+
+def test_two_dimensional_width_four_frames_are_not_treated_as_rgba():
+    frame = np.arange(8, dtype=np.float32).reshape(2, 4) / 8
+
+    frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame])
+
+    expected = np.rint(frame * 255).astype(np.uint8)
+    assert frames.shape == (2, 2, 4)
+    np.testing.assert_array_equal(frames, np.stack([expected, expected]))
+
+
+def test_mixed_float_dtypes_preserve_stacked_rounding_semantics():
+    frames = [
+        np.full((1, 1, 3), 0.1, dtype=np.float16),
+        np.full((1, 1, 3), 0.2, dtype=np.float32),
+    ]
+    stacked = np.stack(frames)
+    expected = np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+    actual = video_api_utils._coerce_video_to_uint8_frames(frames)
+
+    assert expected[0, 0, 0, 0] == 25
+    np.testing.assert_array_equal(actual, expected)
+
+
 @pytest.mark.parametrize("frame_count", [3, 4])
 def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count):
     video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3)

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -240,6 +240,28 @@ def test_serve_cli_accepts_ulysses_mode():
     assert parallel_config.ulysses_mode == "advanced_uaa"
 
 
+def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "MiniMaxAI/MiniMax-H3",
+            "--omni",
+            "--task-type",
+            "fl2va",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+
+    assert args.task_type == "fl2va"
+    assert stage_cfg["engine_args"]["task_type"] == "fl2va"
+
+
 def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     """Ensure diffusion serve CLI exposes the profiler switch."""
     parser = TrackingArgumentParser()

tests/entrypoints/test_omni_snapshot_download.py

@@ -53,6 +53,11 @@ def download_backend(monkeypatch: pytest.MonkeyPatch):
         "download_weights_from_hf_specific",
         lambda **_kwargs: picked.append("huggingface"),
     )
+    monkeypatch.setattr(
+        omni_base,
+        "file_or_path_exists",
+        lambda *_args, **_kwargs: False,
+    )
 
     def run() -> str:
         picked.clear()
@@ -82,3 +87,14 @@ def test_unset_defaults_to_huggingface(monkeypatch, download_backend):
     monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
 
     assert download_backend() == "huggingface"
+
+
+def test_modular_diffusers_defers_component_download(monkeypatch, download_backend):
+    monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
+    monkeypatch.setattr(
+        omni_base,
+        "file_or_path_exists",
+        lambda *_args, **_kwargs: True,
+    )
+
+    assert download_backend() == "none"

tests/helpers/mark.py

@@ -2,10 +2,69 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …)."""
 
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+from pathlib import Path
+
 import pytest
 from vllm.platforms import current_platform
 
-# Re-exported from tests.helpers.env (GPU wait + DeviceMemoryMonitor).
+# Marker description tag in ``pyproject.toml`` ``tool.pytest.ini_options.markers``.
+# Example: ``"H100: [hardware-resource] Tests that require H100 GPU"``.
+_HARDWARE_RESOURCE_MARKER_TAG = "[hardware-resource]"
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parents[2]
+
+
+def _load_pyproject() -> dict:
+    path = _repo_root() / "pyproject.toml"
+    try:
+        try:
+            import tomllib
+        except ModuleNotFoundError:  # Python < 3.11
+            import tomli as tomllib  # type: ignore[no-redef]
+
+        with path.open("rb") as f:
+            return tomllib.load(f)
+    except Exception:
+        return {}
+
+
+@lru_cache(maxsize=1)
+def get_hardware_mark_list() -> frozenset[str]:
+    """Return hardware SKU marker names from ``pyproject.toml``.
+
+    A marker is treated as a hardware resource marker when its description
+    contains the ``[hardware-resource]`` tag, e.g.::
+
+        "H100: [hardware-resource] Tests that require H100 GPU"
+    """
+    data = _load_pyproject()
+    entries = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("markers", [])
+    if not entries:
+        # Fallback when TOML parsing fails: scan tagged marker lines.
+        text = (_repo_root() / "pyproject.toml").read_text(encoding="utf-8")
+        return frozenset(
+            re.findall(
+                rf'^\s*"([A-Za-z0-9_]+)\s*:\s*{re.escape(_HARDWARE_RESOURCE_MARKER_TAG)}',
+                text,
+                flags=re.M,
+            )
+        )
+
+    names: set[str] = set()
+    for entry in entries:
+        text = str(entry)
+        if _HARDWARE_RESOURCE_MARKER_TAG not in text:
+            continue
+        name = text.split(":", 1)[0].strip()
+        if name:
+            names.add(name)
+    return frozenset(names)
 
 
 def cuda_marks(*, res: str, num_cards: int):

tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression coverage for MiniCPM-o 4.5 SigLIP FlashAttention."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from transformers.utils import is_flash_attn_2_available
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni_llm as minicpmo_model
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    SiglipAttention,
+    SiglipEncoder,
+    SiglipFlashAttention2,
+    SiglipVisionConfig,
+)
+
+_VISION_NUM_HEADS = 16
+_VISION_HEAD_DIM = 72
+_VISION_HIDDEN_SIZE = _VISION_NUM_HEADS * _VISION_HEAD_DIM
+_DTYPE = torch.bfloat16
+
+
+def _attention_config(*, hidden_size: int, num_heads: int) -> SiglipVisionConfig:
+    return SiglipVisionConfig(
+        hidden_size=hidden_size,
+        num_attention_heads=num_heads,
+        attention_dropout=0.0,
+    )
+
+
+def _encoder_config(*, num_layers: int = 3) -> SiglipVisionConfig:
+    config = SiglipVisionConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        attention_dropout=0.0,
+    )
+    config._attn_implementation = "flash_attention_2"
+    return config
+
+
+def _capture_output(
+    outputs: dict[str, torch.Tensor], name: str
+) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...], torch.Tensor], None]:
+    def hook(
+        _module: torch.nn.Module,
+        _inputs: tuple[torch.Tensor, ...],
+        output: torch.Tensor,
+    ) -> None:
+        outputs[name] = output
+
+    return hook
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_qkv_bshd_layout_reuses_projection_storage(mocker: MockerFixture) -> None:
+    batch_size = 2
+    seq_len = 7
+    hidden_size = 32
+    num_heads = 4
+    head_dim = hidden_size // num_heads
+    attention = SiglipFlashAttention2(_attention_config(hidden_size=hidden_size, num_heads=num_heads)).eval()
+
+    projected: dict[str, torch.Tensor] = {}
+    attention.q_proj.register_forward_hook(_capture_output(projected, "query"))
+    attention.k_proj.register_forward_hook(_capture_output(projected, "key"))
+    attention.v_proj.register_forward_hook(_capture_output(projected, "value"))
+
+    flash_inputs: dict[str, torch.Tensor] = {}
+
+    def fake_flash_attention(
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        _attention_mask: torch.Tensor | None,
+        _query_length: int,
+        *,
+        dropout: float = 0.0,
+        softmax_scale: float | None = None,
+        unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None,
+    ) -> torch.Tensor:
+        del dropout, softmax_scale, unpadding_metadata
+        assert _attention_mask is None
+        assert _query_length == seq_len
+        flash_inputs.update(
+            query=query_states,
+            key=key_states,
+            value=value_states,
+        )
+        return query_states
+
+    flash_forward = mocker.patch.object(
+        attention,
+        "_flash_attention_forward",
+        side_effect=fake_flash_attention,
+    )
+
+    hidden_states = torch.randn(batch_size, seq_len, hidden_size)
+    with torch.inference_mode():
+        output, _ = attention(hidden_states)
+
+    assert output.shape == hidden_states.shape
+    flash_forward.assert_called_once()
+
+    expected_shape = (batch_size, seq_len, num_heads, head_dim)
+    expected_stride = (
+        seq_len * num_heads * head_dim,
+        num_heads * head_dim,
+        head_dim,
+        1,
+    )
+    for name in ("query", "key", "value"):
+        projected_tensor = projected[name]
+        flash_tensor = flash_inputs[name]
+        assert flash_tensor.shape == expected_shape
+        assert flash_tensor.stride() == expected_stride
+        assert flash_tensor.is_contiguous()
+        assert flash_tensor.storage_offset() == projected_tensor.storage_offset()
+        assert flash_tensor.untyped_storage().data_ptr() == projected_tensor.untyped_storage().data_ptr()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("attention_mask", "expected_calls"),
+    [
+        pytest.param(torch.tensor([[True, True, False], [True, False, False]]), 1, id="padded"),
+        pytest.param(None, 0, id="dense"),
+    ],
+)
+def test_encoder_computes_unpadding_metadata_once_per_forward(
+    mocker: MockerFixture,
+    attention_mask: torch.Tensor | None,
+    expected_calls: int,
+) -> None:
+    encoder = SiglipEncoder(_encoder_config()).eval()
+    metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data")
+    received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = []
+
+    def fake_flash_attention(
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        _attention_mask: torch.Tensor | None,
+        _query_length: int,
+        *,
+        dropout: float = 0.0,
+        softmax_scale: float | None = None,
+        unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None,
+    ) -> torch.Tensor:
+        del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale
+        received_metadata.append(unpadding_metadata)
+        return query_states
+
+    for layer in encoder.layers:
+        mocker.patch.object(layer.self_attn, "_flash_attention_forward", side_effect=fake_flash_attention)
+
+    hidden_states = torch.randn(2, 3, 32)
+    with torch.inference_mode():
+        encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=False,
+        )
+
+    assert metadata_spy.call_count == expected_calls
+    assert len(received_metadata) == len(encoder.layers)
+    if attention_mask is None:
+        assert all(metadata is None for metadata in received_metadata)
+    else:
+        assert received_metadata[0] is not None
+        assert all(metadata is received_metadata[0] for metadata in received_metadata)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_encoder_recomputes_unpadding_metadata_for_new_mask(mocker: MockerFixture) -> None:
+    encoder = SiglipEncoder(_encoder_config(num_layers=1)).eval()
+    metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data")
+    received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = []
+
+    def fake_flash_attention(
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        _attention_mask: torch.Tensor | None,
+        _query_length: int,
+        *,
+        dropout: float = 0.0,
+        softmax_scale: float | None = None,
+        unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None,
+    ) -> torch.Tensor:
+        del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale
+        received_metadata.append(unpadding_metadata)
+        return query_states
+
+    mocker.patch.object(
+        encoder.layers[0].self_attn,
+        "_flash_attention_forward",
+        side_effect=fake_flash_attention,
+    )
+    hidden_states = torch.randn(2, 4, 32)
+    masks = (
+        torch.tensor([[True, True, True, False], [True, True, False, False]]),
+        torch.tensor([[True, True, False, False], [True, False, False, False]]),
+    )
+
+    with torch.inference_mode():
+        for attention_mask in masks:
+            encoder(
+                inputs_embeds=hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=False,
+                output_hidden_states=False,
+                return_dict=False,
+            )
+
+    assert metadata_spy.call_count == len(masks)
+    first_metadata, second_metadata = received_metadata
+    assert first_metadata is not None
+    assert second_metadata is not None
+    assert first_metadata is not second_metadata
+    assert not torch.equal(first_metadata[0], second_metadata[0])
+
+
+def _require_flash_attention() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for FlashAttention execution")
+    if not is_flash_attn_2_available():
+        pytest.skip("Transformers reports that FlashAttention 2 is unavailable")
+
+
+def _make_cuda_attention_pair() -> tuple[SiglipAttention, SiglipFlashAttention2]:
+    config = _attention_config(
+        hidden_size=_VISION_HIDDEN_SIZE,
+        num_heads=_VISION_NUM_HEADS,
+    )
+    torch.manual_seed(0)
+    eager_attention = SiglipAttention(config).eval().to(device="cuda", dtype=_DTYPE)
+    flash_attention = SiglipFlashAttention2(config).eval().to(device="cuda", dtype=_DTYPE)
+    flash_attention.load_state_dict(eager_attention.state_dict())
+    return eager_attention, flash_attention
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"})
+def test_encoder_metadata_reuse_matches_per_layer_recomputation() -> None:
+    _require_flash_attention()
+    encoder = SiglipEncoder(_encoder_config()).eval().to(device="cuda", dtype=_DTYPE)
+    torch.manual_seed(3)
+    hidden_states = torch.randn(2, 64, 32, device="cuda", dtype=_DTYPE)
+    attention_mask = torch.arange(64, device="cuda").unsqueeze(0) < torch.tensor(
+        (64, 43),
+        device="cuda",
+    ).unsqueeze(1)
+
+    with torch.inference_mode():
+        expected = hidden_states
+        for encoder_layer in encoder.layers:
+            expected = encoder_layer(
+                expected,
+                attention_mask,
+                output_attentions=False,
+                unpadding_metadata=None,
+            )[0]
+        actual = encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=False,
+        )[0]
+
+    torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"})
+def test_flash_attention_unpadded_matches_eager_reference() -> None:
+    _require_flash_attention()
+    eager_attention, flash_attention = _make_cuda_attention_pair()
+    torch.manual_seed(1)
+    hidden_states = torch.randn(
+        2,
+        64,
+        _VISION_HIDDEN_SIZE,
+        device="cuda",
+        dtype=_DTYPE,
+    )
+
+    with torch.inference_mode():
+        expected, _ = eager_attention(hidden_states)
+        actual, _ = flash_attention(hidden_states)
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"})
+@pytest.mark.parametrize(
+    "valid_lengths",
+    [
+        pytest.param((64, 43), id="moderate_padding"),
+        pytest.param((64, 17), id="high_padding"),
+    ],
+)
+def test_flash_attention_padded_matches_per_sample_unpadded(
+    valid_lengths: tuple[int, int],
+) -> None:
+    _require_flash_attention()
+    _, flash_attention = _make_cuda_attention_pair()
+    batch_size = len(valid_lengths)
+    seq_len = max(valid_lengths)
+    torch.manual_seed(2)
+    hidden_states = torch.randn(
+        batch_size,
+        seq_len,
+        _VISION_HIDDEN_SIZE,
+        device="cuda",
+        dtype=_DTYPE,
+    )
+    attention_mask = torch.arange(seq_len, device="cuda").unsqueeze(0) < torch.tensor(
+        valid_lengths,
+        device="cuda",
+    ).unsqueeze(1)
+
+    with torch.inference_mode():
+        padded_output, _ = flash_attention(
+            hidden_states,
+            attention_mask=attention_mask,
+        )
+        unpadded_outputs = [
+            flash_attention(hidden_states[index : index + 1, :valid_length])[0]
+            for index, valid_length in enumerate(valid_lengths)
+        ]
+
+    for index, valid_length in enumerate(valid_lengths):
+        torch.testing.assert_close(
+            padded_output[index : index + 1, :valid_length],
+            unpadded_outputs[index],
+            atol=2e-2,
+            rtol=2e-2,
+        )

tests/model_tests/diffusion/test_alignment.py

@@ -33,6 +33,7 @@ EXCLUDED_MODELS = [
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",
     "MiniMaxH3Pipeline",
+    "MiniMaxH3ModularPipeline",
     "StableAudioPipeline",
     "WanImageToVideoPipeline",
     "WanS2VPipeline",

vllm_omni/config/model.py

@@ -17,6 +17,8 @@ import vllm_omni.model_executor.models as me_models
 
 logger = init_logger(__name__)
 
+_QWEN3_TTS_TASK_TYPES = frozenset({"CustomVoice", "VoiceDesign", "Base"})
+
 
 class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase):
     """Config convertor for Omni multi-stage models.
@@ -101,8 +103,8 @@ class OmniModelConfig(ModelConfig):
              "audio", "latents"). If None, output type is inferred.
          stage_connector_config: Stage connector configuration dictionary.
              Contains "name" (connector name), "extra" (extra connector config).
-         task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
-             If not specified, will be inferred from model path.
+         task_type: Model-defined startup task type. Each model validates its
+             supported values and applies the corresponding behavior.
 
 
     The correct way to initialize this class is via vLLM config, as most
@@ -208,6 +210,14 @@ class OmniModelConfig(ModelConfig):
             )
             return get_hf_text_config(self.hf_config)
 
+    def _validate_startup_task_type(self) -> None:
+        """Validate startup-only task selectors owned by an AR model."""
+        if self.model_arch != "Qwen3TTSTalkerForConditionalGenerationARVLLM" or self.task_type is None:
+            return
+        if self.task_type not in _QWEN3_TTS_TASK_TYPES:
+            supported = ", ".join(sorted(_QWEN3_TTS_TASK_TYPES))
+            raise ValueError(f"Qwen3-TTS --task-type must be one of {supported}; got {self.task_type!r}")
+
     def _patch_qwen3_tts(self):
         """Patches the value of `position_id_per_seconds` in Qwen3's
         TTS's talker_config into the this class's codec_frame_rate_hz.
@@ -275,6 +285,7 @@ class OmniModelConfig(ModelConfig):
             omni_cfg._patch_qwen3_tts()
 
         omni_cfg._maybe_override_text_config()
+        omni_cfg._validate_startup_task_type()
 
         if omni_cfg.hf_config is not None:
             omni_cfg.hf_config.architectures = omni_cfg.architectures

vllm_omni/config/omni_config.py

@@ -46,6 +46,29 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str |
 
 _PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS
 
+_NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
+    {
+        "async_chunk",
+        "model",
+        "omni",
+        "output_modalities",
+        "stage_configs_path",
+        "stage_id",
+        "tokenizer",
+    }
+)
+
+# Legacy deploy files can still carry StageConfig metadata in engine_extras.
+# Those values are not backend engine inputs, and the additive typed path keeps
+# sourcing their effective values from the immutable pipeline topology.
+_LEGACY_STAGE_METADATA_EXTRA_FIELDS = frozenset(
+    {
+        "final_output",
+        "final_output_type",
+        "is_comprehension",
+    }
+)
+
 _QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None
 
 
@@ -55,6 +78,21 @@ class _QuantizationEngineOverrides(TypedDict, total=False):
 
 
 class _ModelEngineOverrides(TypedDict, total=False):
+    model: str
+    model_arch: str
+    model_subdir: str
+    tokenizer_subdir: str
+    revision: str
+    tokenizer_revision: str
+    code_revision: str
+    seed: int
+    logits_processors: list[str | type]
+    trust_remote_code: bool
+    dtype: Any
+    attention_backend: Any
+    moe_backend: str
+    hf_overrides: Any
+    limit_mm_per_prompt: dict[str, Any]
     active_stream_window: int
     enable_sleep_mode: bool
     subtalker_sampling_params: dict[str, Any]
@@ -63,6 +101,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
     task_type: str
     codec_frame_rate_hz: float
     enforce_eager: bool
+    max_cudagraph_capture_size: int
     enable_flashinfer_autotune: bool
     compilation_config: dict[str, Any]
     enable_multithread_weight_load: bool
@@ -71,6 +110,9 @@ class _ModelEngineOverrides(TypedDict, total=False):
 
 
 class _LoadEngineOverrides(TypedDict, total=False):
+    tokenizer: str
+    download_dir: str
+    skip_tokenizer_init: bool
     load_format: str
     tokenizer_mode: str
     config_format: str
@@ -78,6 +120,7 @@ class _LoadEngineOverrides(TypedDict, total=False):
 
 
 class _CacheEngineOverrides(TypedDict, total=False):
+    kv_cache_memory_bytes: int
     gpu_memory_utilization: float
     enable_prefix_caching: bool
     disable_hybrid_kv_cache_manager: bool
@@ -93,6 +136,8 @@ class _SchedulerEngineOverrides(TypedDict, total=False):
 
 
 class _RuntimeEngineOverrides(TypedDict, total=False):
+    distributed_executor_backend: Any
+    worker_cls: str
     devices: str
     num_replicas: int
     env: dict[str, Any]
@@ -113,6 +158,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False):
     ulysses_mode: str
     cfg_parallel_size: int
     vae_patch_parallel_size: int
+    vae_parallel_mode: str
     text_encoder_tp_size: int
     use_hsdp: bool
     mask_sp_padding: bool
@@ -125,6 +171,10 @@ class _ParallelEngineOverrides(_ParallelConfigEngineOverrides, total=False):
     parallel_config: _ParallelConfigEngineOverrides | Mapping[str, Any]
 
 
+class _ConnectorEngineOverrides(TypedDict, total=False):
+    omni_kv_config: dict[str, Any]
+
+
 @dataclass(frozen=True)
 class _StageEngineValues:
     """Typed projections of legacy flat per-stage ``yaml_engine_args``."""
@@ -134,6 +184,7 @@ class _StageEngineValues:
     load: _LoadEngineOverrides
     cache: _CacheEngineOverrides
     scheduler: _SchedulerEngineOverrides
+    connector: _ConnectorEngineOverrides
     runtime: _RuntimeEngineOverrides
     parallel: _ParallelEngineOverrides
     diffusion: _DiffusionEngineOverrides
@@ -255,8 +306,21 @@ def _get_deploy_config(
 
 @config
 class OmniStageModelConfig:
-    """Per-stage model behavior."""
+    """Per-stage model behavior and resolved model-engine inputs."""
 
+    model: str | None = None
+    model_arch: str | None = None
+    revision: str | None = None
+    tokenizer_revision: str | None = None
+    code_revision: str | None = None
+    seed: int | None = None
+    logits_processors: list[str | type] | None = None
+    trust_remote_code: bool = False
+    dtype: Any = "auto"
+    attention_backend: Any = None
+    moe_backend: str = "auto"
+    hf_overrides: Any = None
+    limit_mm_per_prompt: dict[str, Any] | None = None
     active_stream_window: int = Field(default=0, ge=0)
     duplex_max_sessions: int = Field(default=1, ge=1)
     enable_sleep_mode: bool = False
@@ -267,6 +331,7 @@ class OmniStageModelConfig:
     task_type: str | None = None
     codec_frame_rate_hz: float | None = None
     enforce_eager: bool = False
+    max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
     enable_flashinfer_autotune: bool | None = None
     compilation_config: dict[str, Any] | None = None
     enable_multithread_weight_load: bool = True
@@ -281,8 +346,11 @@ class OmniStageModelConfig:
 
 @config
 class OmniStageLoadConfig:
-    """Per-stage loading behavior."""
+    """Per-stage loading behavior and resolved tokenizer input."""
 
+    tokenizer: str | None = None
+    download_dir: str | None = None
+    skip_tokenizer_init: bool = False
     load_format: str = "auto"
     tokenizer_mode: str = "auto"
     config_format: str | None = None
@@ -297,9 +365,11 @@ class OmniStageCacheConfig:
     vLLM-Omni diffusion-specific cache backends such as TeaCache and Cache-DiT.
     """
 
-    gpu_memory_utilization: float = Field(default=0.90, gt=0.0, le=1.0)
-    enable_prefix_caching: bool = False
-    disable_hybrid_kv_cache_manager: bool = False
+    kv_cache_memory_bytes: int | None = Field(default=None, ge=0)
+    # None preserves backend-owned defaults; explicit values still project.
+    gpu_memory_utilization: float | None = Field(default=None, gt=0.0, le=1.0)
+    enable_prefix_caching: bool | None = None
+    disable_hybrid_kv_cache_manager: bool | None = None
     mm_processor_cache_gb: float | None = Field(default=None, ge=0.0)
 
 
@@ -307,14 +377,18 @@ class OmniStageCacheConfig:
 class OmniStageSchedulerConfig:
     """Per-stage request scheduling behavior."""
 
-    max_num_seqs: int = Field(default=128, ge=1)
+    max_num_seqs: int | None = Field(default=None, ge=1)
     max_num_batched_tokens: int | None = Field(default=None, ge=1)
     max_model_len: int | None = Field(default=None, ge=-1)
-    enable_chunked_prefill: bool = False
-    async_scheduling: bool = True
+    enable_chunked_prefill: bool | None = None
+    async_scheduling: bool | None = None
 
     def __post_init__(self) -> None:
-        if self.max_num_batched_tokens is not None and self.max_num_batched_tokens < self.max_num_seqs:
+        if (
+            self.max_num_batched_tokens is not None
+            and self.max_num_seqs is not None
+            and self.max_num_batched_tokens < self.max_num_seqs
+        ):
             raise ValueError(
                 f"max_num_batched_tokens ({self.max_num_batched_tokens}) must be >= max_num_seqs ({self.max_num_seqs})"
             )
@@ -322,8 +396,10 @@ class OmniStageSchedulerConfig:
 
 @config
 class OmniStageConnectorConfig:
-    """Per-stage inter-stage connector wiring."""
+    """Per-stage connector wiring and resolved transfer mode."""
 
+    async_chunk: bool = False
+    omni_kv_config: dict[str, Any] | None = None
     stage_connector: dict[str, Any] = field(
         default_factory=lambda: {
             "name": "SharedMemoryConnector",
@@ -336,8 +412,10 @@ class OmniStageConnectorConfig:
 
 @config
 class OmniStageRuntimeConfig:
-    """Per-stage process placement and runtime behavior."""
+    """Per-stage process placement and backend runtime behavior."""
 
+    distributed_executor_backend: Any = None
+    worker_cls: str | None = None
     devices: str | None = None
     num_replicas: int = Field(default=1, ge=1)
     env: dict[str, Any] | None = None
@@ -448,7 +526,12 @@ class _DiffusionConfigProjection:
     stage_id: int = 0
     model: str | None = None
     model_class_name: str | None = None
+    engine_backend: str | type = "default"
+    diffusion_model_runner_cls: str | type | None = None
+    request_batch_max_wait_ms: float = 0.0
+    streaming_output: bool = False
     model_arch: str | None = None
+    task_type: str | None = None
     dtype: Any = "auto"
     trust_remote_code: bool = False
     revision: str | None = None
@@ -456,6 +539,7 @@ class _DiffusionConfigProjection:
     dist_timeout: int | None = None
     nccl_port: int | None = None
     master_port: int | None = None
+    scheduler_port: int | None = None
     host: str | None = None
     port: int | None = None
     model_config: dict[str, Any] = field(default_factory=dict)
@@ -467,6 +551,7 @@ class _DiffusionConfigProjection:
     enable_cache_dit_summary: bool = False
     enable_prompt_embed_cache: bool = False
     prompt_embed_cache_size: int = Field(default=32, ge=1)
+    enable_session_state_manager: bool = False
     diffusion_load_format: str = "default"
     diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
@@ -502,6 +587,7 @@ class _DiffusionConfigProjection:
     step_execution: bool = False
     supports_multimodal_inputs: bool = False
     max_multimodal_image_inputs: int | None = None
+    supports_mixed_reference_inputs: bool = False
     model_paths: dict[str, str] = field(default_factory=dict)
     model_loaded: dict[str, bool] = field(
         default_factory=lambda: {
@@ -679,6 +765,7 @@ _DIFFUSION_SHARED_CONFIG_FIELDS = frozenset(
         "stage_id",
         "model",
         "model_arch",
+        "task_type",
         "dtype",
         "trust_remote_code",
         "revision",
@@ -694,6 +781,7 @@ _DIFFUSION_RUNTIME_CONFIG_FIELDS = frozenset(
         "port",
         "nccl_port",
         "master_port",
+        "scheduler_port",
         "worker_extension_cls",
         "enable_stage_verification",
         "prompt_file_path",
@@ -741,21 +829,85 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
 _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
 _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
 _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
+_CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
 _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
 _PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__)
 _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
+_COMMON_STAGE_ENGINE_FIELDS = (
+    _QUANTIZATION_ENGINE_FIELDS
+    | _MODEL_ENGINE_FIELDS
+    | _LOAD_ENGINE_FIELDS
+    | _CACHE_ENGINE_FIELDS
+    | _SCHEDULER_ENGINE_FIELDS
+    | _CONNECTOR_ENGINE_FIELDS
+    | _RUNTIME_ENGINE_FIELDS
+)
+_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"}
+_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - {
+    "world_size"
+}
+_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
+_DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
+    _COMMON_STAGE_ENGINE_FIELDS
+    | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
+    | _DIFFUSION_STAGE_ENGINE_FIELDS
+    | {"parallel_config"}
+)
+
+# This all-stage union is used only to discover global CLI candidates. Stage
+# validation must use the execution-type-specific sets below.
+_STAGE_ENGINE_FIELDS = _LLM_STAGE_ENGINE_FIELDS | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS
+_STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
+    StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS,
+    StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS,
+    StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS,
+}
+_PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
+    StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
+    StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
+    StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
+}
+
+
+def _validate_stage_engine_override_ownership(
+    stage_id: int,
+    execution_type: StageExecutionType,
+    overrides: Mapping[str, Any],
+) -> None:
+    try:
+        owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
+        parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
+
+    unowned_fields = set(overrides) - owner_fields
+    parallel_config = overrides.get("parallel_config")
+    if isinstance(parallel_config, Mapping):
+        unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields)
+    if unowned_fields:
+        names = ", ".join(sorted(unowned_fields))
+        raise ValueError(
+            f"Stage {stage_id} ({execution_type.value}) has explicit engine argument(s) "
+            f"with no structured config owner: {names}"
+        )
 
 
 def _global_stage_cli_fields() -> frozenset[str]:
     # Lazy import avoids vllm_omni.config -> omni_config -> engine.arg_utils ->
     # vllm_omni.config during package-level config imports.
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
+    from vllm_omni.engine.arg_utils import OmniEngineArgs, orchestrator_field_names
 
-    return (
+    candidates = (
         frozenset(f.name for f in fields(OmniEngineArgs))
         | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS)
         | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS)
-    ) - {"model", "stage_id", "stage_configs_path", "async_chunk"}
+    )
+    externally_consumed = (
+        _NON_STAGE_ENGINE_CLI_FIELDS
+        | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig))
+        | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS)
+    )
+    return candidates - externally_consumed
 
 
 def _mapping_or_empty(value: Any) -> Mapping[str, Any]:
@@ -775,17 +927,33 @@ def _stage_engine_overrides(stage_deploy: StageDeployConfig | None) -> dict[str,
         value = getattr(stage_deploy, name)
         if value is not None:
             overrides[name] = _copy_value(value)
-    overrides.update(_copy_value(stage_deploy.engine_extras))
+    overrides.update(
+        {
+            name: _copy_value(value)
+            for name, value in stage_deploy.engine_extras.items()
+            if name not in _LEGACY_STAGE_METADATA_EXTRA_FIELDS
+        }
+    )
     return overrides
 
 
 def _stage_engine_values(
     stage_deploy: StageDeployConfig | None,
+    topology: StagePipelineConfig,
     stage_cli_overrides: Mapping[str, Any] | None = None,
 ) -> _StageEngineValues:
     engine = _stage_engine_overrides(stage_deploy)
+    # Preserve legacy ordering: topology-owned KV roles override deploy
+    # extras, while an explicit CLI override remains highest priority.
+    if topology.omni_kv_config:
+        engine["omni_kv_config"] = _copy_value(topology.omni_kv_config)
     if stage_cli_overrides:
         engine.update(_copy_value(stage_cli_overrides))
+    _validate_stage_engine_override_ownership(
+        topology.stage_id,
+        topology.execution_type,
+        engine,
+    )
     return _StageEngineValues(
         quantization=cast(
             _QuantizationEngineOverrides,
@@ -798,6 +966,10 @@ def _stage_engine_values(
             _SchedulerEngineOverrides,
             _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS),
         ),
+        connector=cast(
+            _ConnectorEngineOverrides,
+            _select_engine_overrides(engine, _CONNECTOR_ENGINE_FIELDS),
+        ),
         runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)),
         parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
         diffusion=_DiffusionEngineOverrides.from_engine(engine),
@@ -893,9 +1065,10 @@ class BaseVllmOmniStageConfig:
 
     @property
     def scheduler_cls(self) -> str | None:
-        return _resolve_scheduler_path(
+        async_scheduling = self.scheduler_config.async_scheduling
+        return self.stage_pipeline_config.scheduler_cls or _resolve_scheduler_path(
             self.stage_pipeline_config.execution_type,
-            self.scheduler_config.async_scheduling,
+            True if async_scheduling is None else async_scheduling,
         )
 
     @property
@@ -957,11 +1130,14 @@ StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConf
 
 
 def _build_common_stage_config_kwargs(
+    pipeline: PipelineConfig,
     deploy: DeployConfig,
     topology: StagePipelineConfig,
     stage_deploy: StageDeployConfig | None,
     engine: _StageEngineValues,
     parallel_config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig,
+    *,
+    model: str | None,
 ) -> tuple[dict[str, Any], str | None, str | None]:
     input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk))
     quantization_config = _build_quantization_config(deploy, engine.quantization)
@@ -971,16 +1147,36 @@ def _build_common_stage_config_kwargs(
         {
             "stage_pipeline_config": topology,
             "model_config": _build_model_config(
+                pipeline,
+                deploy,
                 topology,
                 stage_deploy,
                 engine.model,
                 duplex_max_sessions=(deploy.duplex_session.max_sessions if deploy.session_mode == "duplex" else 1),
+                model=model,
+            ),
+            "load_config": _build_load_config(topology, engine.load),
+            "cache_config": _build_cache_config(
+                deploy,
+                engine.cache,
+                topology.execution_type,
+            ),
+            "scheduler_config": _build_scheduler_config(
+                deploy,
+                engine.scheduler,
+                topology.execution_type,
+            ),
+            "connector_config": _build_connector_config(
+                deploy,
+                stage_deploy,
+                engine.connector,
+            ),
+            "runtime_config": _build_runtime_config(
+                deploy,
+                stage_deploy,
+                engine.runtime,
+                parallel_config,
             ),
-            "load_config": _build_load_config(engine.load),
-            "cache_config": _build_cache_config(deploy, engine.cache),
-            "scheduler_config": _build_scheduler_config(deploy, engine.scheduler),
-            "connector_config": _build_connector_config(stage_deploy),
-            "runtime_config": _build_runtime_config(stage_deploy, engine.runtime, parallel_config),
             "parallel_config": parallel_config,
             "quantization_config": _copy_value(quantization_config),
         },
@@ -1009,10 +1205,12 @@ def _build_ar_stage_config(
     model: str | None,
 ) -> VllmOmniARStageConfig:
     common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs(
+        pipeline,
         deploy,
         topology,
         stage_deploy,
         engine,
+        model=model,
     )
     return cast(
         VllmOmniARStageConfig,
@@ -1034,10 +1232,12 @@ def _build_generation_stage_config(
     model: str | None,
 ) -> VllmOmniGenerationStageConfig:
     common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs(
+        pipeline,
         deploy,
         topology,
         stage_deploy,
         engine,
+        model=model,
     )
     return cast(
         VllmOmniGenerationStageConfig,
@@ -1059,18 +1259,20 @@ def _build_diffusion_stage_config(
     model: str | None,
 ) -> VllmOmniDiffusionStageConfig:
     common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs(
+        pipeline,
         deploy,
         topology,
         stage_deploy,
         engine,
         OmniStageDiffusionParallelConfig,
+        model=model,
     )
     common_kwargs["diffusion_config"] = _build_diffusion_config_projection(
         pipeline,
         deploy,
         topology,
         engine.diffusion,
-        model=model,
+        model=common_kwargs["model_config"].model,
         quantization_config=common_kwargs["quantization_config"],
     )
     return cast(
@@ -1125,14 +1327,28 @@ def _build_quantization_config(
 
 
 def _build_model_config(
+    pipeline: PipelineConfig,
+    deploy: DeployConfig,
     topology: StagePipelineConfig,
     stage_deploy: StageDeployConfig | None,
     engine: _ModelEngineOverrides,
     *,
     duplex_max_sessions: int,
+    model: str | None,
 ) -> OmniStageModelConfig:
     default_sampling_params = _stage_sampling_params(stage_deploy, topology)
     kwargs = _config_kwargs(engine)
+    kwargs["model"] = _first_defined(kwargs.get("model"), model)
+    if "model_arch" not in kwargs:
+        kwargs["model_arch"] = topology.model_arch or pipeline.model_arch or None
+    if "trust_remote_code" not in kwargs and deploy.trust_remote_code is not None:
+        kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code)
+    if "dtype" not in kwargs and deploy.dtype is not None:
+        kwargs["dtype"] = _copy_value(deploy.dtype)
+    if "active_stream_window" not in kwargs:
+        kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window)
+    if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None:
+        kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir)
     if "has_sampling_extra_args" not in kwargs:
         kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args"))
     if "model_subdir" not in kwargs and topology.model_subdir is not None:
@@ -1146,45 +1362,67 @@ def _build_model_config(
     )
 
 
-def _build_load_config(engine: _LoadEngineOverrides) -> OmniStageLoadConfig:
-    return OmniStageLoadConfig(**_config_kwargs(engine))
+def _build_load_config(
+    topology: StagePipelineConfig,
+    engine: _LoadEngineOverrides,
+) -> OmniStageLoadConfig:
+    kwargs = _config_kwargs(engine)
+    if "skip_mm_profiling" not in kwargs and not topology.requires_multimodal_data:
+        kwargs["skip_mm_profiling"] = True
+    return OmniStageLoadConfig(**kwargs)
 
 
 def _build_cache_config(
     deploy: DeployConfig,
     engine: _CacheEngineOverrides,
+    execution_type: StageExecutionType,
 ) -> OmniStageCacheConfig:
     kwargs = _config_kwargs(engine)
     if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None:
         kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching)
+    if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION:
+        # Match the generation-stage override applied by legacy finalization.
+        kwargs["disable_hybrid_kv_cache_manager"] = True
     return OmniStageCacheConfig(**kwargs)
 
 
 def _build_scheduler_config(
     deploy: DeployConfig,
     engine: _SchedulerEngineOverrides,
+    execution_type: StageExecutionType,
 ) -> OmniStageSchedulerConfig:
     kwargs = _config_kwargs(engine)
     if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None:
         kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill)
+    if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR:
+        kwargs["async_scheduling"] = True
     return OmniStageSchedulerConfig(**kwargs)
 
 
-def _build_connector_config(stage_deploy: StageDeployConfig | None) -> OmniStageConnectorConfig:
+def _build_connector_config(
+    deploy: DeployConfig,
+    stage_deploy: StageDeployConfig | None,
+    engine: _ConnectorEngineOverrides,
+) -> OmniStageConnectorConfig:
     output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
     input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
     return OmniStageConnectorConfig(
+        async_chunk=bool(deploy.async_chunk),
+        omni_kv_config=_copy_value(engine.get("omni_kv_config")),
         output_connectors=_copy_value(output_connectors) if output_connectors else None,
         input_connectors=_copy_value(input_connectors) if input_connectors else None,
     )
 
 
 def _build_runtime_config(
+    deploy: DeployConfig,
     stage_deploy: StageDeployConfig | None,
     engine: _RuntimeEngineOverrides,
     parallel_config: OmniStageParallelConfig,
 ) -> OmniStageRuntimeConfig:
     kwargs = _config_kwargs(engine)
+    if "distributed_executor_backend" not in kwargs and deploy.distributed_executor_backend is not None:
+        kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend)
     if "devices" not in kwargs and stage_deploy is not None and stage_deploy.devices is not None:
         kwargs["devices"] = _copy_value(stage_deploy.devices)
     if "num_replicas" not in kwargs and stage_deploy is not None:
@@ -1234,13 +1472,15 @@ def _build_diffusion_config_projection(
         topology.model_arch,
         pipeline.model_arch,
     )
+    if "model_class_name" not in diffusion_kwargs and topology.model_arch is not None:
+        diffusion_kwargs["model_class_name"] = _copy_value(topology.model_arch)
     if "dtype" not in diffusion_kwargs and deploy.dtype is not None:
         diffusion_kwargs["dtype"] = _copy_value(deploy.dtype)
     if "trust_remote_code" not in diffusion_kwargs and deploy.trust_remote_code is not None:
         diffusion_kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code)
     if "distributed_executor_backend" not in diffusion_kwargs and deploy.distributed_executor_backend is not None:
         diffusion_kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend)
-    if model is not None:
+    if "model" not in diffusion_kwargs and model is not None:
         diffusion_kwargs["model"] = model
     if quantization_config is not None:
         diffusion_kwargs["quantization_config"] = _copy_value(quantization_config)
@@ -1302,6 +1542,7 @@ class VllmOmniConfig:
                 deploy_by_id.get(topology.stage_id),
                 _stage_engine_values(
                     deploy_by_id.get(topology.stage_id),
+                    topology,
                     _stage_cli_overrides(topology.stage_id, cli_overrides),
                 ),
                 model=model,

vllm_omni/diffusion/attention/backends/abstract.py

@@ -20,6 +20,21 @@ class AttentionBackend(ABC):
     # avoid a slower masked-attention plan when tail padding is not semantic.
     supports_prefix_kv_slicing: bool = 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.
+    supported_platforms: tuple[str, ...] | None = None
+
+    @classmethod
+    def validate_available(cls) -> None:
+        """Raise if this backend's optional dependencies are missing.
+
+        Called during platform resolution, i.e. before model construction, so a
+        backend that probes its kernel package lazily still reports the problem
+        while the user can still act on it.
+        """
+        return None
+
     @classmethod
     def supports_attention_mask(cls) -> bool:
         return False
@@ -63,6 +78,23 @@ class QueryRange:
     global_start: int
 
 
+@dataclass(frozen=True, slots=True)
+class VideoTokenLayout:
+    """Where the video segment sits inside a packed multimodal sequence.
+
+    A model that packs its sequence as ``[prefix | t*h*w video rows | padding]``
+    publishes this so backends can recover spatiotemporal locality; the prefix
+    holds everything that is not video (text, visual conditions, audio).
+    Publishing it also asserts that any ``attn_mask`` masks only the trailing
+    padding, so ``prefix_len + t*h*w`` is the used length of the sequence.
+
+    Plain ints, so reading it never forces a device-to-host sync.
+    """
+
+    prefix_len: int
+    latent_grid: tuple[int, int, int]
+
+
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
@@ -95,6 +127,10 @@ class AttentionMetadata:
     full_attn_spans: list[list[tuple[int, int]]] | None = None
     query_ranges: tuple[QueryRange, ...] | None = None
 
+    # Geometry of the video segment for backends that exploit spatiotemporal
+    # locality (block-sparse selection, tiled masks). Dense backends ignore it.
+    video_layout: VideoTokenLayout | None = None
+
 
 T = TypeVar("T", bound=AttentionMetadata)
 

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import extract_layer_index
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
+from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+logger = init_logger(__name__)
+
+# The rf_v2 kernel only implements a 128-token block.
+_BLOCK_SIZE = 128
+
+# Below this many video blocks, the pooling and gather that block selection adds
+# cost more than the QK work it removes, so stay dense.
+_MIN_VIDEO_BLOCKS = 32
+
+# rf_v2's ``input_layout`` describes the caller's tensors, and everything below
+# slices the sequence on dim 1. vLLM-Omni diffusion attention hands the impl
+# [B, S, N, D], so that is the only layout this backend accepts.
+_INPUT_LAYOUT = "BSND"
+
+# rf_v2 precision mode: 0 = high precision, 1 = high performance. Kept at the
+# precise setting because the sparsity knob is the intended perf lever, and on
+# A5 devices the kernel is routed to rf_v3, which overrides this anyway.
+_INNER_PRECISE = 0
+
+_WRONG_PLATFORM = (
+    "RAINFUSION_ATTN runs the MindIE-SD rf_v2 kernel and is available on Ascend NPU only. "
+    "Select FLASH_ATTN or TORCH_SDPA on this platform."
+)
+
+_MISSING_MINDIESD = (
+    "RAINFUSION_ATTN requires MindIE-SD. Please install MindIE-SD to enable RainFusion sparse "
+    "attention on Ascend NPU. For installation details, see https://gitcode.com/Ascend/MindIE-SD "
+    "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN"
+)
+
+
+def _try_extract_layer_index(prefix: str) -> int | None:
+    if not prefix:
+        return None
+    try:
+        return extract_layer_index(prefix)
+    except (AssertionError, ValueError):
+        return None
+
+
+@dataclass(frozen=True)
+class RainFusionConfig:
+    """Resolved RainFusion controls for one attention layer.
+
+    ``sparsity`` is the nominal fraction of key blocks dropped per query block.
+    The realized sparsity is lower because rf_v2 always keeps the prefix rows and
+    the first-frame blocks. ``start_step`` and ``skip_layers`` are the accuracy
+    knobs: early denoise steps and specific DiT blocks stay dense.
+    """
+
+    sparsity: float = 0.0
+    start_step: int = 0
+    skip_layers: frozenset[int] = frozenset()
+
+    @classmethod
+    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> RainFusionConfig:
+        bk = backend_kwargs or {}
+        return cls(
+            sparsity=float(bk.get("sparsity", 0.0)),
+            start_step=int(bk.get("start_step", 0)),
+            skip_layers=frozenset(bk.get("skip_layers") or ()),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return self.sparsity > 0.0
+
+
+@dataclass(frozen=True)
+class RainFusionPlan:
+    """Per-forward geometry handed to the rf_v2 kernel."""
+
+    prefix_len: int
+    used_len: int
+    latent_shape: list[int]
+
+
+class RainFusionAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+    supported_platforms: tuple[str, ...] = ("npu",)
+
+    @classmethod
+    def validate_available(cls) -> None:
+        from importlib.util import find_spec
+
+        if find_spec("mindiesd") is None:
+            raise ValueError(_MISSING_MINDIESD)
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "RAINFUSION_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type[RainFusionAttentionImpl]:
+        return RainFusionAttentionImpl
+
+
+class RainFusionAttentionImpl(AttentionImpl):
+    """Block-sparse video attention via MindIE-SD RainFusion (rf_v2) on Ascend NPU.
+
+    Sparsity applies only to the video segment of a packed multimodal sequence,
+    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
+    other case — warmup denoise steps, exempt layers, a layer that does not declare
+    ``qkv_layout="BSND"``, sequences without a published video segment, video
+    segments too short to pay for block selection or not aligned to the kernel's
+    block size — delegates to FlashAttention, so a model can select this backend
+    unconditionally.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+
+        self.rainfusion = RainFusionConfig.from_backend_kwargs(backend_kwargs)
+        self.layer_idx = _try_extract_layer_index(prefix)
+
+        if self.rainfusion.enabled:
+            self._validate_parallel_config()
+            if causal:
+                raise ValueError(
+                    "RAINFUSION_ATTN does not support causal attention: rf_v2 selects key "
+                    "blocks by pooled relevance and cannot express a causal mask. Select "
+                    "FLASH_ATTN for causal roles."
+                )
+            if qkv_layout is not None and qkv_layout.upper() != _INPUT_LAYOUT:
+                raise ValueError(
+                    f"RAINFUSION_ATTN needs {_INPUT_LAYOUT} tensors to locate the video segment along "
+                    f"the sequence axis, but this layer declares qkv_layout={qkv_layout!r}. Select "
+                    "FLASH_ATTN for this role."
+                )
+
+        self.dense_fallback = FlashAttentionBackend.get_impl_cls()(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+            prefix=prefix,
+            qkv_layout=qkv_layout,
+        )
+
+    def _validate_parallel_config(self) -> None:
+        config = get_current_diffusion_config_or_none()
+        parallel_config = getattr(config, "parallel_config", None)
+        ring_degree = getattr(parallel_config, "ring_degree", 1)
+        if ring_degree > 1:
+            # Ring gives each rank a slice of the sequence, so block selection
+            # would score only local keys and the layer bypasses the backend
+            # entirely (see Attention._run_ring_attention).
+            raise ValueError(
+                "RAINFUSION_ATTN is not compatible with ring sequence parallelism "
+                f"(ring_degree={ring_degree}): rf_v2 needs the whole key sequence to rank "
+                "blocks. Use Ulysses SP (ring_degree=1) instead."
+            )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        # ROCm and MUSA route through forward_cuda by default, so this covers them too.
+        raise NotImplementedError(_WRONG_PLATFORM)
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError(_WRONG_PLATFORM)
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        plan = self._resolve_plan(attn_metadata)
+        if plan is None:
+            return self.dense_fallback.forward_npu(query, key, value, attn_metadata)
+        return self._forward_sparse_npu(query, key, value, plan)
+
+    def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None:
+        """Return the rf_v2 geometry, or None when this forward must stay dense."""
+        rf = self.rainfusion
+        if not rf.enabled:
+            return None
+        if self.layer_idx is not None and self.layer_idx in rf.skip_layers:
+            return None
+        if is_forward_context_available():
+            step_idx = get_forward_context().denoise_step_idx
+            if step_idx is not None and step_idx < rf.start_step:
+                return None
+        if self.qkv_layout is None:
+            # The sparse path reads the sequence off dim 1, which the tensors alone
+            # do not establish, and the dense fallback resolves an absent layout its
+            # own way. Sparsifying on an assumption would put the two paths on
+            # different axes, so an undeclared layout stays dense.
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: this layer does not declare qkv_layout, and rf_v2 "
+                "needs %s to locate the video segment along the sequence axis. Set qkv_layout=%r on "
+                "the Attention layer to enable sparsity.",
+                _INPUT_LAYOUT,
+                _INPUT_LAYOUT,
+            )
+            return None
+
+        if attn_metadata is None:
+            return None
+
+        layout = attn_metadata.video_layout
+        if layout is None:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: this attention role carries no video segment. The "
+                "model must publish AttentionMetadata.video_layout for the sequence to be sparsified."
+            )
+            return None
+        max_seqlen_q = attn_metadata.extra.get("max_seqlen_q")
+        if max_seqlen_q is None:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the "
+                "video segment cannot be confirmed to be the tail of packed document 0."
+            )
+            return None
+
+        prefix_len = int(layout.prefix_len)
+        latent_shape = [int(dim) for dim in layout.latent_grid]
+        # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of
+        # the packed sequence holds those rows; anything past it is alignment
+        # padding that rf_v2 must not see.
+        video_len = math.prod(latent_shape)
+        used_len = prefix_len + video_len
+
+        if used_len != int(max_seqlen_q):
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: prefix (%d) plus latent grid %s does not fill "
+                "packed document 0 (%d rows). rf_v2 requires the video segment to be its tail.",
+                prefix_len,
+                tuple(latent_shape),
+                int(max_seqlen_q),
+            )
+            return None
+        if video_len < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: %d video rows is under the %d-row "
+                "(%d block) threshold where sparse selection pays off.",
+                video_len,
+                _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
+                _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 "
+            "because prefix and first-frame blocks are always kept.",
+            rf.sparsity,
+            rf.start_step,
+            len(rf.skip_layers),
+            tuple(latent_shape),
+            prefix_len,
+            video_len,
+        )
+        return RainFusionPlan(
+            prefix_len=prefix_len,
+            used_len=used_len,
+            latent_shape=latent_shape,
+        )
+
+    def _forward_sparse_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        plan: RainFusionPlan,
+    ) -> torch.Tensor:
+        try:
+            from mindiesd import sparse_attention
+        except ImportError:
+            raise ImportError(_MISSING_MINDIESD)
+
+        used = plan.used_len
+        q, k, v = (tensor[:, :used] for tensor in (query, key, value))
+        # Ulysses has already gathered the full sequence onto this rank and split
+        # the heads, so read the head count off the tensor rather than num_heads.
+        out = sparse_attention(
+            q,
+            k,
+            v,
+            scale=self.softmax_scale,
+            head_num=query.shape[-2],
+            input_layout=_INPUT_LAYOUT,
+            inner_precise=_INNER_PRECISE,
+            sparse_type="rf_v2",
+            txt_len=plan.prefix_len,
+            block_size=_BLOCK_SIZE,
+            latent_shape_q=plan.latent_shape,
+            latent_shape_k=plan.latent_shape,
+            sparsity=self.rainfusion.sparsity,
+        )
+        if used == query.shape[1]:
+            return out
+        padded = torch.zeros_like(query)
+        padded[:, :used] = out
+        return padded

vllm_omni/diffusion/attention/backends/registry.py

@@ -65,6 +65,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
     FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
     TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"
+    RAINFUSION_ATTN = "vllm_omni.diffusion.attention.backends.rainfusion_attn.RainFusionAttentionBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -180,6 +180,7 @@ class TrtllmAttentionImpl(AttentionImpl):
         prefix: str = "",
         qkv_layout: str | None = None,
         backend_kwargs: dict | None = None,
+        role: str = "self",
         **extra_impl_args,
     ) -> None:
         self.num_heads = num_heads
@@ -187,6 +188,7 @@ class TrtllmAttentionImpl(AttentionImpl):
         self.softmax_scale = softmax_scale
         self.causal = causal
         self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.role = role
 
         self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs)
         self._warned_missing_timestep = False
@@ -255,8 +257,16 @@ class TrtllmAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
+        extra = getattr(attn_metadata, "extra", {}) if attn_metadata is not None else {}
+        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+        present_packed_keys = [key for key in packed_keys if key in extra]
+        if present_packed_keys and len(present_packed_keys) != len(packed_keys):
+            missing = sorted(set(packed_keys) - set(present_packed_keys))
+            raise ValueError(f"Incomplete packed TRTLLM attention metadata; missing {missing}")
+        has_packed_metadata = len(present_packed_keys) == len(packed_keys)
+
         attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
-        if attn_mask is not None:
+        if attn_mask is not None and not has_packed_metadata:
             raise ValueError(
                 "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with "
                 "mask_sp_padding=True and a non-divisible seq len). Either set "
@@ -270,30 +280,96 @@ class TrtllmAttentionImpl(AttentionImpl):
                 "another backend via --diffusion-attention-backend."
             )
 
-        batch, q_len, num_q_heads, head_dim = query.shape
+        physical_batch, q_len, num_q_heads, head_dim = query.shape
         kv_len, num_kv_heads = key.shape[1], key.shape[2]
 
         device = query.device
 
-        q = query.reshape(batch * q_len, num_q_heads, head_dim).contiguous()
-        k = key.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous()
-        v = value.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous()
+        q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous()
+        k = key.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
+        v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
+        output_tokens = q.shape[0]
 
-        seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device)
-        cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device)
-        cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device)
+        if has_packed_metadata:
+            cu_seq_lens_q = extra["cu_seqlens_q"]
+            cu_seq_lens_kv = extra["cu_seqlens_k"]
+            if cu_seq_lens_q.ndim != 1 or cu_seq_lens_kv.ndim != 1:
+                raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional")
+            if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2:
+                raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches")
+            if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
+                raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")
+
+            # Drop trailing empty sequences retained by packed producers when
+            # alignment adds no padding tokens.
+            while (
+                cu_seq_lens_q.numel() > 2
+                and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
+                and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
+            ):
+                cu_seq_lens_q = cu_seq_lens_q[:-1]
+                cu_seq_lens_kv = cu_seq_lens_kv[:-1]
+
+            # A packed prefix mask denotes structural suffix padding, not another
+            # document. Exclude it before quantization and attention; in particular,
+            # SAGE block quantization must not span the valid/padding boundary.
+            if attn_mask is not None:
+                if q.shape[0] != k.shape[0] or attn_mask.numel() != q.shape[0]:
+                    raise ValueError("Packed TRTLLM prefix masks require equal Q/K lengths")
+                flat_mask = attn_mask.reshape(-1).to(dtype=torch.bool)
+                valid_tokens = int(flat_mask.sum().item())
+                expected_mask = torch.arange(flat_mask.numel(), device=device) < valid_tokens
+                if not torch.equal(flat_mask, expected_mask):
+                    raise ValueError("TRTLLM_ATTN only supports prefix-valid packed attention masks")
+
+                q_boundary = (cu_seq_lens_q == valid_tokens).nonzero(as_tuple=False)
+                kv_boundary = (cu_seq_lens_kv == valid_tokens).nonzero(as_tuple=False)
+                if q_boundary.numel() != 1 or kv_boundary.numel() != 1:
+                    raise ValueError("Packed TRTLLM prefix length must be a Q/K sequence boundary")
+                q_end = int(q_boundary.item()) + 1
+                kv_end = int(kv_boundary.item()) + 1
+                if q_end != kv_end:
+                    raise ValueError("Packed TRTLLM prefix masks require matching Q/K sequence batches")
+
+                q = q[:valid_tokens]
+                k = k[:valid_tokens]
+                v = v[:valid_tokens]
+                cu_seq_lens_q = cu_seq_lens_q[:q_end].contiguous()
+                cu_seq_lens_kv = cu_seq_lens_kv[:kv_end].contiguous()
+
+            batch = cu_seq_lens_q.numel() - 1
+            seq_lens = (cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous()
+            max_q_len = int(extra["max_seqlen_q"])
+            max_kv_len = int(extra["max_seqlen_k"])
+        else:
+            batch = physical_batch
+            seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device)
+            cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device)
+            cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device)
+            max_q_len = q_len
+            max_kv_len = kv_len
         workspace = self._get_workspace(device)
 
         bmm1_scale = self.softmax_scale
         bmm2_scale = 1.0
 
-        _skip_factor = self._resolve_skip_factor(kv_len)
+        _skip_factor = self._resolve_skip_factor(max_kv_len)
 
         # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when
         # SAGE quant is active (which already requires the kernel, checked at init) so the dense
         # path stays compatible with older builds that lack these parameters.
         sage_kwargs: dict = {}
-        if self.quant.enabled:
+        # The SAGE kernel requires every KV sequence to contain at least one full
+        # quantization block. Small auxiliary attention sites use the dense kernel.
+        use_sage = self.quant.enabled and bool(torch.all(seq_lens >= self.quant.k_block_size).item())
+        if self.quant.enabled and not use_sage:
+            message = (
+                f"TRTLLM_ATTN SAGE quantization is configured for attention role {self.role!r}, but at least one "
+                f"KV sequence is shorter than k_block_size={self.quant.k_block_size}. Falling back to dense "
+                "attention for this input."
+            )
+            logger.warning_once(message)
+        if use_sage:
             q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(q, k, v, self._sage_quantize_fn)
             sage_kwargs["sage_attn_sfs"] = sage_attn_sfs
             sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes
@@ -304,8 +380,8 @@ class TrtllmAttentionImpl(AttentionImpl):
             value=v,
             workspace_buffer=workspace,
             seq_lens=seq_lens,
-            max_q_len=q_len,
-            max_kv_len=kv_len,
+            max_q_len=max_q_len,
+            max_kv_len=max_kv_len,
             bmm1_scale=bmm1_scale,
             bmm2_scale=bmm2_scale,
             o_sf_scale=-1.0,
@@ -319,4 +395,12 @@ class TrtllmAttentionImpl(AttentionImpl):
             skip_softmax_threshold_scale_factor=_skip_factor,
             **sage_kwargs,
         )
-        return out.reshape(batch, q_len, num_q_heads, head_dim)
+        if out.shape[0] != output_tokens:
+            padded_out = torch.zeros(
+                (output_tokens, num_q_heads, head_dim),
+                dtype=out.dtype,
+                device=out.device,
+            )
+            padded_out[: out.shape[0]] = out
+            out = padded_out
+        return out.reshape(physical_batch, q_len, num_q_heads, head_dim)

vllm_omni/diffusion/attention/layer.py

@@ -89,6 +89,15 @@ class Attention(nn.Module):
             role_category=role_category,
             allow_trtllm_default=allow_trtllm_default,
         )
+        parallel_config = getattr(config, "parallel_config", None)
+        allgather_degree = getattr(parallel_config, "allgather_degree", 1)
+        # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
+        # so validation does not depend on backend names.
+        if not skip_sequence_parallel and allgather_degree > 1 and attn_backend_cls.get_name() == "TRTLLM_ATTN":
+            raise ValueError(
+                "TRTLLM_ATTN does not support AllGather-KV sequence parallelism. "
+                "Set --allgather-degree 1 or select another diffusion attention backend."
+            )
         if spec is not None:
             backend_kwargs = spec.backend_kwargs()
             self.backend_pref = spec.backend
@@ -107,6 +116,7 @@ class Attention(nn.Module):
             qkv_layout=qkv_layout,
             prefix=prefix,
             backend_kwargs=backend_kwargs,
+            role=role,
         )
         # Instantiate fallback backend for float32 support
         self.sdpa_fallback = SDPABackend.get_impl_cls()(

vllm_omni/diffusion/cache/cachedit/backend.py

@@ -3,10 +3,13 @@
 
 """Generic Cache-DiT backend lifecycle and integration."""
 
+from __future__ import annotations
+
 from collections.abc import Callable
-from typing import Any, TypeAlias
+from typing import TYPE_CHECKING, Any, TypeAlias, cast
 
 import cache_dit
+import torch.nn as nn
 from cache_dit import BlockAdapter, DBCacheConfig
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from vllm.logger import init_logger
@@ -18,6 +21,9 @@ from vllm_omni.diffusion.cache.cachedit.config import (
 )
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+
 logger = init_logger(__name__)
 
 RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None]
@@ -27,14 +33,18 @@ CacheDiTEnabler: TypeAlias = Callable[[Any, DiffusionCacheConfig], RefreshCacheC
 CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {}
 
 
-def cache_summary(pipeline: Any, details: bool = True) -> None:
+def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]:
+    """Return the pipeline DiT attributes that are present at runtime."""
+    names = getattr(pipeline, "_dit_modules", None)
+    if not isinstance(names, (list, tuple)):
+        names = ("transformer",)
+    return tuple(name for name in names if isinstance(name, str) and getattr(pipeline, name, None) is not None)
+
+
+def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None:
     """Log Cache-DiT statistics for every transformer on the pipeline."""
 
-    transformers = [
-        transformer
-        for attribute in ("transformer", "transformer_2")
-        if (transformer := getattr(pipeline, attribute, None)) is not None
-    ]
+    transformers = [getattr(pipeline, name) for name in _dit_module_names(pipeline)]
     for transformer in transformers:
         cache_dit.summary(transformer, details=details)
 
@@ -42,19 +52,34 @@ def cache_summary(pipeline: Any, details: bool = True) -> None:
         logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute")
 
 
-def _default_get_pipeline_transformer(pipeline: Any) -> Any:
-    return pipeline.transformer
+def _default_get_pipeline_transformer(
+    pipeline: SupportsComponentDiscovery,
+) -> nn.Module:
+    return cast(nn.Module, getattr(pipeline, "transformer"))
+
+
+def _make_pipeline_transformer_getter(
+    name: str,
+) -> Callable[[SupportsComponentDiscovery], nn.Module]:
+    def get_pipeline_transformer(
+        pipeline: SupportsComponentDiscovery,
+    ) -> nn.Module:
+        return cast(nn.Module, getattr(pipeline, name))
+
+    return get_pipeline_transformer
 
 
 def _build_cache_context_refresh(
     cache_config: DiffusionCacheConfig,
-    get_pipeline_transformer: Callable[[Any], Any] = _default_get_pipeline_transformer,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
 ) -> RefreshCacheContextFunc:
     """Build the cache context refresh callback for one transformer."""
 
     projected_config = CacheDiTConfig.from_diffusion_config(cache_config)
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+    def refresh_cache_context(
+        pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True
+    ) -> None:
         transformer = get_pipeline_transformer(pipeline)
 
         # Cache-DiT has no predefined SCM mask for these small step counts.
@@ -80,10 +105,11 @@ def _build_cache_context_refresh(
 
 
 def enable_cache_for_dit(
-    pipeline: Any,
+    pipeline: SupportsComponentDiscovery,
     cache_config: DiffusionCacheConfig,
     block_adapter: BlockAdapter | None = None,
     adapter_cls: type[CachedAdapter] | None = None,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
 ) -> RefreshCacheContextFunc:
     """Enable Cache-DiT for a standard single-transformer DiT pipeline."""
 
@@ -98,7 +124,7 @@ def enable_cache_for_dit(
         db_cache_config.max_warmup_steps,
     )
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     cache_target = transformer if block_adapter is None else block_adapter
     if adapter_cls is not None:
         adapter_cls.apply(
@@ -128,13 +154,16 @@ def enable_cache_for_dit(
             calibrator_config=calibrator_config,
         )
 
-    return _build_cache_context_refresh(cache_config)
+    return _build_cache_context_refresh(cache_config, get_pipeline_transformer)
 
 
-def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None:
+def _maybe_build_block_adapter(
+    pipeline: SupportsComponentDiscovery,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
+) -> BlockAdapter | None:
     """Build the model-declared block adapter, when one is configured."""
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
     if adapter_config is None:
         logger.info(
@@ -160,10 +189,13 @@ def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None:
     )
 
 
-def _maybe_get_cached_adapter_cls(pipeline: Any) -> type[CachedAdapter] | None:
+def _maybe_get_cached_adapter_cls(
+    pipeline: SupportsComponentDiscovery,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
+) -> type[CachedAdapter] | None:
     """Return the custom cached adapter declared by the transformer."""
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
     return None if adapter_config is None else adapter_config.cached_adapter_cls
 
@@ -180,29 +212,37 @@ class CacheDiTBackend(CacheBackend):
             config = cache_config
 
         super().__init__(config)
-        self._refresh_func: RefreshCacheContextFunc | None = None
+        self._refresh_funcs: list[RefreshCacheContextFunc] = []
 
-    def enable(self, pipeline: Any) -> None:
+    def enable(self, pipeline: SupportsComponentDiscovery) -> None:
         pipeline_name = type(pipeline).__name__
         custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name)
         if custom_enabler is not None:
             logger.info("Using custom cache-dit enabler for model: %s", pipeline_name)
-            self._refresh_func = custom_enabler(pipeline, self.config)
+            self._refresh_funcs = [custom_enabler(pipeline, self.config)]
         else:
-            block_adapter = _maybe_build_block_adapter(pipeline)
-            adapter_cls = _maybe_get_cached_adapter_cls(pipeline)
-            self._refresh_func = enable_cache_for_dit(
-                pipeline,
-                self.config,
-                block_adapter,
-                adapter_cls,
-            )
+            self._refresh_funcs = []
+            for name in _dit_module_names(pipeline):
+                get_transformer = _make_pipeline_transformer_getter(name)
+                block_adapter = _maybe_build_block_adapter(pipeline, get_transformer)
+                adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer)
+                self._refresh_funcs.append(
+                    enable_cache_for_dit(
+                        pipeline,
+                        self.config,
+                        block_adapter,
+                        adapter_cls,
+                        get_transformer,
+                    )
+                )
+            if not self._refresh_funcs:
+                raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT")
 
         self.enabled = True
         logger.info("Cache-dit enabled successfully on %s", pipeline_name)
 
-    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if not self.enabled or self._refresh_func is None:
+    def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None:
+        if not self.enabled or not self._refresh_funcs:
             logger.warning("Cache-dit is not enabled. Cannot refresh cache context.")
             return
 
@@ -211,7 +251,8 @@ class CacheDiTBackend(CacheBackend):
                 "Refreshing cache context for transformer with num_inference_steps: %s",
                 num_inference_steps,
             )
-        self._refresh_func(pipeline, num_inference_steps, verbose)
+        for refresh_func in self._refresh_funcs:
+            refresh_func(pipeline, num_inference_steps, verbose)
 
 
 __all__ = [

vllm_omni/diffusion/data.py

@@ -541,16 +541,15 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
     """
     from vllm.transformers_utils.config import get_hf_file_to_dict
 
+    from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
+
     if not model:
         return None
 
     is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
 
-    # Diffusers models: read _class_name from model_index.json.
-    try:
-        model_index = get_hf_file_to_dict("model_index.json", model)
-    except Exception:
-        model_index = None
+    # Diffusers models: read _class_name from the pipeline index.
+    model_index = get_diffusion_model_index(model)
     if model_index is not None:
         return model_index.get("_class_name")
     if diffusion_load_format == "diffusers":
@@ -603,6 +602,10 @@ class OmniDiffusionConfig:
 
     model_class_name: str | None = None
 
+    # Optional model-defined startup task. Pipelines may use this to select
+    # task-specific components or weights before serving requests.
+    task_type: str | None = None
+
     dtype: torch.dtype = torch.bfloat16
 
     model_config: dict[str, Any] = field(default_factory=dict)
@@ -1143,18 +1146,23 @@ class OmniDiffusionConfig:
     def enrich_config(self) -> None:
         """Load model metadata from HuggingFace and populate config fields.
 
-        Diffusers-style models expose ``model_index.json`` with ``_class_name``.
+        Diffusers-style models expose a pipeline index with ``_class_name``.
         Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
         so we fall back to reading that and mapping model_type manually.
         """
         from vllm.transformers_utils.config import get_hf_file_to_dict
 
+        from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
+
         # Default model_class_name for diffusers adapter
         if self.model_class_name is None and self.diffusion_load_format == "diffusers":
             self.model_class_name = "DiffusersAdapterPipeline"
 
         try:
-            config_dict = get_hf_file_to_dict("model_index.json", self.model)
+            config_dict = get_diffusion_model_index(
+                self.model,
+                revision=self.revision,
+            )
             if config_dict is not None:
                 if self.model_class_name is None:
                     self.model_class_name = config_dict.get("_class_name", None)
@@ -1169,7 +1177,7 @@ class OmniDiffusionConfig:
                         self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
                     except (KeyError, AttributeError) as exc:
                         logger.warning(
-                            "Could not find valid _class_name for diffusers pipeline in model_index.json: %s. "
+                            "Could not find a valid _class_name in the Diffusers pipeline index: %s. "
                             "Without the underlying pipeline class the dummy run may omit required inputs.",
                             exc,
                         )
@@ -1182,16 +1190,16 @@ class OmniDiffusionConfig:
                     else:
                         self.set_tf_model_config(TransformerConfig())
             else:
-                raise FileNotFoundError("model_index.json not found")
+                raise FileNotFoundError("Diffusers pipeline index not found")
         except (AttributeError, OSError, ValueError, FileNotFoundError):
             # Skip transformer config loading for diffusers adapter
             # (non-DiT models don't have a separate transformer folder/config)
             if self.diffusion_load_format == "diffusers":
                 self.set_tf_model_config(TransformerConfig())
                 logger.warning(
-                    "Could not find valid model_index.json per diffusers format. "
+                    "Could not find a valid pipeline index per Diffusers format. "
                     "This model is likely unsupported by the diffusers backend. "
-                    "Also, without knowing the underlying diffusers pipeline class from model_index.json, "
+                    "Also, without knowing the underlying pipeline class from its index, "
                     "the dummy run will input only text prompt, which may cause errors for pipelines "
                     "that require additional inputs."
                 )
@@ -1207,7 +1215,7 @@ class OmniDiffusionConfig:
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
                         return
-                    raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
+                    raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}")
 
                 self.set_tf_model_config(TransformerConfig.from_dict(cfg))
                 model_type = cfg.get("model_type")
@@ -1434,6 +1442,10 @@ class SkipSoftmaxSpec:
         if self.target_sparsity is not None and self.threshold is not None:
             raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.")
 
+    @property
+    def enabled(self) -> bool:
+        return self.threshold is not None or self.target_sparsity is not None
+
 
 @dataclass
 class AttnQuantSpec:
@@ -1461,6 +1473,33 @@ class AttnQuantSpec:
         return self.dtype_qk is not None or self.dtype_vo is not None
 
 
+# Backends that select key blocks instead of attending densely, and so accept
+# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel.
+BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})
+
+
+@dataclass
+class BlockSparseSpec:
+    """User-facing controls shared by block-sparse attention backends.
+
+    ``sparsity`` is the nominal fraction of key blocks dropped per query block;
+    ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an
+    index selector such as "0-3,38") exempts individual DiT blocks. Those two are
+    the accuracy knobs to trade back quality at a fixed ``sparsity``.
+    """
+
+    sparsity: float = 0.8
+    start_step: int = 0
+    skip_layers: str | list[int] | None = None
+    skip_layer_indices: set[int] | None = field(default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0
+        if self.start_step < 0:
+            raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.")
+        self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)
+
+
 @dataclass
 class AttentionSpec:
     """Specifies a backend and its typed backend-specific config for one attention role."""
@@ -1468,6 +1507,7 @@ class AttentionSpec:
     backend: str
     skip_softmax: SkipSoftmaxSpec | None = None
     quant: AttnQuantSpec | None = None
+    block_sparse: BlockSparseSpec | None = None
     skip_calibration: dict | None = field(default=None, repr=False)
 
     def __post_init__(self) -> None:
@@ -1475,6 +1515,7 @@ class AttentionSpec:
             raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
         self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax")
         self.quant = self._coerce(self.quant, AttnQuantSpec, "quant")
+        self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse")
         if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
             raise ValueError(
                 f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
@@ -1485,6 +1526,15 @@ class AttentionSpec:
                 f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
                 f"backend={self.backend!r}. Remove quant or set a supported backend."
             )
+        if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
+            # Selecting the backend is the opt-in; without an explicit block the
+            # defaults apply rather than silently running dense.
+            self.block_sparse = self.block_sparse or BlockSparseSpec()
+        elif self.block_sparse is not None:
+            raise ValueError(
+                f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but "
+                f"backend={self.backend!r}. Remove block_sparse or set a supported backend."
+            )
 
     @staticmethod
     def _coerce(value: Any, cls: type, field_name: str) -> Any:
@@ -1517,6 +1567,12 @@ class AttentionSpec:
             if q.flashinfer_backend is not None:
                 quant_kw["flashinfer_backend"] = q.flashinfer_backend
             kw["quant"] = quant_kw
+        if self.block_sparse is not None:
+            bs = self.block_sparse
+            kw["sparsity"] = bs.sparsity
+            kw["start_step"] = bs.start_step
+            if bs.skip_layer_indices:
+                kw["skip_layers"] = sorted(bs.skip_layer_indices)
         return kw or None
 
 
@@ -1586,7 +1642,7 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "skip_softmax", "quant"}
+        spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys:

vllm_omni/diffusion/forward_context.py

@@ -221,9 +221,15 @@ class DenoiseProgressMixin:
         step_idx: int | None,
         timestep=None,
         scheduler=None,
+        normalized_timestep: float | None = None,
     ) -> None:
         set_forward_context_denoise_step_idx(step_idx)
-        if timestep is None or _forward_context is None:
+        if _forward_context is None:
+            return
+        if normalized_timestep is not None:
+            _forward_context.denoise_timestep = float(normalized_timestep)
+            return
+        if timestep is None:
             return
         scheduler = scheduler if scheduler is not None else getattr(self, "scheduler", None)
         ntt = getattr(getattr(scheduler, "config", None), "num_train_timesteps", None)

vllm_omni/diffusion/lora/manager.py

@@ -366,8 +366,8 @@ class DiffusionLoRAManager:
             fully_sharded_loras=False,
         )
 
-        # Default denoising components plus any a pipeline opts into via
-        # ``_lora_components`` (opt-in; other pipelines unchanged).
+        # Default denoising components, declared DiT components, and any a
+        # pipeline opts into via ``_lora_components``.
         #
         # NOTE: SDXL-style pipelines expose the denoiser as ``unet``.
         # Without scanning this component, adapters can load/activate while
@@ -379,8 +379,10 @@ class DiffusionLoRAManager:
             "bagel",
             "unet",
         )
+        declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ())
         extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ())
-        for component_name in (*default_components, *extra_components):
+        component_names = dict.fromkeys((*default_components, *declared_components, *extra_components))
+        for component_name in component_names:
             if not hasattr(self.pipeline, component_name):
                 continue
             component = getattr(self.pipeline, component_name)

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -393,6 +393,13 @@ class DiffusersPipelineLoader:
                             "DLO+AllGather active but model does not support mmap: loading weights via regular loader."
                         )
                     logger.debug("Loading weights on %s ...", load_device)
+                    if offload_after_quant:
+                        marked = self._request_offload_after_quant(model)
+                        if marked:
+                            logger.info(
+                                "Online quantization will return each of %d layers to CPU as it is quantized",
+                                marked,
+                            )
                     if load_format == "diffusers":
                         cast(DiffusersAdapterPipeline, model).load_weights()
                     else:
@@ -406,6 +413,29 @@ class DiffusersPipelineLoader:
         self._apply_skip_softmax_calibration(model)
         return model.eval()
 
+    @staticmethod
+    def _request_offload_after_quant(model: nn.Module) -> int:
+        """Ask online-quant layers to return to host memory once quantized.
+
+        The weights only visit the accelerator so the quant kernels can run on
+        them; ``load_model`` sends the model back to the host afterwards either
+        way. Without this the whole transformer accumulates on device until that
+        final move, which for MiniMax H3 is a ~43 GiB peak that no longer fits
+        beside a resident TP-sharded text encoder — even though layer-wise
+        offload means none of it is supposed to be resident at inference time.
+
+        Only quant methods that advertise ``supports_offload_after_quant`` are
+        asked, since the implementation has to know when a layer is finished.
+        Deferring materialization to the ``meta`` device does not imply that.
+        """
+        marked = 0
+        for module in model.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if getattr(quant_method, "supports_offload_after_quant", False):
+                quant_method.enable_offload_after_quant()
+                marked += 1
+        return marked
+
     @staticmethod
     def _has_online_quant(model: nn.Module) -> bool:
         """Whether any layer uses an online-quant method that defers weight
@@ -460,6 +490,13 @@ class DiffusersPipelineLoader:
             if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
                 continue
 
+            # Layers finished during loading would only be staged onto the target
+            # device for a process call that immediately returns. That round trip
+            # is wasted work in general, and undoes the point of the offload for
+            # layers that already went back to the host.
+            if getattr(module, "_already_called_process_weights_after_loading", False):
+                continue
+
             if has_online_quant:
                 # Online quant may leave straggler params on the ``meta`` device.
                 # Move only real (non-meta) params onto the target device for

vllm_omni/diffusion/model_metadata.py

@@ -41,6 +41,16 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=9,
         supports_mixed_reference_inputs=True,
+        # H3 represents alignment padding as a second packed sequence.  The
+        # packed TRTLLM backend consumes cu_seqlens and isolates that padding.
+        attention_mask_free=True,
+    ),
+    # The modular alias is served by MiniMaxH3Pipeline and has the same
+    # Ref2VA request contract. Keep admission limits in sync with it.
+    "MiniMaxH3ModularPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=9,
+        supports_mixed_reference_inputs=True,
     ),
     "WanPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -14,9 +14,13 @@ from typing import Any
 
 import torch
 
+from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
+from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
+    minimax_h3_validate_finite_tensors_enabled,
 )
 
 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
@@ -87,6 +91,13 @@ class MiniMaxH3DenoiseBranch:
                 "max_seqlen_q": text_len,
             },
         }
+        # Where the video segment sits in the packed sequence. Resolved to plain
+        # ints here so the attention layers never sync on it per step.
+        grid = packed["latent_grid"].tolist()
+        self.static_kwargs["video_token_layout"] = VideoTokenLayout(
+            prefix_len=int(packed["video_row_start"]),
+            latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])),
+        )
 
     def forward_kwargs(
         self,
@@ -139,7 +150,8 @@ def minimax_h3_denoise_loop(
     device: torch.device,
     imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP,
     audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
-    on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
+    on_step_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
+    on_step_start: Callable[[int, float, float], None] | None = None,
     step_profiler: Callable[[int], AbstractContextManager] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Run the full denoise loop; returns final (video_rows, audio_rows).
@@ -187,12 +199,19 @@ def minimax_h3_denoise_loop(
     if audio_anchor is not None:
         audio_rows[~audio_update] = audio_anchor
 
+    validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled()
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
+            # Publish the step index so step-gated attention features (e.g. the
+            # dense warmup of RAINFUSION_ATTN) can see where we are.
+            set_forward_context_denoise_step_idx(step)
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
             s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
+            if on_step_start is not None:
+                # Attention gates use the scheduler-style descending timestep.
+                on_step_start(step, s_v, s_a)
             t_v, t_a = 1.0 - s_v, 1.0 - s_a
             imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference))
             audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference))
@@ -214,8 +233,15 @@ def minimax_h3_denoise_loop(
                 video_rows[update],
                 mv_video_t,
                 torch.tensor(t_v, dtype=torch.float32, device=device),
+                validate_finite_tensors=validate_finite_tensors,
+            )
+            new_target = minimax_h3_euler_eta0_step(
+                video_rows[update],
+                x0_video,
+                sigma_curr=s_v,
+                sigma_next=s_v_next,
+                validate_finite_tensors=validate_finite_tensors,
             )
-            new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next)
             video_rows = video_rows.clone()
             video_rows[update] = new_target
             if cond_anchor is not None:
@@ -225,17 +251,23 @@ def minimax_h3_denoise_loop(
                 audio_rows[audio_update],
                 mv_audio_t,
                 torch.tensor(t_a, dtype=torch.float32, device=device),
+                validate_finite_tensors=validate_finite_tensors,
             )
             new_audio = minimax_h3_euler_eta0_step(
-                audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next
+                audio_rows[audio_update],
+                x0_audio,
+                sigma_curr=s_a,
+                sigma_next=s_a_next,
+                validate_finite_tensors=validate_finite_tensors,
             )
             audio_rows = audio_rows.clone()
             audio_rows[audio_update] = new_audio
             if audio_anchor is not None:
                 audio_rows[~audio_update] = audio_anchor  # per-step audio ref reset
-            if on_step is not None:
-                on_step(step, video_rows, audio_rows)
+            if on_step_end is not None:
+                on_step_end(step, video_rows, audio_rows)
 
+    set_forward_context_denoise_step_idx(None)
     return video_rows, audio_rows
 
 

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
@@ -134,6 +134,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset(
         "img_pos_for_infer_output_info",
         "packed_seq_params",
         "refiner_packed_seq_params",
+        "video_token_layout",
     }
 )
 
@@ -310,6 +311,8 @@ class MiniMaxH3Attention(nn.Module):
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
+        role: str = "self",
+        role_category: str | None = None,
         skip_sequence_parallel: bool = False,
     ) -> None:
         super().__init__()
@@ -348,7 +351,12 @@ class MiniMaxH3Attention(nn.Module):
             head_size=self.head_dim,
             softmax_scale=self.softmax_scale,
             causal=False,
+            # Packed rows reach the impl as [B, S, N, D].
+            qkv_layout="BSND",
+            role=role,
+            role_category=role_category,
             skip_sequence_parallel=skip_sequence_parallel,
+            prefix=prefix,
         )
 
     def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
@@ -374,6 +382,7 @@ class MiniMaxH3Attention(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """Run packed attention as a small eager island.
 
@@ -407,6 +416,7 @@ class MiniMaxH3Attention(nn.Module):
                 "max_seqlen_k": max_seqlen,
                 "valid_kv_length": used,
             },
+            video_layout=video_layout,
         )
         return self.attention(
             q.unsqueeze(0),
@@ -424,6 +434,7 @@ class MiniMaxH3Attention(nn.Module):
         max_seqlen: int,
         packed_total: int | None = None,
         sp_seq_lens: list[int] | None = None,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """x: [T, hidden] packed thd rows -> [T, hidden].
 
@@ -463,6 +474,7 @@ class MiniMaxH3Attention(nn.Module):
             # backend receives the global sequence after all-to-all, so carry
             # its Python length explicitly instead of inferring it from q.
             packed_total=packed_total if packed_total is not None else q.shape[0],
+            video_layout=video_layout,
         )
         out = out.reshape(total, self.num_heads * self.head_dim)
         out, _ = self.out_proj(out)
@@ -573,6 +585,8 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
             arch,
             quant_config,
             prefix=f"{prefix}.attn",
+            role="minimax_h3.token_refiner",
+            role_category="self",
             skip_sequence_parallel=True,
         )
         self.mlp = MiniMaxH3MLP(
@@ -642,6 +656,8 @@ class MiniMaxH3DiTBlock(nn.Module):
         super().__init__()
         self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
         self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps)
+        # The prefix also carries the block index that block-sparse attention
+        # backends match against their skip_layers selector.
         self.attn = MiniMaxH3Attention(
             arch,
             quant_config,
@@ -672,6 +688,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         max_seqlen: int,
         packed_total: int,
         sp_seq_lens: list[int] | None = None,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T]
         (= inverse_indices * modality_num + token_tags.clamp(min=0)).
@@ -699,6 +716,7 @@ class MiniMaxH3DiTBlock(nn.Module):
             max_seqlen=max_seqlen,
             packed_total=packed_total,
             sp_seq_lens=sp_seq_lens,
+            video_layout=video_layout,
         )
         x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE)
 
@@ -831,6 +849,10 @@ class MiniMaxH3DiTModel(nn.Module):
         },
         "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2),
     }
+    # The checkpoint already stores qkv and the MLP gate/up as single tensors
+    # (see the reordering in load_weights), so there are no unfused names for
+    # quantization or LoRA to map onto. Address the fused layers directly, e.g.
+    # ignored_layers=["blocks.0.attn.qkv_proj"].
     packed_modules_mapping = {}
 
     def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None:
@@ -1097,6 +1119,7 @@ class MiniMaxH3DiTModel(nn.Module):
         refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
         refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
         refiner_max = int(self._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)}")
@@ -1145,6 +1168,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,
+                video_layout=video_layout,
             )
         hidden = self.sp_gather(hidden)
 

vllm_omni/diffusion/models/minimax_h3/packed_sequence.py

@@ -236,6 +236,9 @@ def minimax_h3_packed_sequence(
         "token_tags": token_tags,
         "cu_seqlens": cu,
         "document_id": doc,
+        # Extent of the video segment within the packed sequence.
+        "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
+        "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
     }
 
 
@@ -554,6 +557,9 @@ def minimax_h3_packed_sequence_ref2va_blocks(
         "token_tags": token_tags,
         "cu_seqlens": cu,
         "document_id": doc,
+        # Extent of the video segment within the packed sequence.
+        "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
+        "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
     }
 
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -9,6 +9,7 @@ import os
 import tempfile
 from collections.abc import Iterable, Mapping, Sequence
 from contextlib import contextmanager
+from itertools import groupby
 from pathlib import Path
 from typing import Any, ClassVar
 
@@ -27,6 +28,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_world_group,
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import (
     DiffusersPipelineLoader,
 )
@@ -43,6 +45,9 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
 )
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.errors import OmniClientError
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
 from vllm_omni.platforms import current_omni_platform
 
 from .condition_noise import (
@@ -106,6 +111,57 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024
 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"})
 MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0
 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0
+MINIMAX_H3_DOWNLOAD_PATTERNS = [
+    "FL2VA/**",
+    "Ref2VA/model_index.json",
+    "Ref2VA/transformer/**",
+]
+MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = {
+    "fl2va": ["FL2VA/**"],
+    "ref2va": ["Ref2VA/**"],
+}
+
+
+def _minimax_h3_partition_for_task(
+    task_type: str | None,
+    model: str | None = None,
+) -> str:
+    task = str(task_type or "auto").lower()
+    if task == "auto" and model is not None:
+        path = Path(model)
+        if path.is_dir() and path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file():
+            return path.name.lower()
+    if task in {"auto", "combined"}:
+        return "combined"
+    if task in {"t2va", "fl2va"}:
+        return "fl2va"
+    if task == "ref2va":
+        return "ref2va"
+    raise ValueError(f"MiniMax-H3 task_type must be one of auto, t2va, fl2va, or ref2va; got {task_type!r}")
+
+
+def _resolve_minimax_h3_model_root(
+    model: str,
+    revision: str | None,
+    partition: str,
+) -> Path:
+    path = Path(model)
+    if path.is_dir():
+        if path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file():
+            return path.parent
+        return path
+    allow_patterns = (
+        MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
+    )
+    return Path(
+        download_weights_from_hf_specific(
+            model_name_or_path=model,
+            cache_dir=None,
+            allow_patterns=allow_patterns,
+            revision=revision,
+            require_all=True,
+        )
+    )
 
 
 def _resolve_component_quant_config(quant_config, component: str):
@@ -446,6 +502,7 @@ class _SingleRankEncoderGroup:
 
 class MiniMaxH3Pipeline(
     nn.Module,
+    DenoiseProgressMixin,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
     SupportImageInput,
@@ -455,11 +512,7 @@ class MiniMaxH3Pipeline(
 ):
     """CFG-distilled joint video/audio generation for MiniMax H3."""
 
-    _dit_modules: ClassVar[list[str]] = ["transformer"]
-    # TODO(offload): Reconcile this static plan with the modular H3 pipeline
-    # (#5720). Its combined mode has ``transformer`` and ``transformers_ref``;
-    # build the resident/staging plan from the selected task partition so a
-    # consumer-GPU single-task server never stages both DiTs.
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"]
     _offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
@@ -498,24 +551,64 @@ class MiniMaxH3Pipeline(
         if int(self.parallel_config.cfg_parallel_size) != 1:
             raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1")
         self.device = get_local_device()
-        model_path = str(od_config.model)
-        model_index = json.loads((Path(model_path) / "model_index.json").read_text(encoding="utf-8"))
+        self.partition = _minimax_h3_partition_for_task(
+            getattr(od_config, "task_type", None),
+            str(od_config.model),
+        )
+        model_root = _resolve_minimax_h3_model_root(
+            str(od_config.model),
+            od_config.revision,
+            self.partition,
+        )
+        model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
+        model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
         release = model_index.get("_minimax_h3") or {}
-        self.partition = str(release.get("partition", ""))
-        self.supported_tasks = frozenset(release.get("tasks") or ())
+        partition = str(release.get("partition", "")).lower()
+        expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va"
+        if partition != expected_partition:
+            raise ValueError(f"invalid MiniMax-H3 {expected_partition} partition at {model_path}")
+
+        supported_tasks = {str(task).lower() for task in release.get("tasks", [])}
+        if not supported_tasks:
+            supported_tasks = {"ref2va"} if partition == "ref2va" else {"t2va", "fl2va"}
+        ref2va_model_path = None
+        if self.partition == "combined":
+            ref2va_model_path = model_root / "Ref2VA"
+            ref2va_index_path = ref2va_model_path / "model_index.json"
+            if not ref2va_index_path.is_file():
+                raise ValueError(f"Ref2VA partition not found at {ref2va_model_path}")
+            ref2va_index = json.loads(ref2va_index_path.read_text(encoding="utf-8"))
+            ref2va_release = ref2va_index.get("_minimax_h3") or {}
+            if str(ref2va_release.get("partition", "")).lower() != "ref2va":
+                raise ValueError(f"invalid MiniMax-H3 ref2va partition at {ref2va_model_path}")
+            supported_tasks.update(str(task).lower() for task in ref2va_release.get("tasks", ["ref2va"]))
+
+        self.supported_tasks = frozenset(supported_tasks)
         shifts = release.get("sigma_shift_scales") or {}
         self.default_video_shift = float(shifts.get("video", 12.0))
         self.default_audio_shift = float(shifts.get("audio", 3.0))
 
         self.weights_sources = [
             DiffusersPipelineLoader.ComponentSource(
-                model_or_path=model_path,
+                model_or_path=str(model_path),
                 subfolder="transformer",
                 revision=od_config.revision,
                 prefix="transformer.",
                 fall_back_to_pt=False,
             )
         ]
+        self._dit_modules = ["transformer"]
+        if ref2va_model_path is not None:
+            self.weights_sources.append(
+                DiffusersPipelineLoader.ComponentSource(
+                    model_or_path=str(ref2va_model_path),
+                    subfolder="transformer",
+                    revision=od_config.revision,
+                    prefix="transformers_ref.",
+                    fall_back_to_pt=False,
+                )
+            )
+            self._dit_modules.append("transformers_ref")
         transformer_quant_config = _resolve_component_quant_config(
             od_config.quantization_config,
             "transformer",
@@ -524,14 +617,19 @@ class MiniMaxH3Pipeline(
             od_config,
             quant_config=transformer_quant_config,
         )
+        if ref2va_model_path is not None:
+            self.transformers_ref = MiniMaxH3DiTModel(
+                od_config,
+                quant_config=transformer_quant_config,
+            )
 
         self.tokenizer = Qwen2TokenizerFast.from_pretrained(
-            model_path,
+            str(model_path),
             subfolder="tokenizer",
             local_files_only=os.path.isdir(model_path),
         )
         self.processor = Qwen3VLProcessor.from_pretrained(
-            model_path,
+            str(model_path),
             subfolder="processor",
             local_files_only=os.path.isdir(model_path),
         )
@@ -586,16 +684,23 @@ class MiniMaxH3Pipeline(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
     ) -> set[str]:
-        prefix = "transformer."
+        def source_prefix(item: tuple[str, torch.Tensor]) -> str:
+            name, _ = item
+            prefix = name.partition(".")[0] + "."
+            if prefix in {"transformer.", "transformers_ref."}:
+                return prefix
+            raise ValueError(f"unexpected MiniMax-H3 weight {name!r}")
 
-        def transformer_weights():
-            for name, tensor in weights:
-                if name.startswith(prefix):
-                    yield name[len(prefix) :], tensor
-
-        loaded = self.transformer.load_weights(transformer_weights())
-        self.transformer.post_load_weights()
-        loaded_with_prefix = {prefix + name for name in loaded}
+        loaded_with_prefix: set[str] = set()
+        loaded_prefixes: set[str] = set()
+        for prefix, grouped_weights in groupby(weights, key=source_prefix):
+            if prefix in loaded_prefixes:
+                raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
+            loaded_prefixes.add(prefix)
+            transformer = getattr(self, prefix.removesuffix("."))
+            loaded = transformer.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
+            transformer.post_load_weights()
+            loaded_with_prefix.update(prefix + name for name in loaded)
         # The text encoder and both VAEs load eagerly in ``__init__`` rather
         # than through ``weights_sources``. Record them for the runner's strict
         # missing-parameter check.
@@ -604,14 +709,23 @@ class MiniMaxH3Pipeline(
             loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
         return loaded_with_prefix
 
+    def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel:
+        if task == "ref2va" and hasattr(self, "transformers_ref"):
+            return self.transformers_ref
+        return self.transformer
+
     def _resolve_task(
         self,
         requested: str | None,
         multi_modal_data: dict[str, Any],
     ) -> str:
         if requested is None:
+            # A Ref2VA-only startup has no FL2VA transformer; preserve its
+            # historical implicit default even for image-only references.
             if self.partition == "ref2va":
                 requested = "ref2va"
+            elif multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None:
+                requested = "ref2va"
             elif multi_modal_data.get("image") is not None:
                 requested = "fl2va"
             else:
@@ -1226,14 +1340,14 @@ class MiniMaxH3Pipeline(
         return video_rows, audio_rows
 
     @contextmanager
-    def _resident_dit_layers_on_device(self):
+    def _resident_dit_layers_on_device(self, *, enabled: bool = True):
         controller = getattr(self, "_dlo_residency_controller", None)
-        if controller is not None:
+        if controller is not None and enabled:
             controller.load_resident_layers()
         try:
             yield
         finally:
-            if controller is not None:
+            if controller is not None and enabled:
                 controller.offload_resident_layers()
 
     def diffuse(
@@ -1357,10 +1471,15 @@ class MiniMaxH3Pipeline(
             num_steps=num_steps,
             shift_scale=audio_shift,
         )
-        with self._resident_dit_layers_on_device():
+        transformer = self._transformer_for_task(task)
+        # The static DLO plan keeps leading blocks resident only for the
+        # primary ``transformer``. In combined mode ``transformers_ref`` is
+        # fully streamed, so a Ref2VA request must not stage the inactive
+        # FL2VA transformer resident blocks.
+        with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
             with self.progress_bar(total=len(video_sigmas) - 1) as progress:
                 video_rows, audio_rows = minimax_h3_denoise_loop(
-                    model=self.transformer,
+                    model=transformer,
                     positive=branch,
                     initial_video_rows=initial_video,
                     initial_audio_rows=initial_audio,
@@ -1371,7 +1490,11 @@ class MiniMaxH3Pipeline(
                     device=self.device,
                     imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
                     audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
-                    on_step=lambda step, video, audio: progress.update(),
+                    on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step(
+                        step,
+                        normalized_timestep=video_sigma,
+                    ),
+                    on_step_end=lambda step, video, audio: progress.update(),
                 )
 
         target_video = video_rows[branch.update_mask_dev]

vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py

@@ -2,10 +2,24 @@
 from __future__ import annotations
 
 import math
+import os
 from typing import Any
 
 import torch
 
+MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV = "VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS"
+
+
+def minimax_h3_validate_finite_tensors_enabled() -> bool:
+    value = os.getenv(MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV, "")
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _resolve_validate_finite_tensors(value: bool | None) -> bool:
+    if value is None:
+        return minimax_h3_validate_finite_tensors_enabled()
+    return bool(value)
+
 
 def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None:
     if not bool(torch.isfinite(tensor).all().item()):
@@ -50,6 +64,8 @@ def minimax_h3_rf_v_to_x0(
     xt: torch.Tensor,
     v: torch.Tensor,
     timestep: torch.Tensor,
+    *,
+    validate_finite_tensors: bool | None = None,
 ) -> torch.Tensor:
     if xt.shape != v.shape:
         raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}")
@@ -57,15 +73,18 @@ def minimax_h3_rf_v_to_x0(
         raise ValueError("xt must be a floating point tensor")
     if not torch.is_floating_point(v):
         raise ValueError("v must be a floating point tensor")
-    _require_finite_tensor(xt, "xt")
-    _require_finite_tensor(v, "v")
+    validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors)
+    if validate_finite_tensors:
+        _require_finite_tensor(xt, "xt")
+        _require_finite_tensor(v, "v")
     _validate_unit_timestep(timestep, "timestep")
     cond_t = timestep.to(device=xt.device, dtype=xt.dtype)
     while cond_t.ndim < xt.ndim:
         cond_t = cond_t.unsqueeze(-1)
     sigma_t = 1 - cond_t
     x0 = xt + sigma_t * v
-    _require_finite_tensor(x0, "x0")
+    if validate_finite_tensors:
+        _require_finite_tensor(x0, "x0")
     return x0
 
 
@@ -75,6 +94,7 @@ def minimax_h3_euler_eta0_step(
     *,
     sigma_curr: float,
     sigma_next: float,
+    validate_finite_tensors: bool | None = None,
 ) -> torch.Tensor:
     if state.shape != denoised.shape:
         raise ValueError(f"state and denoised shapes must match, got {state.shape} vs {denoised.shape}")
@@ -82,8 +102,10 @@ def minimax_h3_euler_eta0_step(
         raise ValueError("state must be a floating point tensor")
     if not torch.is_floating_point(denoised):
         raise ValueError("denoised must be a floating point tensor")
-    _require_finite_tensor(state, "state")
-    _require_finite_tensor(denoised, "denoised")
+    validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors)
+    if validate_finite_tensors:
+        _require_finite_tensor(state, "state")
+        _require_finite_tensor(denoised, "denoised")
     sigma_curr = _validate_sigma(sigma_curr, "sigma_curr")
     sigma_next = _validate_sigma(sigma_next, "sigma_next")
     if sigma_curr == 0.0:
@@ -98,7 +120,8 @@ def minimax_h3_euler_eta0_step(
     sigma_ratio = sigma_next_t / sigma_curr_t
     out = sigma_ratio * state.to(dtype=compute_dtype) + (1.0 - sigma_ratio) * denoised.to(dtype=compute_dtype)
     out = out.to(dtype=state.dtype)
-    _require_finite_tensor(out, "euler_eta0_step output")
+    if validate_finite_tensors:
+        _require_finite_tensor(out, "euler_eta0_step output")
     return out
 
 
@@ -127,6 +150,7 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
         audio_sigma_curr: float | None = None,
         audio_sigma_next: float | None = None,
     ) -> dict[str, torch.Tensor]:
+        validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled()
         visual_timestep = timestep if video_timestep is None else video_timestep
         audio_timestep = timestep if audio_timestep is None else audio_timestep
         visual_sigma_curr = sigma_curr if video_sigma_curr is None else video_sigma_curr
@@ -148,11 +172,13 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
             input_visual_latent,
             noise_pred_visual,
             visual_timestep,
+            validate_finite_tensors=validate_finite_tensors,
         )
         denoised_audio = minimax_h3_rf_v_to_x0(
             input_audio_latent,
             noise_pred_audio,
             audio_timestep,
+            validate_finite_tensors=validate_finite_tensors,
         )
         return {
             "output_visual_latent": minimax_h3_euler_eta0_step(
@@ -160,12 +186,14 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
                 denoised_visual,
                 sigma_curr=visual_sigma_curr,
                 sigma_next=visual_sigma_next,
+                validate_finite_tensors=validate_finite_tensors,
             ),
             "output_audio_latent": minimax_h3_euler_eta0_step(
                 input_audio_latent,
                 denoised_audio,
                 sigma_curr=audio_sigma_curr,
                 sigma_next=audio_sigma_next,
+                validate_finite_tensors=validate_finite_tensors,
             ),
         }
 
@@ -173,7 +201,9 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
 EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter
 
 __all__ = [
+    "MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV",
     "MiniMaxH3EulerAncestralEta0SchedulerAdapter",
     "minimax_h3_euler_eta0_step",
     "minimax_h3_rf_v_to_x0",
+    "minimax_h3_validate_finite_tensors_enabled",
 ]

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -51,7 +51,14 @@ def _load_remote_component(
         class_reference,
         component_path,
     )
-    return component_cls.from_pretrained(component_path)
+    # Build on the host regardless of the ambient default device. Online
+    # quantization wraps pipeline construction in a `with torch.device(<accel>)`
+    # block for the DiT's quantized linears, and the checkpoint's own VAE code
+    # builds constants with ops that have no accelerator kernel (BigVGAN's
+    # anti-aliasing filters call torch.kaiser_window). Callers place the module
+    # explicitly right after this returns, so nothing depends on the context.
+    with torch.device("cpu"):
+        return component_cls.from_pretrained(component_path)
 
 
 class _AudioVAEDeterminismContext(AbstractContextManager):

vllm_omni/diffusion/registry.py

@@ -91,6 +91,11 @@ _DIFFUSION_MODELS = {
         "pipeline_minimax_h3",
         "MiniMaxH3Pipeline",
     ),
+    "MiniMaxH3ModularPipeline": (
+        "minimax_h3",
+        "pipeline_minimax_h3",
+        "MiniMaxH3Pipeline",
+    ),
     "StableAudioPipeline": (
         "stable_audio",
         "pipeline_stable_audio",
@@ -427,9 +432,11 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
         if sp_size <= 1:
             return
 
-        # Find transformer model(s) in the pipeline that have _sp_plan
-        # Include transformer_2 for two-stage models (e.g., Wan MoE)
-        transformer_attrs = ["transformer", "transformer_2", "dit", "unet"]
+        # Prefer the pipeline's declared DiT components so custom component
+        # names receive the same SP hooks as conventional transformer names.
+        transformer_attrs = getattr(model, "_dit_modules", None)
+        if not transformer_attrs:
+            transformer_attrs = ("transformer", "transformer_2", "dit", "unet")
         applied_count = 0
 
         for attr in transformer_attrs:
@@ -508,6 +515,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
     "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
     "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func",
+    "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func",
     "StableAudioPipeline": "get_stable_audio_post_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_post_process_func",
     "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func",

vllm_omni/diffusion/utils/hf_utils.py

@@ -7,6 +7,24 @@ from vllm.transformers_utils.config import get_hf_file_to_dict
 
 logger = init_logger(__name__)
 
+DIFFUSION_MODEL_INDEX_FILES = (
+    "model_index.json",
+    "modular_model_index.json",
+)
+
+
+def get_diffusion_model_index(
+    model_name: str,
+    *,
+    revision: str | None = None,
+) -> dict | None:
+    """Read the first standard Diffusers pipeline index available."""
+    for filename in DIFFUSION_MODEL_INDEX_FILES:
+        config = get_hf_file_to_dict(filename, model_name, revision=revision)
+        if isinstance(config, Mapping):
+            return dict(config)
+    return None
+
 
 def load_diffusers_config(model_name) -> dict:
     from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -58,33 +76,32 @@ def is_diffusion_model(model_name: str) -> bool:
     """Check if a model is a diffusion model.
 
     Uses multiple fallback strategies to detect diffusion models:
-    1. Check local file system for model_index.json (fastest, no imports)
+    1. Check the local file system for a standard Diffusers index
     2. Check using vllm's get_hf_file_to_dict utility
     3. Try the standard diffusers approach (may fail due to import issues)
     """
     # Strategy 1: Check local file system first (fastest, avoids import issues)
     if os.path.isdir(model_name):
-        model_index_path = os.path.join(model_name, "model_index.json")
-        if os.path.exists(model_index_path):
+        for filename in DIFFUSION_MODEL_INDEX_FILES:
+            model_index_path = os.path.join(model_name, filename)
+            if not os.path.exists(model_index_path):
+                continue
             try:
                 import json
 
                 with open(model_index_path) as f:
                     config_dict = json.load(f)
                 if config_dict.get("_class_name") and config_dict.get("_diffusers_version"):
-                    logger.debug("Detected diffusion model via local model_index.json")
+                    logger.debug("Detected diffusion model via local %s", filename)
                     return True
             except Exception as e:
-                logger.debug("Failed to read local model_index.json: %s", e)
+                logger.debug("Failed to read local %s: %s", filename, e)
 
     # Strategy 2: Check using vllm's utility (works for both local and remote models)
-    try:
-        config_dict = get_hf_file_to_dict("model_index.json", model_name)
-        if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"):
-            logger.debug("Detected diffusion model via model_index.json")
-            return True
-    except Exception as e:
-        logger.debug("Failed to check model_index.json via get_hf_file_to_dict: %s", e)
+    config_dict = get_diffusion_model_index(model_name)
+    if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"):
+        logger.debug("Detected diffusion model via a standard Diffusers index")
+        return True
 
     # Strategy 3: Try the standard diffusers approach (may fail due to import issues)
     # This is last because it requires importing diffusers/xformers/flash_attn

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -280,8 +280,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             exc,
                         )
                 else:
-                    self._compile_transformer("transformer")
-                    self._compile_transformer("transformer_2")
+                    transformer_attrs = getattr(self.pipeline, "_dit_modules", None)
+                    if not transformer_attrs:
+                        transformer_attrs = ("transformer", "transformer_2")
+                    for attr_name in transformer_attrs:
+                        self._compile_transformer(attr_name)
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",

vllm_omni/engine/arg_utils.py

@@ -140,8 +140,9 @@ class OmniEngineArgs(EngineArgs):
         stage_connector_spec: Extra configuration for stage connector
         async_chunk: If set to True, perform async chunk
         worker_type: Model Type, e.g., "ar" or "generation"
-        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
-            If not specified, will be inferred from model path.
+        task_type: Model-defined startup task type. Consumers validate the
+            supported values and decide whether it selects request behavior,
+            task-specific weights, or both.
         omni_master_address: TCP address that the OmniMasterServer (running
             inside AsyncOmniEngine) listens on for engine core registrations.
             Required when single-stage mode is active.

vllm_omni/engine/async_omni_engine.py

@@ -1019,6 +1019,7 @@ class AsyncOmniEngine:
             "max_num_seqs": kwargs.get("max_num_seqs") or 1,
             "parallel_config": parallel_config,
             "model_class_name": kwargs.get("model_class_name", None),
+            "task_type": kwargs.get("task_type", None),
             "model_config": kwargs.get("model_config", None),
             "additional_config": kwargs.get("additional_config", None),
             "step_execution": kwargs.get("step_execution", False),

vllm_omni/engine/stage_init_utils.py

@@ -8,6 +8,7 @@ out of StageEngineCoreClient into reusable functions.
 
 from __future__ import annotations
 
+import copy
 import fcntl
 import importlib
 import multiprocessing as mp
@@ -15,7 +16,7 @@ import os
 import time
 from collections.abc import Callable, Generator, Sequence
 from contextlib import contextmanager
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, fields, replace
 from typing import Any, Literal, cast
 
 from vllm.logger import init_logger
@@ -25,7 +26,10 @@ from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
 
-from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
+from vllm_omni.config.omni_config import (
+    BaseVllmOmniStageConfig,
+    VllmOmniDiffusionStageConfig,
+)
 from vllm_omni.config.stage_config import StageType
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
@@ -348,6 +352,30 @@ class StageMetadata:
     replica_id: int = 0
 
 
+def _apply_rocm_attention_backend(
+    engine_args: dict[str, Any],
+    stage_type: str | StageType,
+) -> None:
+    """Preserve Omni's ROCm attention-backend compatibility default."""
+    if (
+        not current_omni_platform.is_rocm()
+        or stage_type == StageType.DIFFUSION
+        or engine_args.get("attention_backend") is not None
+    ):
+        return
+
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    if rocm_aiter_ops.is_enabled():
+        engine_args["attention_backend"] = "ROCM_AITER_FA"
+    # Before vLLM v0.19.0, the default attention backend is TRITON_ATTN for ROCm.
+    # Since vLLM v0.19.0, the default attention backend is ROCM_ATTN for ROCm.
+    # However, the compatibility of ROCM_ATTN with Omni is not guaranteed.
+    # Therefore, we still use TRITON_ATTN as the default attention backend,
+    # when the selected_backend is not specified.
+    engine_args["attention_backend"] = "TRITON_ATTN"
+
+
 def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
     """Extract metadata through the active production legacy path.
 
@@ -358,18 +386,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
     stage_type: Literal["llm", "diffusion"] = _get_attr_or_item(stage_config, "stage_type", "llm")
     engine_args = stage_config.engine_args
 
-    if current_omni_platform.is_rocm():
-        if stage_type != "diffusion" and engine_args.get("attention_backend") is None:
-            from vllm._aiter_ops import rocm_aiter_ops
-
-            if rocm_aiter_ops.is_enabled():
-                engine_args["attention_backend"] = "ROCM_AITER_FA"
-            # Before vLLM v0.19.0, the default attention backend is TRITON_ATTN for ROCm.
-            # Since vLLM v0.19.0, the default attention backend is ROCM_ATTN for ROCm.
-            # However, the compatibility of ROCM_ATTN with Omni is not guaranteed.
-            # Therefore, we still use TRITON_ATTN as the default attention backend,
-            # when the selected_backend is not specified.
-            engine_args["attention_backend"] = "TRITON_ATTN"
+    _apply_rocm_attention_backend(engine_args, stage_type)
 
     runtime_cfg = stage_config.runtime
     engine_input_source: list[int] = _get_attr_or_item(stage_config, "engine_input_source", [])
@@ -748,25 +765,110 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None,
                 os.environ[key] = old_value
 
 
-def build_engine_args_dict(
-    stage_config: Any,
-    model: str,
-    stage_connector_spec: dict[str, Any] | None = None,
-    cli_tokenizer: str | None = None,
+def _project_omni_config_fields(
+    config: Any,
+    *,
+    exclude: frozenset[str] = frozenset(),
 ) -> dict[str, Any]:
-    """Build the normalized engine args dict for one stage."""
-    engine_args = stage_config.engine_args
-    # HACK (Alex) Tensor parallel size should not be passed as None;
-    # remove it if this is the case so that we fall back to default
-    # creation from vLLM's engine args.
-    # NOTE: This will be fixed more generically in ongoing work for engine arg filtering.
-    if "tensor_parallel_size" in engine_args and engine_args["tensor_parallel_size"] is None:
-        del engine_args["tensor_parallel_size"]
+    """Copy defined typed config fields into backend adapter kwargs."""
+    projected: dict[str, Any] = {}
+    for config_field in fields(config):
+        name = config_field.name
+        if name in exclude:
+            continue
+        value = getattr(config, name)
+        if value is not None:
+            projected[name] = copy.deepcopy(value)
+    return projected
 
-    stage_type = _get_attr_or_item(stage_config, "stage_type", "llm")
-    stage_id = stage_config.stage_id
 
-    engine_args_dict = _to_dict(engine_args)
+def _project_omni_stage_engine_args(
+    stage_config: BaseVllmOmniStageConfig,
+) -> dict[str, Any]:
+    """Read backend inputs from one structured stage config."""
+    engine_args: dict[str, Any] = {}
+    is_diffusion = isinstance(stage_config, VllmOmniDiffusionStageConfig)
+
+    if is_diffusion:
+        engine_args.update(_project_omni_config_fields(stage_config.diffusion_config))
+
+    for config, excluded_fields in (
+        (
+            stage_config.model_config,
+            frozenset({"default_sampling_params", "has_sampling_extra_args"}),
+        ),
+        (stage_config.load_config, frozenset()),
+        (stage_config.cache_config, frozenset()),
+        (stage_config.scheduler_config, frozenset()),
+        (
+            stage_config.runtime_config,
+            frozenset({"devices", "num_replicas", "env", "num_gpus"}),
+        ),
+    ):
+        engine_args.update(
+            _project_omni_config_fields(
+                config,
+                exclude=excluded_fields,
+            )
+        )
+
+    # The legacy builder always emits this key, including for pipelines such
+    # as Audex that intentionally defer architecture discovery to HF config.
+    engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch)
+
+    topology = stage_config.stage_pipeline_config
+    topology_engine_args = {
+        "model_stage": stage_config.model_stage,
+        "worker_type": stage_config.worker_type,
+        "scheduler_cls": stage_config.scheduler_cls,
+        "hf_config_name": stage_config.hf_config_name,
+        "engine_output_type": stage_config.engine_output_type,
+        "custom_process_next_stage_input_func": stage_config.custom_process_next_stage_input_func,
+        "retains_state_across_chunks": topology.retains_state_across_chunks,
+    }
+    engine_args.update(
+        {name: copy.deepcopy(value) for name, value in topology_engine_args.items() if value is not None}
+    )
+
+    connector_config = stage_config.connector_config
+    engine_args["async_chunk"] = connector_config.async_chunk
+    if connector_config.omni_kv_config is not None:
+        engine_args["omni_kv_config"] = copy.deepcopy(connector_config.omni_kv_config)
+
+    if is_diffusion:
+        engine_args["parallel_config"] = _project_omni_config_fields(
+            stage_config.parallel_config,
+            exclude=frozenset({"world_size"}),
+        )
+    else:
+        engine_args.update(
+            _project_omni_config_fields(
+                stage_config.parallel_config,
+                exclude=frozenset({"world_size"}),
+            )
+        )
+
+    quantization_config = stage_config.quantization_config
+    if quantization_config is not None:
+        quantization_key = (
+            "quantization" if isinstance(quantization_config, str) and not is_diffusion else "quantization_config"
+        )
+        engine_args[quantization_key] = copy.deepcopy(quantization_config)
+
+    return engine_args
+
+
+def _finalize_engine_args_dict(
+    engine_args_dict: dict[str, Any],
+    *,
+    stage_type: str | StageType,
+    stage_id: int,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None,
+    cli_tokenizer: str | None,
+    has_sampling_extra_args: bool,
+) -> dict[str, Any]:
+    """Apply representation-independent engine adapter behavior."""
     pipeline_model_root = model
     model = engine_args_dict.pop("model", None) or model
     stage_defines_tokenizer = (
@@ -826,15 +928,15 @@ def build_engine_args_dict(
     if stage_connector_spec:
         engine_args_dict["stage_connector_spec"] = dict(stage_connector_spec or {})
 
-    if stage_type == "diffusion":
+    is_diffusion = stage_type == StageType.DIFFUSION
+    if is_diffusion:
         from vllm_omni.diffusion.data import parse_attention_config
 
         if engine_args_dict.get("diffusion_attention_config") is not None:
             engine_args_dict["diffusion_attention_config"] = parse_attention_config(
-                engine_args_dict.get("diffusion_attention_config"),
+                engine_args_dict["diffusion_attention_config"],
             )
-
-    if stage_type != "diffusion":
+    else:
         resolve_worker_cls(engine_args_dict)
 
     if engine_args_dict.get("worker_type") == "generation":
@@ -844,17 +946,86 @@ def build_engine_args_dict(
         engine_args_dict.setdefault("disable_hybrid_kv_cache_manager", True)
         engine_args_dict.setdefault("enable_prefix_caching", False)
 
-    # Check whether the stage's default_sampling_params defines extra_args.
-    default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {}))
-    engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args"))
+    engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args
 
     # TODO: Remove this after the performance regression is fixed
     # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression
     _maybe_set_qwen3_omni_moe_env(engine_args_dict)
-
     return engine_args_dict
 
 
+def build_legacy_engine_args_dict(
+    stage_config: Any,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None = None,
+    cli_tokenizer: str | None = None,
+) -> dict[str, Any]:
+    """Implement engine-argument building for the legacy stage representation."""
+    engine_args_dict = _to_dict(stage_config.engine_args)
+    # Legacy configs can materialize an omitted optional TP size as None.
+    # Remove it from the detached adapter dict so the backend default applies
+    # without mutating stage_config.engine_args.
+    if engine_args_dict.get("tensor_parallel_size") is None:
+        engine_args_dict.pop("tensor_parallel_size", None)
+
+    default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {}))
+    return _finalize_engine_args_dict(
+        engine_args_dict,
+        stage_type=_get_attr_or_item(stage_config, "stage_type", "llm"),
+        stage_id=stage_config.stage_id,
+        model=model,
+        stage_connector_spec=stage_connector_spec,
+        cli_tokenizer=cli_tokenizer,
+        has_sampling_extra_args=bool(default_sp.get("extra_args")),
+    )
+
+
+def build_engine_args_dict(
+    stage_config: Any,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None = None,
+    cli_tokenizer: str | None = None,
+) -> dict[str, Any]:
+    """Build engine arguments through the stable production entry point.
+
+    Production inputs still use the legacy stage representation. Keep that
+    compatibility choice behind this function so callers do not bind directly
+    to a representation-specific implementation.
+    """
+    return build_legacy_engine_args_dict(
+        stage_config,
+        model,
+        stage_connector_spec=stage_connector_spec,
+        cli_tokenizer=cli_tokenizer,
+    )
+
+
+def build_engine_args_dict_from_omni_stage_config(
+    stage_config: BaseVllmOmniStageConfig,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None = None,
+    cli_tokenizer: str | None = None,
+) -> dict[str, Any]:
+    """Project one typed stage config into backend engine arguments.
+
+    This projection is prepared for the RFC #4021 stage-init cutover. Current
+    production startup reaches the legacy implementation through
+    ``build_engine_args_dict`` while strategy and startup-plan inputs still
+    use the legacy representation.
+    """
+    engine_args_dict = _project_omni_stage_engine_args(stage_config)
+    _apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type)
+    return _finalize_engine_args_dict(
+        engine_args_dict,
+        stage_type=stage_config.stage_type,
+        stage_id=stage_config.stage_id,
+        model=model,
+        stage_connector_spec=stage_connector_spec,
+        cli_tokenizer=cli_tokenizer,
+        has_sampling_extra_args=stage_config.model_config.has_sampling_extra_args,
+    )
+
+
 def build_vllm_config(
     stage_config: Any,
     model: str,

vllm_omni/engine/stage_runtime.py

@@ -370,6 +370,9 @@ class StageRuntime:
             executor_class = None
             engine_args_dict = None
             if base_metadata.stage_type != "diffusion":
+                # The stable adapter entry point still receives the same
+                # legacy stage object as replica planning. Its implementation
+                # switches only at the coordinated RFC #4021 cutover.
                 engine_args_dict = build_engine_args_dict(
                     stage_cfg,
                     self._model,

vllm_omni/entrypoints/cli/serve.py

@@ -213,9 +213,10 @@ class OmniServeCommand(CLISubcommand):
             "--task-type",
             type=str,
             default=None,
-            choices=["CustomVoice", "VoiceDesign", "Base"],
-            help="Default task type for TTS models (CustomVoice, VoiceDesign, or Base). "
-            "If not specified, will be inferred from model path.",
+            help="Model-defined startup task type. The selected model validates "
+            "supported values; for example, TTS models accept CustomVoice, "
+            "VoiceDesign, or Base, while diffusion models may use it to select "
+            "task-specific weights. If omitted, the model default is used.",
         )
         # Forced aligner / word timestamps. --forced-aligner is the opt-in
         # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len)
@@ -934,6 +935,9 @@ def run_headless(args: TrackingNamespace) -> None:
     # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally
     # bracket each replica's spawn below with setup_stage_devices so they
     # don't all stack on cuda:0 (see ``per_replica_devices`` above).
+    # Headless startup still supplies the legacy OmegaConf stage shape through
+    # the stable adapter entry point. The implementation switches only when
+    # RFC #4021 threads structured stage configs through the launch plan.
     engine_args_dict = build_engine_args_dict(
         stage_cfg,
         model,

vllm_omni/entrypoints/omni_base.py

@@ -9,6 +9,7 @@ from typing import Any, Literal
 import huggingface_hub
 import vllm.envs as envs
 from vllm.logger import init_logger
+from vllm.transformers_utils.repo_utils import file_or_path_exists
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -73,6 +74,16 @@ def omni_snapshot_download(model_id: str) -> str:
 
         return snapshot_download(model_id)
 
+    # Modular Diffusers repositories describe independently loadable
+    # components. Let the selected pipeline download only its component
+    # sources instead of eagerly materializing the entire repository.
+    try:
+        if file_or_path_exists(model_id, "modular_model_index.json", revision=None):
+            return model_id
+    except (huggingface_hub.errors.GatedRepoError, huggingface_hub.errors.RepositoryNotFoundError):
+        # Preserve the more helpful errors raised by the full-download path.
+        pass
+
     try:
         download_weights_from_hf_specific(
             model_name_or_path=model_id,

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -481,18 +481,32 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
     if not frames:
         raise ValueError("No frames found to encode.")
 
-    frames_np = np.stack(frames, axis=0)
-    if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
-        frames_np = frames_np[..., :3]
+    frame_shape = frames[0].shape
+    has_alpha = len(frame_shape) == 3 and frame_shape[-1] == 4
+    output_shape = (*frame_shape[:-1], 3) if has_alpha else frame_shape
+    frames_u8 = np.empty((len(frames), *output_shape), dtype=np.uint8)
+    common_dtype = np.result_type(*(frame.dtype for frame in frames))
 
-    if frames_np.dtype == np.uint8:
-        frames_u8 = frames_np
-    else:
-        frames_np = np.clip(frames_np, 0.0, 1.0)
-        frames_np *= 255.0
-        frames_u8 = np.round(frames_np).astype(np.uint8)
+    # Convert one frame at a time instead of stacking the normalized float
+    # payload first. Long videos can otherwise require another full-size
+    # float array plus conversion temporaries before encoding.
+    for index, frame in enumerate(frames):
+        if frame.shape != frame_shape:
+            raise ValueError("All video frames must have the same shape.")
+        frame = frame[..., :3] if has_alpha else frame
+        if frame.dtype == np.uint8:
+            frames_u8[index] = frame
+            continue
 
-    return np.ascontiguousarray(frames_u8)
+        # np.stack(), used by the previous implementation, promoted mixed
+        # frame dtypes before scaling. Preserve those rounding semantics
+        # without allocating a full-video float buffer.
+        scaled = np.clip(frame.astype(common_dtype, copy=False), 0.0, 1.0)
+        scaled *= 255.0
+        np.rint(scaled, out=scaled)
+        frames_u8[index] = scaled
+
+    return frames_u8
 
 
 def _encode_video_bytes(

vllm_omni/entrypoints/utils.py

@@ -16,7 +16,10 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import _DEPLOY_DIR
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
-from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
+from vllm_omni.diffusion.utils.hf_utils import (
+    _looks_like_dreamzero,
+    get_diffusion_model_index,
+)
 from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -74,9 +77,9 @@ def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
     Returns:
         Model type string if found, None otherwise
     """
-    model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
+    model_index = get_diffusion_model_index(model)
     if model_index and isinstance(model_index, dict) and "_class_name" in model_index:
-        logger.debug(f"Found model_type '{model_index['_class_name']}' in model_index.json")
+        logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'")
         return model_index["_class_name"]
 
     return None
@@ -278,12 +281,12 @@ def resolve_model_config_path(model: str) -> str | None:
         model_type = hf_config.model_type
     except (ValueError, Exception):
         # If standard transformers format fails, try diffusers format
-        if file_or_path_exists(model, "model_index.json", revision=None):
+        if get_diffusion_model_index(model) is not None:
             model_type = _try_get_class_name_from_diffusers_config(model)
             if model_type is None:
                 raise ValueError(
                     f"Could not determine model_type for diffusers model: {model}. "
-                    f"Please ensure the model has 'model_type' in transformer/config.json or model_index.json"
+                    "Please ensure its Diffusers pipeline index contains '_class_name'"
                 )
         elif file_or_path_exists(model, "config.json", revision=None):
             # Try to read config.json manually for custom models like Bagel that fail get_config
@@ -308,7 +311,7 @@ def resolve_model_config_path(model: str) -> str | None:
             if model_type is None:
                 raise ValueError(
                     f"Could not determine model_type for model: {model}. "
-                    f"Model is not in standard transformers format and does not have model_index.json. "
+                    "Model is not in standard transformers or Diffusers format. "
                     f"Please ensure the model has proper configuration files with 'model_type' field"
                 )
 

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -770,8 +770,11 @@ AutoImageProcessor.register(MiniCPMOConfig, MiniCPMVImageProcessor, exist_ok=Tru
 
 
 # ============== SigLIP Vision Transformer Classes ==============
+FlashAttentionUnpaddingMetadata: TypeAlias = tuple[torch.Tensor, torch.Tensor, int]
+
+
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
+def _get_unpad_data(attention_mask: torch.Tensor) -> FlashAttentionUnpaddingMetadata:
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
     indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
     max_seqlen_in_batch = seqlens_in_batch.max().item()
@@ -1116,6 +1119,7 @@ class SiglipFlashAttention2(SiglipAttention):
         past_key_value: tuple[torch.Tensor] | None = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
         output_attentions = False
@@ -1178,7 +1182,13 @@ class SiglipFlashAttention2(SiglipAttention):
             value_states = value_states.to(target_dtype)
 
         attn_output = self._flash_attention_forward(
-            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=dropout_rate,
+            unpadding_metadata=unpadding_metadata,
         )
 
         attn_output = attn_output.reshape(bsz, q_len, self.embed_dim).contiguous()
@@ -1190,7 +1200,15 @@ class SiglipFlashAttention2(SiglipAttention):
         return attn_output, attn_weights
 
     def _flash_attention_forward(
-        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        query_length,
+        dropout=0.0,
+        softmax_scale=None,
+        unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None,
     ):
         """
         Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
@@ -1218,7 +1236,12 @@ class SiglipFlashAttention2(SiglipAttention):
         if attention_mask is not None:
             batch_size = query_states.shape[0]
             query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
-                query_states, key_states, value_states, attention_mask, query_length
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                query_length,
+                unpadding_metadata=unpadding_metadata,
             )
 
             cu_seqlens_q, cu_seqlens_k = cu_seq_lens
@@ -1245,8 +1268,18 @@ class SiglipFlashAttention2(SiglipAttention):
 
         return attn_output
 
-    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
-        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+    def _upad_input(
+        self,
+        query_layer,
+        key_layer,
+        value_layer,
+        attention_mask,
+        query_length,
+        unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None,
+    ):
+        if unpadding_metadata is None:
+            unpadding_metadata = _get_unpad_data(attention_mask)
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = unpadding_metadata
         batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
 
         key_layer = index_first_axis(
@@ -1316,6 +1349,7 @@ class SiglipEncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
         output_attentions: bool | None = False,
+        unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None,
     ) -> tuple[torch.FloatTensor]:
         """
         Args:
@@ -1330,11 +1364,19 @@ class SiglipEncoderLayer(nn.Module):
         residual = hidden_states
 
         hidden_states = self.layer_norm1(hidden_states)
-        hidden_states, attn_weights = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-        )
+        if self._use_flash_attention_2:
+            hidden_states, attn_weights = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                unpadding_metadata=unpadding_metadata,
+            )
+        else:
+            hidden_states, attn_weights = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+            )
         hidden_states = residual + hidden_states
 
         residual = hidden_states
@@ -1435,6 +1477,7 @@ class SiglipEncoder(nn.Module):
         self.config = config
         self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
         self.gradient_checkpointing = False
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
 
     # Ignore copy
     def forward(
@@ -1475,6 +1518,9 @@ class SiglipEncoder(nn.Module):
         all_attentions = () if output_attentions else None
 
         hidden_states = inputs_embeds
+        unpadding_metadata = (
+            _get_unpad_data(attention_mask) if self._use_flash_attention_2 and attention_mask is not None else None
+        )
         for encoder_layer in self.layers:
             if output_hidden_states:
                 encoder_states = encoder_states + (hidden_states,)
@@ -1484,12 +1530,14 @@ class SiglipEncoder(nn.Module):
                     hidden_states,
                     attention_mask,
                     output_attentions,
+                    unpadding_metadata,
                 )
             else:
                 layer_outputs = encoder_layer(
                     hidden_states,
                     attention_mask,
                     output_attentions=output_attentions,
+                    unpadding_metadata=unpadding_metadata,
                 )
 
             hidden_states = layer_outputs[0]

vllm_omni/platforms/cuda/platform.py

@@ -134,6 +134,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 try:
                     importlib.import_module("kernels")

vllm_omni/platforms/interface.py

@@ -127,6 +127,29 @@ class OmniPlatform(Platform):
         """
         raise NotImplementedError
 
+    @classmethod
+    def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
+        """Reject an explicitly selected backend this platform cannot run.
+
+        Platforms call this from ``get_diffusion_attn_backend_cls`` so that a
+        backend restricted to other hardware, or one whose kernel package is
+        missing, fails during resolution rather than at the first forward.
+        """
+        from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+
+        backend_upper = selected_backend.upper()
+        backend_cls = DiffusionAttentionBackendEnum[backend_upper].get_class()
+
+        supported = backend_cls.supported_platforms
+        platform = cls._omni_enum.value
+        if supported is not None and platform not in supported:
+            raise ValueError(
+                f"The {backend_upper} diffusion attention backend runs on {', '.join(supported)} only, "
+                f"but the current platform is {platform}. Select a backend supported here, "
+                "such as FLASH_ATTN or TORCH_SDPA."
+            )
+        backend_cls.validate_available()
+
     @classmethod
     def supports_torch_inductor(cls) -> bool:
         """Check if the platform supports torch.compile with inductor backend."""

vllm_omni/platforms/musa/platform.py

@@ -80,6 +80,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/npu/platform.py

@@ -130,6 +130,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/rocm/platform.py

@@ -102,6 +102,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/xpu/platform.py

@@ -53,6 +53,7 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/quantization/int8_config.py

@@ -50,9 +50,55 @@ if TYPE_CHECKING:
 # Dynamic quantization is supported first.
 ACTIVATION_SCHEMES = ["dynamic"]
 
+# Ascend's npu_quant_matmul (QuantBatchMatmulV3) refuses a weight whose last
+# dimension exceeds this, and it fails at the first forward rather than at load
+# time. Layers still wider than the limit after TP sharding stay unquantized.
+NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535
+
 logger = init_logger(__name__)
 
 
+def _fell_back_to_unquantized_npu(
+    layer: torch.nn.Module,
+    input_size_per_partition: int,
+    output_partition_sizes: list[int],
+    input_size: int,
+    output_size: int,
+    params_dtype: torch.dtype,
+    **extra_weight_attrs,
+) -> bool:
+    """Swap a layer to unquantized weights when npu_quant_matmul cannot run its shape.
+
+    The check uses the per-partition output size because that is what the kernel
+    actually sees; a layer over the limit on one rank can be within it at a
+    higher TP degree. Returns True when the layer was swapped and its weights
+    created, in which case the caller must not create its own.
+    """
+    output_size_per_partition = sum(output_partition_sizes)
+    if output_size_per_partition <= NPU_QUANT_MATMUL_MAX_OUT_FEATURES:
+        return False
+
+    logger.warning_once(
+        "Keeping a %d-wide linear unquantized: npu_quant_matmul rejects an output dimension past "
+        "%d. Tensor parallelism shrinks this per-rank dimension, so a higher TP degree brings such "
+        "layers back into range.",
+        output_size_per_partition,
+        NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+    )
+    fallback = UnquantizedLinearMethod()
+    layer.quant_method = fallback
+    fallback.create_weights(
+        layer,
+        input_size_per_partition,
+        output_partition_sizes,
+        input_size,
+        output_size,
+        params_dtype,
+        **extra_weight_attrs,
+    )
+    return True
+
+
 def create_weight_parameter(
     output_size_per_partition: int,
     input_size_per_partition: int,
@@ -231,6 +277,21 @@ class LazyWeightMixin:
 
     uses_meta_device: bool = True
 
+    # This mixin knows when a layer's weight is final, so it can hand the layer
+    # back to the host right there. Loaders that intend to offload the whole
+    # model after loading opt in per layer via ``enable_offload_after_quant``.
+    supports_offload_after_quant: bool = True
+    _offload_after_quant: bool = False
+
+    def enable_offload_after_quant(self) -> None:
+        """Return each layer to host memory as soon as it has been quantized.
+
+        Caps the load-time device footprint at one layer instead of the whole
+        model. A quant method instance belongs to a single layer, so this is not
+        a global switch.
+        """
+        self._offload_after_quant = True
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -287,6 +348,11 @@ class LazyWeightMixin:
                 # anything
                 layer._already_called_process_weights_after_loading = True
 
+                # This layer's weight is final, so nothing needs it on the
+                # accelerator until inference.
+                if self._offload_after_quant:
+                    layer.to("cpu")
+
                 # Note that we keep `layer._loaded_numel` around just in case
                 # there is logic added to vllm in the future which calls a
                 # weight loader twice - we do not want to re-initialize in
@@ -354,6 +420,11 @@ class NPUInt8LinearMethod(BaseInt8LinearMethod):
     def __init__(self, quant_config: DiffusionInt8Config):
         super().__init__(quant_config)
 
+    def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None:
+        if _fell_back_to_unquantized_npu(layer, *args, **kwargs):
+            return
+        super().create_weights(layer, *args, **kwargs)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         layer.weight.data = layer.weight.data.t().contiguous()
         layer.weight_scale.data = layer.weight_scale.data.squeeze()
@@ -424,6 +495,13 @@ class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod):
     and quantized the weights during loading.
     """
 
+    def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None:
+        # NPUInt8LinearMethod's override is unreachable from here: LazyWeightMixin
+        # comes first in the MRO and does not call super().
+        if _fell_back_to_unquantized_npu(layer, *args, **kwargs):
+            return
+        super().create_weights(layer, *args, **kwargs)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return