[Model] Add soundfile fallback for MiniMax H3 audio loading (#5699)

Signed-off-by: brandneway <gyuan4892@gmail.com>

.buildkite/cuda/test-nightly.yml

@@ -457,14 +457,20 @@ steps:
   - group: ":card_index_dividers: Diffusion Test"
     key: nightly-diffusion-test-group
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     steps:
-      - label: ":full_moon: Diffusion · Distributed Test"
+      - label: ":full_moon: Diffusion · Distributed Test with L4"
         timeout_in_minutes: 60
         commands:
-          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
         mirror_hardwares: l4_4
 
+      - label: ":full_moon: Diffusion · Distributed Test with H100"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
+        mirror_hardwares: h100_4
+
       - label: ":full_moon: Diffusion · Offloader Test"
         timeout_in_minutes: 60
         commands:

.buildkite/npu/test-npu-nightly.yml

@@ -95,6 +95,80 @@ steps:
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
             exit $$EXIT
+  - group: ":card_index_dividers: Diffusion I2I Model Test"
+    key: nightly-HunyuanImage3-x2i-test-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
+    steps:
+      - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
+          HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model
+  - group: ":card_index_dividers: Diffusion Benchmark Perf Test"
+    key: nightly-diffusion-benchmark-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
+    steps:
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.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
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.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
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4.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
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
     depends_on: upload-nightly-pipeline
@@ -103,6 +177,14 @@ steps:
       build.pull_request.labels includes "nightly-test" ||
       build.pull_request.labels includes "diffusion-x2v-test"
     steps:
+      - label: ":full_moon: Diffusion X2V · Function Test"
+        mirror_hardwares: a2b3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
       - label: ":full_moon: Diffusion X2V · Perf Test"
         key: nightly-diffusion-x2v-performance
         timeout_in_minutes: 180

docker/Dockerfile.npu.ci.a3

@@ -45,6 +45,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "hyperpyyaml>=1.2.0" \
       "onnx>=1.16.0" \
       "librosa>=0.10.0" \
+      "torchmetrics>=1.9.0" \
       --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60 && \

docs/models/supported_models.md

@@ -44,6 +44,7 @@ th {
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
+| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion/attention_backends.md

@@ -14,8 +14,8 @@ 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 and **SAGE** quantized attention — see below. Requires `flashinfer`. |
-| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
+| `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`. |
+| `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`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
@@ -222,6 +222,30 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py
 
 Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
 
+## Known Limitations
+
+### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
+
+LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
+
+**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 ...
+```
+
+### FlashAttention-4 on Blackwell
+
+Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
+
+```bash
+pip install 'vllm-omni[fa4]'
+```
+
+`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
+when it is unavailable. Version `4.0.0b18` is required; earlier beta
+wheels had known JIT failures on Blackwell.
 
 ## Choosing a Backend Manually
 

pyproject.toml

@@ -96,6 +96,12 @@ quack = [
     "quack-kernels>=0.3.11",
 ]
 
+# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because
+# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific.
+fa4 = [
+    "flash-attn-4[cu13]==4.0.0b18",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -0,0 +1,162 @@
+# MiniMax H3 on Ascend NPU
+
+> Joint video and audio generation with text, first-frame, image/audio, or
+> multi-video conditions — Ascend NPU deployment guide
+
+## Summary
+
+- Vendor: MiniMaxAI
+- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
+- Tasks: T2VA, FL2VA, and Ref2VA
+- Mode: OpenAI-compatible `/v1/videos` HTTP serving
+- Hardware: Atlas 800I A3
+- Maintainer: Community
+
+This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU
+environments. Differences from the GPU path:
+
+- Runs on a CPU-only (aarch64) PyTorch build with `torch_npu`; no CUDA
+  runtime is present.
+- Audio loading does **not** require TorchCodec (whose aarch64 wheels are
+  built against CUDA torch and fail to load on CPU-only builds). vLLM-Omni
+  automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio
+  inputs.
+
+## Prerequisites
+
+### Checkpoint
+
+Same as the GPU recipe — Hugging Face access approval is required:
+
+```bash
+hf auth login
+export MODEL_ROOT=/path/to/MiniMax-H3
+hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+```
+
+### Environment
+
+- Ascend driver & firmware: 25.5.1
+- CANN toolkit: 9.0.1
+- Python: 3.12
+- PyTorch: 2.10.0+cpu
+- torch_npu: 2.10.0.post2
+- Install vLLM-Omni from a checkout with MiniMax H3 support:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e .
+```
+
+Install the **mindie-sd** operator library to enable Ascend-optimized fused
+operators (`adalayernorm`, etc.):
+
+```bash
+git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
+
+# Comment out the tik_ops build step (not needed for this use case)
+sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
+
+python setup.py bdist_wheel
+cd dist
+pip install mindiesd-*.whl
+```
+
+- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
+  reference-video preparation and MP4 output.
+- Reference-video decoding uses `decord` when available and falls back to
+  PyAV otherwise.
+- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded
+  at native sample rate through the soundfile / ffmpeg fallback.
+
+## Start a server
+
+One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
+and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
+
+### Multi-NPU: 768P validated 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:
+
+```bash
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export PORT=9098
+export MODEL="${MODEL_ROOT}/FL2VA"
+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 \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 8 \
+  --usp 8 \
+  --ring 1 \
+  --text-encoder-tp-size 8 \
+  --enable-layerwise-offload \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --vae-patch-parallel-size 8 \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+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:
+
+```bash
+export MODEL="${MODEL_ROOT}/Ref2VA"
+```
+
+## HTTP API examples
+
+The request format is identical to the GPU recipe; see
+[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples).
+Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU.
+
+## Key parameters
+
+Same as the GPU recipe; see
+[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters).
+The validated resolution on NPU is 768P (e.g. 1344x768).
+
+## Validated NPU evidence
+
+Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1,
+PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU
+configuration above:
+
+| Workload | Configuration |
+|----------|---------------|
+| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+
+These measurements describe the validated shapes rather than a general
+throughput guarantee.
+
+## Known limitations
+
+- Each server process loads only one checkpoint partition.
+- 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.
+- Image+audio Ref2VA accepts exactly one image and one audio reference.
+- Video Ref2VA accepts one or more video files, but not an additional
+  standalone audio reference.
+- VAE patch parallelism requires size 1 or the full DiT group size and
+  supports the H3 native `tile` mode only.
+
+## Additional resources
+
+- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)

recipes/MiniMaxAI/MiniMax-H3.md

@@ -0,0 +1,329 @@
+# MiniMax H3
+
+> Joint video and audio generation with text, first-frame, image/audio, or
+> multi-video conditions
+
+## Summary
+
+- Vendor: MiniMaxAI
+- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
+- Tasks: T2VA, FL2VA, and Ref2VA
+- Mode: OpenAI-compatible `/v1/videos` HTTP serving
+- Maintainer: Community
+
+MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
+checkpoint has two independently served partitions:
+
+- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
+  (`fl2va`)
+- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
+  (`ref2va`)
+
+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:
+
+```bash
+hf auth login
+export MODEL_ROOT=/path/to/MiniMax-H3
+hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+```
+
+Install vLLM-Omni from the checkout containing MiniMax H3 support. On
+Blackwell, install the optional FlashAttention-4 dependency:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e '.[fa4]'
+```
+
+`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
+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.
+
+### Single GPU: accuracy and memory first
+
+The single-GPU configuration uses model-level CPU offload.
+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 PORT=8091
+
+CUDA_VISIBLE_DEVICES=0 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 1 \
+  --enable-cpu-offload \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+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.
+
+### Four GPUs: measured best-practice throughput
+
+The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
+
+- 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;
+- FlashAttention, with Ring and TP left at 1.
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+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}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 4 \
+  --usp 4 \
+  --ring 1 \
+  --vae-patch-parallel-size 4 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+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`.
+
+### Text encoder tensor parallelism
+
+The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
+default fully resident on the DiT main rank.  On multi-GPU no-offload runs that
+rank becomes the peak-memory hotspot.  Add `--text-encoder-tp-size N` to shard
+the encoder across the first `N` DiT ranks (the encoder is implemented with
+vLLM-style tensor-parallel layers and runs with distributed collectives over
+its own encoder process group):
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+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}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 4 \
+  --usp 4 \
+  --ring 1 \
+  --text-encoder-tp-size 4 \
+  --vae-patch-parallel-size 4 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
+valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs).  The
+encoder TP rank set is the first `N` DiT ranks; on 4 GPUs
+`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main
+rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the
+other ranks each gain `~51.5/N` GB.  The encoder output remains identical to
+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"
+```
+
+## HTTP API examples
+
+The following requests use the synchronous endpoint so the returned body can
+be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can
+also be used when job polling is preferred.
+
+All four tasks use 24 FPS, 50 sigma points, seed values from the validated
+workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3.
+Decimal durations are passed through `extra_params`.
+
+Set the endpoint once:
+
+```bash
+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:
+
+```bash
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=1101' \
+  -F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \
+  -o t2va.mp4
+```
+
+### 2. FL2VA: first frame to video and audio
+
+Run this request against the `FL2VA` partition. When width and height are
+omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
+edge.
+
+```bash
+export FIRST_FRAME=/path/to/fl2va_first_frame.png
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=2101' \
+  -F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \
+  -F "input_reference=@${FIRST_FRAME};type=image/png" \
+  -o fl2va.mp4
+```
+
+### 3. Ref2VA: image and audio reference
+
+Run this request against the `Ref2VA` partition. `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 \
+  --bind 127.0.0.1 \
+  --directory /path/to/reference_assets
+```
+
+Then submit the image and audio request from another terminal:
+
+```bash
+export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
+export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=3101' \
+  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
+  -F "input_reference=@${REF_IMAGE};type=image/png" \
+  -F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \
+  -o ref2va_image_audio.mp4
+```
+
+The requested duration should cover the complete audio. If `duration` is
+shorter, the reference soundtrack is truncated to the generated clip.
+
+### 4. Ref2VA: two video references
+
+Run this request against the `Ref2VA` partition. Repeat the
+`input_references` multipart field once per source video. H3 consumes the
+videos in form order and preserves their original soundtracks during
+conditioning.
+
+```bash
+export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4
+export BACKGROUND_VIDEO=/path/to/fairytale_background.mov
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=3101' \
+  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
+  -F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \
+  -F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \
+  -o ref2va_video_video.mp4
+```
+
+The server stores uploaded multi-video references only for the lifetime of the
+request and deletes the temporary files after generation. Video Ref2VA uses
+the source-video soundtracks and does not accept a separate
+`audio_reference`.
+
+## Key parameters
+
+| Parameter | Recommended value | Notes |
+|-----------|-------------------|-------|
+| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
+| `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 |
+| `flow_shift` | `12` | Video sigma shift |
+| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
+| `seed` | Task-specific | Use a fixed value for reproducibility |
+| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
+
+## Validated four-GPU evidence
+
+The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
+excluded warmup followed by three requests.
+
+| Workload | Configuration | Observed result |
+|----------|---------------|-----------------|
+| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
+| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
+
+These measurements describe the validated shapes rather than a general
+throughput guarantee. Multi-video Ref2VA is much slower because the two
+reference videos expand both the Qwen3-VL vision sequence and the packed DiT
+attention sequence.
+
+## Known limitations
+
+- Each server process loads only one checkpoint partition.
+- 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.
+- Image+audio Ref2VA accepts exactly one image and one audio reference.
+- Video Ref2VA accepts one or more video files, but not an additional standalone
+  audio reference.
+- VAE patch parallelism requires size 1 or the full DiT group size and supports
+  the H3 native `tile` mode only.
+
+## Additional resources
+
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)
+- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)

recipes/README.md

@@ -45,6 +45,7 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
+| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |

tests/config/test_omni_config.py

@@ -440,6 +440,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
+        "text_encoder_tp_size",
         "vae_parallel_mode",
         "use_hsdp",
         "mask_sp_padding",

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }
@@ -27,16 +28,9 @@
             "async_chunk": false,
             "trust_remote_code": true,
             "connectors": {
-              "rdma_connector": {
-                "name": "MooncakeTransferEngineConnector",
-                "extra": {
-                  "host": "auto",
-                  "zmq_port": 50051,
-                  "protocol": "rdma",
-                  "device_name": "",
-                  "memory_pool_size": 17179869184,
-                  "memory_pool_device": "cpu"
-                }
+              "shared_mem_connector": {
+                "name": "SharedMemoryConnector",
+                "extra": {}
               }
             },
             "stages": [
@@ -66,7 +60,7 @@
                   "need_send_cache": true
                 },
                 "output_connectors": {
-                  "to_stage_1": "rdma_connector"
+                  "to_stage_1": "shared_mem_connector"
                 },
                 "default_sampling_params": {
                   "temperature": 0.0,
@@ -96,7 +90,7 @@
                 },
                 "vae_use_tiling": true,
                 "input_connectors": {
-                  "from_stage_0": "rdma_connector"
+                  "from_stage_0": "shared_mem_connector"
                 },
                 "default_sampling_params": {
                   "num_inference_steps": 50,

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/diffusion/attention/test_flash_attn.py

@@ -20,7 +20,6 @@ from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
-
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813
@@ -330,6 +329,48 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     assert calls[0]["softmax_scale"] == 0.5
 
 
+def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
+    calls = []
+
+    def fake_varlen_func(q, k, v, **kwargs):
+        calls.append((q.shape, k.shape, v.shape, kwargs))
+        return q
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True] * 6 + [False] * 2]),
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert len(calls) == 1
+    assert calls[0][0] == torch.Size([8, 2, 4])
+    assert calls[0][3]["cu_seqlens_q"] is cu_seqlens
+    assert calls[0][3]["max_seqlen_q"] == 6
+
+
+def test_packed_varlen_metadata_must_be_complete(monkeypatch):
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_ring_fa4.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.ring import ring_kernels
+from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
+from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch):
+    expected_out = object()
+    expected_lse = object()
+    captured = {}
+
+    def fake_fa4(q, k, v, **kwargs):
+        captured.update(q=q, k=k, v=v, **kwargs)
+        return expected_out, expected_lse
+
+    monkeypatch.setattr(ring_kernels, "HAS_FA4", True)
+    monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4)
+
+    out, lse = ring_kernels.flash_attn4_func_forward(
+        "q",
+        "k",
+        "v",
+        softmax_scale=0.125,
+        causal=True,
+        window_size=(-1, -1),
+    )
+
+    assert (out, lse) == (expected_out, expected_lse)
+    assert captured == {
+        "q": "q",
+        "k": "k",
+        "v": "v",
+        "softmax_scale": 0.125,
+        "causal": True,
+        "window_size": (None, None),
+        "softcap": 0.0,
+        "return_lse": True,
+    }
+
+
+def test_ring_prefers_fa4_on_blackwell(monkeypatch):
+    captured = {}
+
+    def fake_ring_flash_attn_func(*args, **kwargs):
+        captured.update(kwargs)
+        return "output"
+
+    monkeypatch.setattr(ring_parallel, "HAS_FA4", True)
+    monkeypatch.setattr(ring_parallel, "HAS_FA3", True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+
+    from vllm_omni.diffusion.attention.backends import ring_flash_attn
+
+    monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func)
+
+    query = SimpleNamespace(
+        dtype=torch.bfloat16,
+        shape=(1, 16, 4, 128),
+        is_cuda=True,
+        device=torch.device("cuda:0"),
+    )
+    strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object()))
+
+    assert strategy.run_attention(query, "key", "value", None) == "output"
+    assert captured["attn_type"] is AttnType.FA4

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -39,7 +39,7 @@ MODEL_CASES = [
             "needs_image": False,
             "is_moe": False,
             "mean_threshold": 3e-2,
-            "p99_threshold": 3e-2,
+            "p99_threshold": 4e-2,
         },
         id="wan22_t2v_a14b",
     ),
@@ -146,7 +146,7 @@ def _run_generate(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4})
 @pytest.mark.parametrize("model_case", MODEL_CASES)
 def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
     if current_omni_platform.is_npu():

tests/diffusion/models/minimax_h3/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_pipeline_import_registry_and_component_discovery():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+    )
+
+    assert _DIFFUSION_MODELS["MiniMaxH3Pipeline"] == (
+        "minimax_h3",
+        "pipeline_minimax_h3",
+        "MiniMaxH3Pipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
+    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
+    assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
+    assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
+
+
+def test_joint_postprocess_is_multiprocessing_picklable():
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        get_minimax_h3_post_process_func,
+    )
+
+    postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
+    postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
+    video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+    audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
+
+    result = postprocess((video, audio), output_type="np")
+
+    assert isinstance(result["video"], list)
+    assert result["video"][0].shape == (2, 4, 5, 3)
+    np.testing.assert_array_equal(result["audio"], audio.numpy())
+    assert result["audio_sample_rate"] == 32000
+    assert result["fps"] == 24
+
+
+def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    od_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(cfg_parallel_size=2),
+    )
+    with pytest.raises(ValueError, match="CFG-distilled"):
+        MiniMaxH3Pipeline(od_config=od_config)
+
+
+def test_shape_contract_matches_three_reference_tasks():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        MINIMAX_H3_SHAPE_PLANNER,
+        minimax_h3_align_frame_count,
+    )
+
+    assert minimax_h3_align_frame_count(round(8.7 * 24)) == 209
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(209) == 62
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(8.7) == 348
+    assert minimax_h3_align_frame_count(round(5.0 * 24)) == 124
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(124) == 37
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(5.0) == 200
+    assert MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(62) == 209
+
+
+def test_shifted_sigma_schedule_matches_reference_values():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        minimax_h3_time_shift_sigmas,
+    )
+
+    sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0)
+
+    assert sigmas == pytest.approx(
+        [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0],
+        abs=1e-7,
+    )
+
+
+def test_reference_image_resize_contract():
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _reference_image_shape,
+    )
+
+    assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == (
+        2048,
+        2720,
+    )
+    with pytest.raises(ValueError, match="aspect ratio"):
+        _reference_image_shape(Image.new("RGB", (100, 501)))
+
+
+def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values = torch.ones(1, 4)
+    image_grid_thw = torch.tensor([[1, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+
+def test_encoder_forward_forwards_video_inputs():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values_videos = torch.ones(1, 4)
+    video_grid_thw = torch.tensor([[2, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=None,
+        image_grid_thw=None,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+
+def test_reference_video_shape_uses_h3_adapt_shape_policy():
+    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+        _reference_video_shape,
+    )
+
+    assert _reference_video_shape(1280, 720) == (1344, 768)
+    assert _reference_video_shape(3844, 2160) == (1344, 768)
+
+
+def test_text_encoder_stub_constructs_without_group_or_weights():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = MiniMaxH3Qwen3VLEncoder(
+        "/nonexistent/text_encoder",
+        device=torch.device("cpu"),
+        load_model=False,
+        encoder_group=None,
+    )
+    assert not encoder.is_loaded
+    assert encoder.tp_size == 1
+    # The stub has no parameters, so it never contributes to the runner's
+    # strict missing-parameter check on non-encoder ranks.
+    assert list(encoder.named_parameters()) == []
+
+
+def test_no_offload_keeps_text_encoder_resident():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    input_ids = torch.tensor([1, 2])
+
+    actual = pipeline._encode_text_hidden(input_ids, {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
+
+
+def test_model_offload_uses_hooked_text_encoder_call():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+    )
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    vision_kwargs = {"pixel_values": torch.ones(1, 4)}
+
+    actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
+
+    assert actual is expected
+    pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
+    pipeline.text_encoder.load_to_device.assert_not_called()
+
+
+def test_layerwise_offload_releases_text_encoder():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=True,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+
+def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = torch.nn.Linear(1, 1).half()
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("cpu"),
+    )
+
+    assert next(video_vae.parameters()).dtype == torch.float32
+
+
+def test_video_vae_encode_uses_configured_parallel_tiling():
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        MiniMaxH3VideoVAE,
+    )
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.parallel_tiling = True
+            self.encode_calls = []
+
+        def encode_videos(self, frames, *, use_fp16_latent):
+            assert self.parallel_tiling
+            self.encode_calls.append((frames, use_fp16_latent))
+            return [torch.ones(1, 1, 2, 2, 2)]
+
+    video_vae = object.__new__(MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(video_vae)
+    video_vae.model = FakeModel()
+    video_vae.config_dict = {
+        "latent_channels": 1,
+        "latents_mean": [0.0],
+        "latents_std": [1.0],
+    }
+
+    rows, shape = video_vae.encode_video("frames")
+
+    assert video_vae.model.parallel_tiling
+    assert video_vae.model.encode_calls == [("frames", True)]
+    assert rows.shape == (2, 4)
+    assert shape == (2, 2, 2)
+
+
+def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    prepared = [
+        {"prepared_path": "video-1.mp4"},
+        {"prepared_path": "video-2.mp4"},
+    ]
+
+    class FakeVideoVAE:
+        def __init__(self):
+            self.calls = []
+
+        def is_distributed_enabled(self):
+            return True
+
+        def encode_video(self, frames):
+            self.calls.append(frames)
+            index = len(self.calls)
+            return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3)
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = FakeVideoVAE()
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "_dit_rank_world",
+        lambda: ("dit-group", 1, 4),
+    )
+
+    def fake_broadcast_object_list(values, *, src, group, device):
+        assert values == [None]
+        assert (src, group, device) == (0, "dit-group", torch.device("cpu"))
+        values[0] = prepared
+
+    monkeypatch.setattr(
+        pipeline_module.dist,
+        "broadcast_object_list",
+        fake_broadcast_object_list,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "load_video_frames",
+        lambda path: f"frames:{path}",
+    )
+
+    rows, shapes = pipeline._encode_video_conditions(None, count=2)
+
+    assert pipeline.video_vae.calls == [
+        "frames:video-1.mp4",
+        "frames:video-2.mp4",
+    ]
+    torch.testing.assert_close(
+        rows,
+        torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
+    )
+    assert shapes == [(1, 2, 3), (2, 2, 3)]

tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import numpy as np
+import pytest
+
+pytestmark = [
+    pytest.mark.local_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
+
+
+@pytest.mark.skipif(
+    not os.environ.get(MODEL_ENV),
+    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
+)
+def test_minimax_h3_t2va_ulysses8_smoke():
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    if torch.accelerator.device_count() < 8:
+        pytest.skip("MiniMax H3 Ulysses smoke requires eight GPUs")
+
+    engine = Omni(
+        model=os.environ[MODEL_ENV],
+        parallel_config=DiffusionParallelConfig(ulysses_degree=8),
+        trust_remote_code=True,
+        enable_cpu_offload=True,
+        enforce_eager=True,
+    )
+    try:
+        outputs = engine.generate(
+            "A quiet cinematic night scene with matching ambient sound.",
+            OmniDiffusionSamplingParams(
+                height=256,
+                width=448,
+                num_frames=29,
+                fps=24,
+                num_inference_steps=2,
+                seed=42,
+                output_type="np",
+                extra_args={
+                    "task": "t2va",
+                    "flow_shift": 12.0,
+                    "audio_flow_shift": 3.0,
+                },
+            ),
+            use_tqdm=False,
+        )
+    finally:
+        engine.close()
+
+    assert len(outputs) == 1
+    frames = np.asarray(outputs[0].images[0])
+    assert frames.shape == (39, 256, 448, 3)
+    multimodal = outputs[0].multimodal_output
+    assert multimodal is not None
+    assert np.asarray(multimodal["audio"]).shape[1] == 2
+    assert multimodal["audio_sample_rate"] == 32000
+    assert multimodal["fps"] == 24

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py

@@ -0,0 +1,146 @@
+# 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_video_patchify_round_trip_preserves_values():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_patchify_video_latent,
+        minimax_h3_unpatchify_video_tokens,
+    )
+
+    latent = torch.arange(2 * 3 * 2 * 4 * 6).reshape(2, 3, 2, 4, 6)
+    rows = minimax_h3_patchify_video_latent(
+        latent,
+        patch_size=(1, 2, 2),
+    )
+    restored = minimax_h3_unpatchify_video_tokens(
+        rows,
+        latent_shape=(2, 2, 3, 3),
+        patch_size=(1, 2, 2),
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_audio_pack_round_trip_preserves_channel_major_order():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_pack_audio_latent,
+        minimax_h3_unpack_audio_tokens,
+    )
+
+    latent = torch.arange(2 * 4 * 5).reshape(2, 4, 5)
+    rows = minimax_h3_pack_audio_latent(latent)
+    restored = minimax_h3_unpack_audio_tokens(
+        rows,
+        audio_t=10,
+        audio_channel=2,
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_t2va_and_fl2va_packing_keep_update_rows_separate():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence,
+    )
+
+    common = dict(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+    )
+    t2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=False,
+    )
+    fl2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=True,
+        keyframe_frame_indices=[0],
+        frame_count=5,
+    )
+
+    assert int(t2va["seq_len"]) == 64
+    assert t2va["img_pos"].numel() == 12
+    assert t2va["update_mask"].all()
+    assert fl2va["img_pos"].numel() == 18
+    assert fl2va["update_mask"].sum().item() == 12
+    assert (~fl2va["update_mask"]).sum().item() == 6
+    assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
+
+
+def test_ref2va_packing_tracks_video_and_audio_update_masks():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+
+    packed = minimax_h3_packed_sequence_ref2va_blocks(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 2},
+        ],
+    )
+
+    assert int(packed["seq_len"]) == 64
+    assert packed["img_pos"].numel() == 16
+    assert packed["update_mask"].sum().item() == 12
+    assert packed["audio_pos"].numel() == 10
+    assert packed["audio_update_mask"].sum().item() == 6
+    assert packed["cu_seqlens"].tolist() == [0, 30, 64]
+
+
+def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
+    from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
+        minimax_h3_audio_cond_noise_aug_rows,
+        minimax_h3_imgvid_cond_noise_aug_rows,
+    )
+
+    image_rows = torch.randn(4, 96)
+    audio_rows = torch.randn(6, 32)
+
+    torch.testing.assert_close(
+        minimax_h3_imgvid_cond_noise_aug_rows(
+            image_rows,
+            condition_shapes=[(1, 4, 4)],
+            target_latent_t=2,
+            imgvid_cond_num_frames=1,
+            seed=42,
+            noise_aug=1.0,
+        ),
+        image_rows,
+    )
+    torch.testing.assert_close(
+        minimax_h3_audio_cond_noise_aug_rows(
+            audio_rows,
+            condition_audio_t=[3],
+            seed=42,
+            noise_aug=1.0,
+        ),
+        audio_rows,
+    )
+
+    first = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    second = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    torch.testing.assert_close(first, second)

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_grouped_qkv_checkpoint_reorder():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        _reorder_grouped_qkv_to_qkv,
+    )
+
+    # Two groups with rows [q, k, v] become [q0, q1, k0, k1, v0, v1].
+    grouped = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+    reordered = _reorder_grouped_qkv_to_qkv(
+        grouped,
+        num_query_groups=2,
+        heads_per_group=1,
+        head_dim=1,
+    )
+
+    assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+
+
+def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
+    from cache_dit import ForwardPattern
+
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    assert MiniMaxH3DiTModel._repeated_blocks == ["MiniMaxH3DiTBlock"]
+    assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
+    assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
+    assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
+    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.token_refiner = nn.Module()
+    model.token_refiner.blocks = nn.ModuleList([nn.Linear(4, 4)])
+    model.final_layer = nn.Linear(4, 4)
+
+    matched = [
+        name
+        for name, module in model.named_modules()
+        if any(condition(name, module) for condition in MiniMaxH3DiTModel._hsdp_shard_conditions)
+    ]
+    assert matched == ["blocks.0", "blocks.1"]
+
+
+def test_packed_attention_is_a_regional_compile_boundary():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
+
+
+@pytest.mark.parametrize(
+    ("tp_size", "message"),
+    [
+        (3, "num_attention_heads"),
+        (5, "num_attention_heads"),
+    ],
+)
+def test_tp_rejects_non_divisible_head_counts(tp_size, message):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    with pytest.raises(ValueError, match=message):
+        model._validate_tp_config(
+            arch=MiniMaxH3DiTArchConfig(),
+            tp_size=tp_size,
+        )
+
+
+def test_tp_accepts_checkpoint_supported_sizes():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    arch = MiniMaxH3DiTArchConfig()
+    for tp_size in (1, 2, 4, 7):
+        model._validate_tp_config(arch=arch, tp_size=tp_size)

tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py

@@ -0,0 +1,257 @@
+"""Pure-TP KV receive consensus (#5627) — CPU-testable mechanism.
+
+In pure TP (LOCAL role) each rank fetches its own KV shard independently and
+``distribute_kv_cache`` is a no-op, so one rank's silent receive timeout used
+to send only that rank into the pipeline's no-KV prefill — desyncing TP
+collectives until the NCCL watchdog killed the worker. These tests cover the
+identical-state verdict exchange: skip guards, both rank roles, KV dropping
+on any cross-rank divergence (primary or CFG companion roles), that the
+exchange runs on the TP group and never on ``world`` (which also holds DP/PP
+ranks working on other requests), and the wiring from both distributed-receive
+entry points.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+    ReceiveRole,
+    _TransferTopoConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeGroup:
+    """Scripted stand-in for one rank's view of a GroupCoordinator."""
+
+    def __init__(self, rank_in_group, world_size, peer_signatures=None, verdict=None):
+        self.rank_in_group = rank_in_group
+        self.world_size = world_size
+        self._peer_signatures = list(peer_signatures or [])
+        self._verdict = verdict
+        self.sent = []
+        self.broadcasts = []
+
+    def recv_object(self, src):
+        return self._peer_signatures.pop(0)
+
+    def send_object(self, obj, dst):
+        self.sent.append((obj, dst))
+
+    def broadcast_object(self, obj=None, src=0):
+        self.broadcasts.append(obj)
+        return obj if obj is not None else self._verdict
+
+    @property
+    def touched(self):
+        return bool(self.sent or self.broadcasts)
+
+
+def _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None, world=None):
+    tp_size = tp_group.world_size if tp_group is not None else 1
+    return _TransferTopoConfig(
+        role=role,
+        tp_active=tp_active,
+        cfg_size=1,
+        cfg_rank=0,
+        cfg_group=None,
+        sp_size=1,
+        sp_rank=0,
+        sp_group=None,
+        world=world,
+        tp_group=tp_group,
+        tp_size=tp_size,
+    )
+
+
+def _manager(topo, need_recv_cache=True):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig(need_recv_cache=need_recv_cache))
+    mgr._topo_config = topo
+    return mgr
+
+
+def _req_with_kv(rid="r1", cfg_text=False):
+    kv = SimpleNamespace(key_cache=[], value_cache=[])
+    sp = SimpleNamespace(past_key_values=kv, kv_metadata={"ropes": [3]})
+    if cfg_text:
+        sp.cfg_text_past_key_values = SimpleNamespace(key_cache=[], value_cache=[])
+    return SimpleNamespace(
+        request_id=rid,
+        past_key_values=kv,
+        kv_metadata={"ropes": [3]},
+        sampling_params=sp,
+    )
+
+
+def _signature_of(req, received=True):
+    return (received, tuple(sorted(OmniKVTransferManager._collect_request_kv_payload(req))))
+
+
+# --------------------------------------------------------------------------- #
+#  Skip guards — no exchange may happen (a one-sided exchange would deadlock)
+# --------------------------------------------------------------------------- #
+
+
+def test_pass_through_outside_pure_tp():
+    group = _FakeGroup(0, 2, peer_signatures=[(False, ())])
+    for topo in (
+        _topo(role=ReceiveRole.LEADER, tp_active=True, tp_group=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=False, tp_group=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None),
+    ):
+        mgr = _manager(topo)
+        assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+        assert not group.touched
+
+
+def test_pass_through_single_rank_group_no_recv_config_and_dummy():
+    group = _FakeGroup(0, 1)
+    assert _manager(_topo(tp_group=group))._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    group2 = _FakeGroup(0, 2, peer_signatures=[(False, ())])
+    mgr = _manager(_topo(tp_group=group2), need_recv_cache=False)
+    assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    mgr3 = _manager(_topo(tp_group=group2))
+    dummy = _req_with_kv(rid=DUMMY_DIFFUSION_REQUEST_ID)
+    assert mgr3._tp_local_receive_consensus(dummy, False) is False
+    assert not group2.touched
+
+
+def test_world_group_is_never_used_for_the_exchange():
+    """DP replicas / PP stages live in ``world`` and work on other requests —
+    the exchange must be scoped to the TP group (review #5636, P1)."""
+    tp_group = _FakeGroup(0, 2, peer_signatures=[(True, ())])
+    world = _FakeGroup(0, 4, peer_signatures=[(False, ()), (False, ()), (False, ())])
+    mgr = _manager(_topo(tp_group=tp_group, world=world))
+    req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace())
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert tp_group.touched
+    assert not world.touched
+
+
+# --------------------------------------------------------------------------- #
+#  Rank 0: collects signatures, decides, broadcasts
+# --------------------------------------------------------------------------- #
+
+
+def test_rank0_drops_kv_when_peer_missed_primary():
+    req = _req_with_kv()
+    peer_sig = (False, ())
+    group = _FakeGroup(0, 2, peer_signatures=[peer_sig])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.broadcasts == [False]  # peers were told to fall back too
+    assert req.past_key_values is None
+    assert req.kv_metadata is None
+    assert req.sampling_params.past_key_values is None
+    assert req.sampling_params.kv_metadata is None
+
+
+def test_rank0_keeps_kv_when_all_ranks_match():
+    req = _req_with_kv()
+    group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.broadcasts == [True]
+    assert req.past_key_values is not None
+    assert req.sampling_params.past_key_values is not None
+
+
+def test_rank0_drops_kv_on_companion_role_divergence():
+    """Primary received everywhere but cfg_text only landed on this rank —
+    divergent branch caches would desync TP collectives (review #5636, P1)."""
+    req = _req_with_kv(cfg_text=True)
+    peer_sig = _signature_of(_req_with_kv(cfg_text=False))  # peer lacks cfg_text
+    group = _FakeGroup(0, 2, peer_signatures=[peer_sig])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.broadcasts == [False]
+    assert req.past_key_values is None
+    assert req.sampling_params.cfg_text_past_key_values is None
+
+
+def test_rank0_keeps_symmetric_partial_state():
+    """cfg_text absent on every rank is consistent — Bagel's empty
+    unconditional branch — and must NOT be cleared or failed."""
+    req = _req_with_kv(cfg_text=False)
+    group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert req.past_key_values is not None
+
+
+def test_rank0_own_miss_fails_verdict_even_if_peers_received():
+    req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace())
+    group = _FakeGroup(0, 3, peer_signatures=[(True, ("past_key_values",)), (True, ("past_key_values",))])
+    mgr = _manager(_topo(tp_group=group))
+    assert mgr._tp_local_receive_consensus(req, False) is False
+    assert group.broadcasts == [False]
+
+
+# --------------------------------------------------------------------------- #
+#  Rank > 0: reports signature, honors verdict
+# --------------------------------------------------------------------------- #
+
+
+def test_follower_rank_drops_kv_on_negative_verdict():
+    req = _req_with_kv()
+    sig_before_clear = _signature_of(req)
+    group = _FakeGroup(1, 2, verdict=False)
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.sent == [(sig_before_clear, 0)]
+    assert req.past_key_values is None
+    assert req.sampling_params.past_key_values is None
+
+
+def test_follower_rank_keeps_kv_on_positive_verdict():
+    req = _req_with_kv()
+    group = _FakeGroup(1, 2, verdict=True)
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.sent == [(_signature_of(req), 0)]
+    assert req.past_key_values is not None
+
+
+# --------------------------------------------------------------------------- #
+#  Wiring — both distributed-receive entry points run the consensus
+# --------------------------------------------------------------------------- #
+
+
+def _wired_manager(monkeypatch, consensus_calls):
+    mgr = _manager(_topo(tp_group=_FakeGroup(0, 2)))
+
+    def fake_consensus(req, received):
+        consensus_calls.append(received)
+        return False
+
+    monkeypatch.setattr(mgr, "_tp_local_receive_consensus", fake_consensus)
+    monkeypatch.setattr(mgr, "receive_multi_kv_cache", lambda *a, **k: True)
+    return mgr
+
+
+def test_receive_multi_kv_cache_distributed_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.receive_multi_kv_cache_distributed(_req_with_kv()) is False
+    assert calls == [True]
+
+
+def test_consume_and_distribute_kv_cache_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.consume_and_distribute_kv_cache(_req_with_kv()) is False
+    assert calls == [True]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -682,6 +682,10 @@ class TestDistributedReceive:
         mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
         req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
         world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
+        # broadcast_value=True: rank 0's verdict says every TP rank has the
+        # same KV state, so the per-rank receive result stands (see #5627
+        # consensus). The exchange runs on the TP group, not world.
+        tp_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, broadcast_value=True)
         mgr.receive_multi_kv_cache = MagicMock(return_value=True)
 
         with (
@@ -704,10 +708,19 @@ class TestDistributedReceive:
                 return_value=0,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
+            patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group),
         ):
             assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
+        # Pure TP still fetches per-rank (no leader→follower distribution),
+        # but each rank now reports its KV-state signature on the TP group
+        # for the identical-state verdict instead of silently diverging on a
+        # miss (#5627). World must stay untouched — it holds DP/PP ranks.
+        assert len(tp_group.send_calls) == 1
+        dst, sig = tp_group.send_calls[0]
+        assert dst == 0 and sig[0] is True
+        assert world_group.send_calls == []
 
     # ── SP-only scenarios ────────────────────────────────────────────
 

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -32,10 +32,28 @@ PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
 SSIM_THRESHOLD = 0.97
-PSNR_THRESHOLD = 30.0
+PSNR_THRESHOLD_CUDA = 30.0
+PSNR_THRESHOLD_NPU = 26.0
+
+
+def _psnr_threshold() -> float:
+    from vllm_omni.platforms import current_omni_platform
+
+    return PSNR_THRESHOLD_NPU if current_omni_platform.is_npu() else PSNR_THRESHOLD_CUDA
+
 
-BASELINE_PATH = get_asset_path("hunyuan/hunyuan_baseline.png")
 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
+
+
+def _baseline_path() -> Path:
+    from vllm_omni.platforms import current_omni_platform
+
+    suffix = "_npu" if current_omni_platform.is_npu() else ""
+    path = get_asset_path(f"hunyuan/hunyuan_baseline{suffix}.png")
+    assert path.exists(), f"Baseline image not found at {path}"
+    return path
+
+
 _OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
 
 # DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
@@ -74,6 +92,23 @@ _DEPLOY_CONFIG = {
             },
         },
     ],
+    "platforms": {
+        "npu": {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "gpu_memory_utilization": 0.65,
+                    "devices": "0,1,2,3",
+                    "moe_backend": "auto",
+                    "max_num_batched_tokens": 32768,
+                    "parallel_config": {
+                        "tensor_parallel_size": 4,
+                        "enable_expert_parallel": True,
+                    },
+                },
+            ],
+        },
+    },
 }
 
 
@@ -216,8 +251,8 @@ def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, out
 
 
 def _assert_against_baseline(image: Image.Image, label: str) -> None:
-    assert BASELINE_PATH.exists(), f"Baseline image not found at {BASELINE_PATH}"
-    baseline_image = Image.open(BASELINE_PATH).convert("RGB")
+    baseline_path = _baseline_path()
+    baseline_image = Image.open(baseline_path).convert("RGB")
 
     assert_images_pixel_close(
         model_name=f"{MODEL_NAME} ({label} vs baseline)",
@@ -231,11 +266,11 @@ def _assert_against_baseline(image: Image.Image, label: str) -> None:
         vllm_image=image,
         diffusers_image=baseline_image,
         ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
+        psnr_threshold=_psnr_threshold(),
     )
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
     model = _model_name()
     output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
@@ -249,7 +284,7 @@ def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> N
     _assert_against_baseline(image, "online")
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
     model = _model_name()
     output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)

tests/engine/test_single_stage_mode.py

@@ -12,6 +12,7 @@ import pytest
 from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
@@ -476,6 +477,89 @@ class TestSingleStageModeDetection:
         assert not hasattr(engine, "_omni_master_server")
 
 
+# ---------------------------------------------------------------------------
+# Endpoint-restriction resolution vs. tri-state trust_remote_code (#5495)
+# ---------------------------------------------------------------------------
+
+
+class TestEndpointRestrictionsTrustRemoteCode:
+    """Regression tests for issue #5495.
+
+    ``AsyncOmniEngine.__init__`` takes ``trust_remote_code: bool | None`` where
+    ``None`` means "not specified" (the CLI default, since ``--trust-remote-code``
+    is ``store_true`` and its absence is mapped to ``None``). Endpoint-restriction
+    resolution eagerly loads the HF config via vLLM's ``get_config``, which does
+    ``trust_remote_code |= ...`` internally and raises ``TypeError`` on ``None``.
+    That used to be swallowed, leaving ``endpoint_restrictions`` empty so
+    ``/v1/completions`` was never shut down and the request crashed stage-1.
+
+    ``__init__`` must collapse the ``None`` "not specified" case to ``False`` at
+    the ``get_pipeline_config`` call site so the restriction is still computed.
+    """
+
+    @pytest.fixture(autouse=True)
+    def clear_config_factory_caches(self):
+        """Isolate the process-wide ``functools.cache`` on StageConfigFactory.
+
+        ``get_hf_config`` / ``try_infer_model_type`` cache per-model results for
+        the whole process. These tests patch ``get_config`` with a synthetic
+        Qwen3-Omni config for ``"fake-model"``; without clearing on teardown the
+        synthetic entry would leak to later tests in the same worker (after the
+        mock is restored), making the suite order-dependent. Clear both before
+        and after.
+        """
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+        yield
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+
+    def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine:
+        mocker.patch.object(
+            AsyncOmniEngine,
+            "_resolve_stage_configs",
+            return_value=("/fake/path", [_make_stage_cfg(0)]),
+        )
+        mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
+        mocker.patch("threading.Thread")
+        mocker.patch("concurrent.futures.Future")
+        return AsyncOmniEngine(model="fake-model", **kwargs)
+
+    @staticmethod
+    def _install_qwen3_omni_config(mocker: MockerFixture) -> None:
+        """Patch get_config to a thinker+talker config, mirroring vLLM's
+        ``trust_remote_code |= ...`` (raises TypeError on None)."""
+        from transformers import Qwen3OmniMoeConfig
+
+        def fake_get_config(model, trust_remote_code, **_):
+            trust_remote_code |= False  # TypeError if None reaches here
+            return Qwen3OmniMoeConfig(enable_audio_output=True)
+
+        mocker.patch("vllm_omni.config.config_factory.get_config", side_effect=fake_get_config)
+
+    @pytest.mark.parametrize("trc_kwargs", [{}, {"trust_remote_code": False}, {"trust_remote_code": True}])
+    def test_completions_restriction_survives(self, mocker: MockerFixture, trc_kwargs: dict):
+        """COMPLETIONS restriction is computed regardless of how (or whether)
+        trust_remote_code is passed — including the unset case that resolves to
+        the engine's ``None`` default."""
+        from vllm_omni.config.endpoint_policy import OmniServingCapability
+
+        self._install_qwen3_omni_config(mocker)
+        engine = self._make_engine_no_thread(mocker, **trc_kwargs)
+        assert any(r.capability is OmniServingCapability.COMPLETIONS for r in engine.endpoint_restrictions), (
+            f"COMPLETIONS restriction dropped for {trc_kwargs or 'unset (None default)'}"
+        )
+
+    def test_unset_trust_remote_code_defaults_to_none(self, mocker: MockerFixture):
+        """Guard the tri-state premise: the engine keeps ``None`` (not ``False``)
+        when trust_remote_code is not passed, so the call-site coercion is what
+        prevents the crash."""
+        import inspect
+
+        default = inspect.signature(AsyncOmniEngine.__init__).parameters["trust_remote_code"].default
+        assert default is None
+
+
 # ---------------------------------------------------------------------------
 # AsyncOmniEngine single-stage initialization paths
 # ---------------------------------------------------------------------------

tests/entrypoints/openai_api/test_video_server.py

@@ -12,6 +12,7 @@ import os
 import threading
 import time
 from contextlib import asynccontextmanager
+from pathlib import Path
 from types import SimpleNamespace
 
 import numpy as np
@@ -62,11 +63,19 @@ class FakeAsyncOmni:
         self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
         self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
         self.captured_prompt = None
+        self.captured_reference_video_bytes = None
         self.captured_sampling_params_list = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_prompt = prompt
         self.captured_sampling_params_list = sampling_params_list
+        reference_videos = prompt.get("multi_modal_data", {}).get("video")
+        if (
+            isinstance(reference_videos, list)
+            and reference_videos
+            and all(isinstance(item, str) for item in reference_videos)
+        ):
+            self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos]
         num_outputs = sampling_params_list[0].num_outputs_per_prompt
         videos = [object() for _ in range(num_outputs)]
         yield MockVideoResult(videos)
@@ -501,6 +510,44 @@ def test_v2v_video_generation_with_video_reference_form(test_client, mocker: Moc
     assert input_video[0].size == (32, 24)
 
 
+@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"])
+def test_multi_video_generation_preserves_uploaded_files_until_generation(
+    endpoint,
+    test_client,
+    mocker: MockerFixture,
+):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        endpoint,
+        data={
+            "prompt": "Composite the subject from the first video into the second.",
+            "extra_params": json.dumps({"task": "ref2va", "duration": 15.0}),
+        },
+        files=[
+            ("input_references", ("subject.mp4", b"subject-video", "video/mp4")),
+            ("input_references", ("background.mov", b"background-video", "video/quicktime")),
+        ],
+    )
+
+    assert response.status_code == 200
+    if endpoint.endswith("/sync"):
+        assert response.content == b"fake-video"
+    else:
+        video_id = response.json()["id"]
+        _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_videos = engine.captured_prompt["multi_modal_data"]["video"]
+    assert engine.captured_reference_video_bytes == [b"subject-video", b"background-video"]
+    assert len(input_videos) == 2
+    assert all(not Path(path).exists() for path in input_videos)
+    assert engine.captured_sampling_params_list[0].extra_args["task"] == "ref2va"
+    assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
+
+
 def test_decode_video_bytes_can_keep_last_frames():
     from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
 

tests/model_tests/diffusion/test_alignment.py

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

tests/utils/test_tracking_parser.py

@@ -219,6 +219,46 @@ def test_store_true_explicit():
     assert ns.verbose is True
 
 
+def test_unset_trust_remote_code_falls_back_to_engine_default_none():
+    """Regression guard for issue #5495 (the CLI-plumbing half).
+
+    The CI test ``test_completions_rejected_for_thinker_talker`` does NOT pass
+    ``--trust-remote-code``. Because it is a ``store_true`` flag, an omitted
+    value is dropped from ``get_explicit_kwargs_dict()`` (it is not user-typed),
+    so it is never forwarded to ``AsyncOmni(**kwargs)`` and the engine falls back
+    to ``AsyncOmniEngine.__init__``'s default of ``None`` — NOT argparse's
+    ``False``. That ``None`` is what used to break endpoint-restriction
+    resolution (see ``TestQwen3OmniPipeline.
+    test_endpoint_restrictions_survive_unset_trust_remote_code``).
+
+    This test locks in the "CI does not pass the flag => effective value is
+    None" premise so a future refactor cannot silently change it.
+    """
+    import inspect
+
+    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+    p = TrackingArgumentParser()
+    # Mirror how --trust-remote-code is registered upstream (store_true).
+    p.add_argument("--trust-remote-code", action="store_true")
+
+    # Scenario A: CI — flag omitted.
+    ns = p.parse_args([])
+    assert "trust_remote_code" not in ns.explicit_keys
+    kwargs = ns.get_explicit_kwargs_dict()
+    assert "trust_remote_code" not in kwargs
+    engine_default = inspect.signature(AsyncOmniEngine.__init__).parameters["trust_remote_code"].default
+    assert engine_default is None
+    effective = kwargs.get("trust_remote_code", engine_default)
+    assert effective is None  # <- the value that reached the buggy restriction path
+
+    # Scenario B: flag passed (what earlier repros used) => True, restrictions fine.
+    ns2 = p.parse_args(["--trust-remote-code"])
+    assert "trust_remote_code" in ns2.explicit_keys
+    kwargs2 = ns2.get_explicit_kwargs_dict()
+    assert kwargs2.get("trust_remote_code", engine_default) is True
+
+
 def test_store_false_default():
     """Ensure that store false is handled correctly when omitted."""
     p = TrackingArgumentParser()

vllm_omni/config/omni_config.py

@@ -113,6 +113,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False):
     ulysses_mode: str
     cfg_parallel_size: int
     vae_patch_parallel_size: int
+    text_encoder_tp_size: int
     use_hsdp: bool
     mask_sp_padding: bool
     hsdp_shard_size: int
@@ -370,6 +371,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     ulysses_mode: str = "strict"
     cfg_parallel_size: int = Field(default=1, ge=1)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
+    text_encoder_tp_size: int = Field(default=1, ge=1)
     vae_parallel_mode: str = "tile"
     use_hsdp: bool = False
     mask_sp_padding: bool = False

vllm_omni/config/stage_config.py

@@ -377,6 +377,7 @@ class StageDeployConfig:
     cfg_parallel_size: int | None = None
     vae_patch_parallel_size: int | None = None
     vae_parallel_mode: str | None = None
+    text_encoder_tp_size: int | None = None
     use_hsdp: bool | None = None
     hsdp_shard_size: int | None = None
     hsdp_replicate_size: int | None = None

vllm_omni/diffusion/attention/backends/abstract.py

@@ -79,6 +79,10 @@ class AttentionMetadata:
     # Well-known optional keys (convention, not required on all forwards):
     #   "kv_cache_dtype": str | None — quantized KV dtype (e.g. "fp8"); backends
     #     decide whether/how to apply.
+    #   "cu_seqlens_q" / "cu_seqlens_k": int32 CUDA tensors describing packed
+    #     variable-length query/key sequences for FlashAttention.
+    #   "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with
+    #     the packed cu_seqlens tensors.
 
     # Piecewise attention metadata (mixed causal/full masks).
     # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -141,6 +141,46 @@ class FlashAttentionImpl(AttentionImpl):
         out_unpad = self._unwrap_flash_output(out_unpad)
         return _pad_input(out_unpad, indices_q, query.size(0), query_length)
 
+    def _forward_varlen_packed(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        *,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+    ) -> torch.Tensor:
+        """Run FlashAttention directly on an already packed sequence.
+
+        Some diffusion transformers already maintain exact packed-document
+        boundaries. Reusing those boundaries avoids rebuilding boolean masks
+        and gathering/scattering Q/K/V in every attention layer.
+        """
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            flash_attn_varlen_func,
+        )
+
+        if flash_attn_varlen_func is None:
+            raise ImportError("Packed variable-length attention requires flash_attn_varlen_func")
+        if query.shape[0] != 1 or key.shape[0] != 1 or value.shape[0] != 1:
+            raise ValueError("Packed variable-length attention currently requires batch size 1")
+
+        out = flash_attn_varlen_func(
+            q=query.flatten(0, 1),
+            k=key.flatten(0, 1),
+            v=value.flatten(0, 1),
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        out = self._unwrap_flash_output(out)
+        return out.reshape_as(query)
+
     def _forward_varlen_dense(
         self,
         query: torch.Tensor,
@@ -206,6 +246,7 @@ class FlashAttentionImpl(AttentionImpl):
 
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+        extra = attn_metadata.extra if attn_metadata is not None else {}
 
         # Try piecewise attention
         if full_attn_spans is not None:
@@ -233,6 +274,22 @@ class FlashAttentionImpl(AttentionImpl):
                 query_ranges=attn_metadata.query_ranges,
             )
 
+        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:
+            if len(present_packed_keys) != len(packed_keys):
+                missing = sorted(set(packed_keys) - set(present_packed_keys))
+                raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+            return self._forward_varlen_packed(
+                query,
+                key,
+                value,
+                cu_seqlens_q=extra["cu_seqlens_q"],
+                cu_seqlens_k=extra["cu_seqlens_k"],
+                max_seqlen_q=extra["max_seqlen_q"],
+                max_seqlen_k=extra["max_seqlen_k"],
+            )
+
         if attention_mask is not None and torch.any(~attention_mask):
             return self._forward_varlen_masked(
                 query,

vllm_omni/diffusion/attention/backends/ring/ring_globals.py

@@ -23,6 +23,19 @@ HAS_FA3 = False
 fa3_fwd_func = None  # Low-level forward function (_flash_attn_forward)
 fa3_attn_func = None  # High-level attention function (flash_attn_func)
 
+# FA4 detection. The CuTe API returns LSE directly, so it can participate in
+# Ring Attention's numerically stable block-wise output accumulation.
+HAS_FA4 = False
+fa4_attn_func = None
+try:
+    from flash_attn.cute import flash_attn_func as fa4_attn_func  # noqa: F401
+
+    HAS_FA4 = True
+except Exception:
+    # Optional CuTe/CUTLASS/Quack components can be importable but
+    # ABI-incompatible. Treat the whole optional backend as unavailable.
+    pass
+
 # Try flash_attn_interface first (from flash-attention source build)
 try:
     from flash_attn_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401

vllm_omni/diffusion/attention/backends/ring/ring_kernels.py

@@ -9,9 +9,11 @@ import torch
 from .ring_globals import (
     HAS_AITER,
     HAS_FA3,
+    HAS_FA4,
     HAS_FLASH_ATTN,
     HAS_FLASHINFER,
     fa3_fwd_func,
+    fa4_attn_func,
 )
 
 _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention
@@ -190,6 +192,40 @@ def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap,
 flash_attn3_func_forward = fa3_forward
 
 
+def flash_attn4_func_forward(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    softcap=None,
+    alibi_slopes=None,
+    return_softmax=False,
+):
+    """FA4 forward pass returning the LSE required by Ring Attention."""
+    del dropout_p, alibi_slopes, return_softmax
+    assert HAS_FA4, "FA4 is not available"
+    assert fa4_attn_func is not None, "FA4 high-level API is not available"
+
+    left, right = window_size if window_size else (-1, -1)
+    out, softmax_lse = fa4_attn_func(
+        q,
+        k,
+        v,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size=(
+            None if left == -1 else left,
+            None if right == -1 else right,
+        ),
+        softcap=softcap or 0.0,
+        return_lse=True,
+    )
+    return out, softmax_lse
+
+
 def flash_attn_forward_aiter(
     q,
     k,

vllm_omni/diffusion/attention/backends/ring/ring_selector.py

@@ -15,6 +15,7 @@ from .ring_globals import (
 )
 from .ring_kernels import (
     flash_attn3_func_forward,
+    flash_attn4_func_forward,
     flash_attn_forward,
     flash_attn_forward_aiter,
     flashinfer_attn_forward,
@@ -32,6 +33,7 @@ class AttnType(Enum):
     AITER = "aiter"
     FA = "fa"
     FA3 = "fa3"
+    FA4 = "fa4"
     FLASHINFER = "flashinfer"
     TORCH = "torch"
     SAGE_AUTO = "sage_auto"
@@ -77,6 +79,9 @@ def select_flash_attn_impl(
     elif impl_type == AttnType.FA3:
         return flash_attn3_func_forward
 
+    elif impl_type == AttnType.FA4:
+        return flash_attn4_func_forward
+
     elif impl_type == AttnType.FLASHINFER:
         return flashinfer_attn_forward
 

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -22,6 +22,15 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+
+def _is_blackwell() -> bool:
+    try:
+        capability = current_omni_platform.get_device_capability()
+    except Exception:
+        return False
+    return capability is not None and capability.major >= 10
+
+
 # Flash Attention function detection with fallback chain
 flash_attn_func = None
 flash_attn_varlen_func = None
@@ -48,12 +57,27 @@ elif current_omni_platform.is_musa():
     except (ImportError, ModuleNotFoundError):
         pass
 else:
-    # CUDA: try FA3 -> FA2 fallback chain
+    # CUDA: prefer the CuTe-based FA4 package when present.  This matters on
+    # Blackwell, where an importable Hopper-only FA3 wheel may otherwise be
+    # selected and fail at launch with "no kernel image is available".
+    if _is_blackwell():
+        try:
+            from flash_attn.cute import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+
+            logger.info("Using CuTe FlashAttention-4 on Blackwell")
+        except Exception as exc:
+            # Optional FA4 dependencies may be present but ABI-incompatible
+            # (for example, mismatched CUTLASS DSL and Quack builds). Treat
+            # that as unavailable so backend discovery can continue.
+            logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+
+    # Try FA3 -> FA2 fallback chain.
     # Try FA3 from fa3-fwd PyPI package
-    try:
-        from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
-    except (ImportError, ModuleNotFoundError):
-        pass
+    if flash_attn_func is None:
+        try:
+            from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
 
     # Fallback: Try FA3 from flash-attention source build
     if flash_attn_func is None:
@@ -100,7 +124,16 @@ def is_flash_attn_installed() -> bool:
     Shared by CUDA/ROCm/MUSA platforms.
     """
     try:
-        # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
+        # Check for any FA backend: FA4 (flash_attn.cute), FA3
+        # (fa3_fwd_interface, flash_attn_interface), or FA2 (flash_attn).
+        if _is_blackwell():
+            try:
+                from flash_attn.cute import flash_attn_varlen_func  # noqa: F401
+
+                return True
+            except Exception as exc:
+                logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+
         # Try FA3 from fa3-fwd PyPI package
         try:
             import fa3_fwd_interface  # noqa: F401

vllm_omni/diffusion/attention/parallel/ring.py

@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FA4, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -145,9 +145,12 @@ class RingParallelAttention:
 
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
-        # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
-        # On ROCm, use AITER
-        if HAS_FA3:
+        # Prefer FA4 on Blackwell. An importable Hopper-only FA3 wheel can
+        # otherwise be selected and fail at launch with "no kernel image".
+        if HAS_FA4 and query.is_cuda and torch.cuda.get_device_capability(query.device)[0] >= 10:
+            attn_type = AttnType.FA4
+        # Prefer FA3 over FA2 on Ampere/Ada/Hopper. On ROCm, use AITER.
+        elif HAS_FA3:
             attn_type = AttnType.FA3
         elif HAS_AITER:
             attn_type = AttnType.AITER

vllm_omni/diffusion/data.py

@@ -192,6 +192,16 @@ class DiffusionParallelConfig:
     vae_patch_parallel_size: int = 1
     """Number of ranks used for VAE patch/tile parallelism (decode/encode)."""
 
+    text_encoder_tp_size: int = 1
+    """Number of ranks used to tensor-parallel shard the diffusion text encoder.
+
+    Ranks are the first ``text_encoder_tp_size`` DiT ranks.  Defaults to 1,
+    which keeps the encoder fully resident on the DiT main rank (historical
+    behavior).  Values > 1 shard the Qwen3-VL encoder across the encoder TP
+    ranks and run the encode with distributed collectives over the encoder
+    process group.
+    """
+
     vae_parallel_mode: str = "tile"
     """VAE parallel decode strategy.
 

vllm_omni/diffusion/model_metadata.py

@@ -36,6 +36,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES,
     ),
+    "MiniMaxH3Pipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=1,
+    ),
     "WanPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),

vllm_omni/diffusion/models/minimax_h3/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 audio-video diffusion support."""
+
+from .pipeline_minimax_h3 import (
+    MiniMaxH3Pipeline,
+    get_minimax_h3_post_process_func,
+)
+
+__all__ = [
+    "MiniMaxH3Pipeline",
+    "get_minimax_h3_post_process_func",
+]

vllm_omni/diffusion/models/minimax_h3/condition_noise.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 visual/audio condition-noise augmentation.
+
+The request's condition timestep is applied to both the tensor value and the
+DiT timestep. Tokenizer artifacts remain clean
+and reusable; this module materializes the fixed noised anchors immediately
+before the denoise loop.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from .packed_tokens import minimax_h3_patchify_video_latent
+
+# Channel-major packed audio rows always carry a stereo layout.
+MINIMAX_H3_AUDIO_COND_CHANNELS = 2
+
+
+def minimax_h3_imgvid_cond_noise_aug_rows(
+    clean_rows: torch.Tensor,
+    *,
+    condition_shapes: Sequence[Sequence[int]],
+    target_latent_t: int,
+    imgvid_cond_num_frames: int,
+    seed: int,
+    noise_aug: float,
+) -> torch.Tensor:
+    """Apply the imgvid-condition RF noise recipe to packed clean rows.
+
+    ``condition_shapes`` contains ``(latent_t, latent_h, latent_w)`` in packed
+    visual-condition order. A new CPU generator with the same row seed is
+    created for every condition. Under the dependent-noise policy, each draw
+    uses the target temporal length plus the template's imgvid-condition frame
+    count, then slices the prefix matching the current condition.
+    """
+
+    noise_aug = float(noise_aug)
+    if not 0.0 <= noise_aug <= 1.0:
+        raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
+    if noise_aug == 1.0:
+        return clean_rows
+    if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 96:
+        raise ValueError(f"clean imgvid condition rows must have shape [n, 96], got {list(clean_rows.shape)}")
+
+    target_latent_t = int(target_latent_t)
+    imgvid_cond_num_frames = int(imgvid_cond_num_frames)
+    if target_latent_t <= 0:
+        raise ValueError(f"target_latent_t must be positive, got {target_latent_t}")
+    if imgvid_cond_num_frames <= 0:
+        raise ValueError(
+            f"imgvid_cond_num_frames must be positive when condition rows exist, got {imgvid_cond_num_frames}"
+        )
+
+    parsed_shapes: list[tuple[int, int, int]] = []
+    expected_rows = 0
+    for raw_shape in condition_shapes:
+        if len(raw_shape) != 3:
+            raise ValueError(
+                f"each imgvid condition shape must be (latent_t, latent_h, latent_w), got {list(raw_shape)}"
+            )
+        latent_t, latent_h, latent_w = (int(value) for value in raw_shape)
+        if latent_t <= 0 or latent_h <= 0 or latent_w <= 0:
+            raise ValueError(f"imgvid condition shape must be positive, got {list(raw_shape)}")
+        if latent_h % 2 or latent_w % 2:
+            raise ValueError(
+                f"imgvid condition spatial dimensions must be divisible by 2, got {(latent_t, latent_h, latent_w)}"
+            )
+        parsed_shapes.append((latent_t, latent_h, latent_w))
+        expected_rows += latent_t * (latent_h // 2) * (latent_w // 2)
+    if not parsed_shapes:
+        raise ValueError("condition_shapes must not be empty")
+    if int(clean_rows.shape[0]) != expected_rows:
+        raise ValueError(
+            f"clean imgvid condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}"
+        )
+
+    output_device = clean_rows.device
+    clean_rows = clean_rows.detach().to(device="cpu", dtype=torch.float32)
+    out: list[torch.Tensor] = []
+    row_offset = 0
+    timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu")
+    for latent_t, latent_h, latent_w in parsed_shapes:
+        full_t = target_latent_t + imgvid_cond_num_frames
+        if full_t < latent_t:
+            raise ValueError(f"condition latent_t {latent_t} exceeds the noise draw length {full_t}")
+        generator = torch.Generator(device="cpu").manual_seed(int(seed))
+        noise = torch.randn(
+            1,
+            24,
+            full_t,
+            latent_h,
+            latent_w,
+            generator=generator,
+            dtype=torch.float32,
+            device="cpu",
+        )[:, :, :latent_t]
+        noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to(dtype=torch.float32)
+        row_count = int(noise_rows.shape[0])
+        clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32)
+        out.append(timestep * clean_part + (1.0 - timestep) * noise_rows)
+        row_offset += row_count
+    return torch.cat(out, dim=0).to(device=output_device, dtype=torch.float32).contiguous()
+
+
+def minimax_h3_audio_cond_noise_aug_rows(
+    clean_rows: torch.Tensor,
+    *,
+    condition_audio_t: Sequence[int],
+    seed: int,
+    noise_aug: float,
+) -> torch.Tensor:
+    """Apply the audio-condition RF noise recipe to packed clean rows.
+
+    ``condition_audio_t`` contains the latent T of each audio-bearing
+    condition in canonical request order. Noise is drawn per condition
+    element, with a fresh CPU generator seeded with ``seed + 1`` for every
+    element.  Consequently each condition restarts the
+    same RNG stream; concatenating the rows and drawing once would be
+    numerically different for ordered multi-reference requests.
+
+    The mix is intentionally evaluated on CPU in fp32 before the packed rows
+    are transferred to the DiT device.
+    """
+
+    noise_aug = float(noise_aug)
+    if not 0.0 <= noise_aug <= 1.0:
+        raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
+    if noise_aug == 1.0:
+        return clean_rows
+    if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 32:
+        raise ValueError(f"clean audio condition rows must have shape [n, 32], got {list(clean_rows.shape)}")
+
+    audio_channels = MINIMAX_H3_AUDIO_COND_CHANNELS
+    parsed_audio_t = [int(value) for value in condition_audio_t]
+    if not parsed_audio_t:
+        raise ValueError("condition_audio_t must not be empty")
+    if any(value <= 0 for value in parsed_audio_t):
+        raise ValueError(f"condition audio latent lengths must be positive, got {parsed_audio_t}")
+    expected_rows = audio_channels * sum(parsed_audio_t)
+    if int(clean_rows.shape[0]) != expected_rows:
+        raise ValueError(f"clean audio condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}")
+
+    out: list[torch.Tensor] = []
+    row_offset = 0
+    timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu")
+    for audio_t in parsed_audio_t:
+        row_count = audio_channels * audio_t
+        clean_part = clean_rows[row_offset : row_offset + row_count].detach().to(device="cpu", dtype=torch.float32)
+        generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1)
+        noise = torch.randn(
+            clean_part.shape,
+            generator=generator,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        out.append(timestep * clean_part + (1.0 - timestep) * noise)
+        row_offset += row_count
+    return torch.cat(out, dim=0).to(device=clean_rows.device, dtype=torch.float32).contiguous()
+
+
+__all__ = [
+    "minimax_h3_audio_cond_noise_aug_rows",
+    "minimax_h3_imgvid_cond_noise_aug_rows",
+]

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 cfg-distilled full denoise loop.
+
+Per step, the positive presentation is forwarded exactly once. Video and audio
+target rows chain through the Euler-eta0 update while visual and audio condition
+rows stay pinned to their noised step-0 anchors.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from contextlib import AbstractContextManager, nullcontext
+from typing import Any
+
+import torch
+
+from .scheduling_minimax_h3_euler_ancestral import (
+    minimax_h3_euler_eta0_step,
+    minimax_h3_rf_v_to_x0,
+)
+
+MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
+# ref2va audio reference anchor timestep
+MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0
+# Packed row widths: video rows are [1,2,2]-patchified 24-channel latents
+# (24 * 1 * 2 * 2 = 96); audio rows carry the 32-dim audio latent.
+MINIMAX_H3_VIDEO_ROW_WIDTH = 96
+MINIMAX_H3_AUDIO_ROW_WIDTH = 32
+
+
+class MiniMaxH3DenoiseBranch:
+    """Static per-branch state: packed layout + fixed forward kwargs.
+
+    `packed` is a minimax_h3_packed_sequence(...) result (or equivalent layout
+    dict); `text_embeddings` is the branch's [text_len, 5120] hidden states;
+    `token_tags` must already carry any fl2va vision-span overrides.
+    """
+
+    def __init__(
+        self,
+        *,
+        packed: dict[str, torch.Tensor],
+        text_embeddings: torch.Tensor,
+        token_tags: torch.Tensor,
+        device: torch.device,
+    ) -> None:
+        seq_len = int(packed["seq_len"])
+        self.seq_len = seq_len
+        self.img_pos = packed["img_pos"].view(-1).to(torch.long)
+        self.audio_pos = packed["audio_pos"].view(-1).to(torch.long)
+        self.update_mask = packed["update_mask"].view(-1).to(torch.bool)
+        # ref2va: audio_pos may include reference-audio anchor rows
+        # (audio_update_mask False); absent means all rows are targets.
+        if "audio_update_mask" in packed:
+            self.audio_update_mask = packed["audio_update_mask"].view(-1).to(torch.bool)
+        else:
+            self.audio_update_mask = torch.ones(self.audio_pos.shape[0], dtype=torch.bool)
+        text_len = int(packed["text_pos"].view(-1).shape[0])
+        if list(text_embeddings.shape)[0] != text_len:
+            raise ValueError(f"text_embeddings rows {list(text_embeddings.shape)} != packed text_len {text_len}")
+        if int(token_tags.view(-1).shape[0]) != seq_len:
+            raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}")
+        cu = packed["cu_seqlens"].to(torch.int32)
+        self.img_pos_dev = self.img_pos.to(device)
+        self.audio_pos_dev = self.audio_pos.to(device)
+        self.update_mask_dev = self.update_mask.to(device)
+        self.audio_update_mask_dev = self.audio_update_mask.to(device)
+        self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
+        self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
+        text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
+        self.static_kwargs: dict[str, Any] = {
+            "img_position_ids": packed["img_position_ids"][None].to(device),
+            "update_mask": self.update_mask_dev,
+            "token_tags": token_tags.view(-1).to(torch.long).to(device),
+            "skip_mask_out_condition": False,
+            "prompt_embeds": text_embeddings.to(device),
+            "img_pos_info": {"position_ids": self.img_pos_dev},
+            "audio_pos_info": {"position_ids": self.audio_pos_dev},
+            "text_pos_info": {"position_ids": text_pos_dev},
+            "img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev},
+            "packed_seq_params": {
+                "cu_seqlens_q": cu.to(device),
+                "max_seqlen_q": int(cu[1]),
+            },
+            "refiner_packed_seq_params": {
+                "cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device),
+                "max_seqlen_q": text_len,
+            },
+        }
+
+    def forward_kwargs(
+        self,
+        *,
+        video_rows: torch.Tensor,
+        audio_rows: torch.Tensor,
+        t_video: float,
+        t_audio: float,
+        imgvid_cond_timestep: float,
+        audio_ref_cond_timestep: float,
+    ) -> dict[str, Any]:
+        x = self.x_base.clone()
+        x[0].index_copy_(0, self.img_pos_dev, video_rows)
+        audio_x = self.audio_x_base.clone()
+        audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows)
+        # Packed-sequence timestep semantics: non-media rows (text and
+        # padding) inherit the current video timestep. Later steps must reuse
+        # the previous step's updated rows; re-initializing from zeros is only
+        # valid at step 0.
+        timesteps = torch.full(
+            (self.seq_len,),
+            float(t_video),
+            dtype=torch.float32,
+            device=x.device,
+        )
+        timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video
+        timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep
+        timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio
+        timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep
+        unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True)
+        return {
+            **self.static_kwargs,
+            "x": x,
+            "audio_x": audio_x,
+            "unique_timesteps": unique_timesteps,
+            "inverse_indices": inverse_indices,
+        }
+
+
+def minimax_h3_denoise_loop(
+    *,
+    model: Any,
+    positive: MiniMaxH3DenoiseBranch,
+    initial_video_rows: torch.Tensor,
+    initial_audio_rows: torch.Tensor,
+    keyframe_cond_rows: torch.Tensor | None,
+    audio_ref_rows: torch.Tensor | None = None,
+    sigmas_video: list[float],
+    sigmas_audio: list[float],
+    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,
+    step_profiler: Callable[[int], AbstractContextManager] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run the full denoise loop; returns final (video_rows, audio_rows).
+
+    ``initial_video_rows`` covers all image rows of the positive layout. For a
+    conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to
+    pin those rows across every step. The model's raw positive velocity is the
+    update signal; MiniMax H3 only supports cfg-distilled checkpoints.
+    """
+    if len(sigmas_video) != len(sigmas_audio):
+        raise ValueError("video/audio sigma schedules must have equal length")
+    if len(sigmas_video) < 2:
+        raise ValueError("sigma schedules need at least 2 entries")
+    n_cond = int((~positive.update_mask).sum())
+    if keyframe_cond_rows is None:
+        if n_cond != 0:
+            raise ValueError(f"layout has {n_cond} cond rows but keyframe_cond_rows is None")
+    else:
+        if int(keyframe_cond_rows.shape[0]) != n_cond:
+            raise ValueError(f"keyframe_cond_rows {int(keyframe_cond_rows.shape[0])} != layout cond rows {n_cond}")
+    video_rows = initial_video_rows.to(device=device, dtype=torch.float32).clone()
+    audio_rows = initial_audio_rows.to(device=device, dtype=torch.float32).clone()
+    update = positive.update_mask_dev
+    audio_update = positive.audio_update_mask_dev
+    if int(video_rows.shape[0]) != int(positive.img_pos.shape[0]):
+        raise ValueError(
+            f"initial video rows {int(video_rows.shape[0])} != positive layout rows {int(positive.img_pos.shape[0])}"
+        )
+    if int(audio_rows.shape[0]) != int(positive.audio_pos.shape[0]):
+        raise ValueError(
+            f"initial audio rows {int(audio_rows.shape[0])} != positive layout rows {int(positive.audio_pos.shape[0])}"
+        )
+    n_audio_ref = int((~positive.audio_update_mask).sum())
+    if audio_ref_rows is None:
+        if n_audio_ref != 0:
+            raise ValueError(f"layout has {n_audio_ref} audio ref rows but audio_ref_rows is None")
+        audio_anchor = None
+    else:
+        if int(audio_ref_rows.shape[0]) != n_audio_ref:
+            raise ValueError(f"audio_ref_rows {int(audio_ref_rows.shape[0])} != layout audio ref rows {n_audio_ref}")
+        audio_anchor = audio_ref_rows.to(device=device, dtype=torch.float32)
+    cond_anchor = keyframe_cond_rows.to(device=device, dtype=torch.float32) if keyframe_cond_rows is not None else None
+    if cond_anchor is not None:
+        video_rows[~update] = cond_anchor
+    if audio_anchor is not None:
+        audio_rows[~audio_update] = audio_anchor
+
+    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:
+            s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
+            s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
+            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))
+
+            fk = positive.forward_kwargs(
+                video_rows=video_rows,
+                audio_rows=audio_rows,
+                t_video=t_v,
+                t_audio=t_a,
+                imgvid_cond_timestep=imgvid_cond_t,
+                audio_ref_cond_timestep=audio_ref_cond_t,
+            )
+            with torch.inference_mode():
+                v_video, v_audio = model(**fk)
+            mv_video_t = v_video.float()[update]
+            mv_audio_t = v_audio.float()[audio_update]
+
+            x0_video = minimax_h3_rf_v_to_x0(
+                video_rows[update],
+                mv_video_t,
+                torch.tensor(t_v, dtype=torch.float32, device=device),
+            )
+            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:
+                video_rows[~update] = cond_anchor  # per-step imgvid cond reset
+
+            x0_audio = minimax_h3_rf_v_to_x0(
+                audio_rows[audio_update],
+                mv_audio_t,
+                torch.tensor(t_a, dtype=torch.float32, device=device),
+            )
+            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_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)
+
+    return video_rows, audio_rows
+
+
+__all__ = [
+    "MINIMAX_H3_AUDIO_REF_COND_TIMESTEP",
+    "MINIMAX_H3_AUDIO_ROW_WIDTH",
+    "MINIMAX_H3_IMGVID_COND_TIMESTEP",
+    "MINIMAX_H3_VIDEO_ROW_WIDTH",
+    "MiniMaxH3DenoiseBranch",
+    "minimax_h3_denoise_loop",
+]

vllm_omni/diffusion/models/minimax_h3/packed_sequence.py

@@ -0,0 +1,572 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 packed-sequence materialization for FL2VA and Ref2VA tasks.
+
+Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P].
+Builder rules:
+- placeholder ids; block-derived pos infos/masks/cu_seqlens/document_id
+- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter
+  continues text_len with temporal interp spans (frame_rescale 5/3 x
+  frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced
+  coordinates excluding the right endpoint, then scales them by INTERP;
+  audio channel-major blocks pinned to the w-grid extremes.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+import torch
+
+MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = (
+    (0,),
+    (-1,),
+    (0, -1),
+)
+
+MINIMAX_H3_TEXT_ID = -5
+MINIMAX_H3_IMGVID_COND_ID = -11
+MINIMAX_H3_AUDIO_REF_COND_ID = -17
+MINIMAX_H3_AUDIO_FIRST_ID = -15
+MINIMAX_H3_AUDIO_ID = -14
+MINIMAX_H3_VIDEO_FIRST_ID = -3
+MINIMAX_H3_VIDEO_ID = -2
+MINIMAX_H3_VIDEO_LAST_ID = -4
+MINIMAX_H3_PAD_ID = -1
+
+_INTERP = 32
+_T_GROUP = 5
+_FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
+_FRAME_RESCALE = 5.0 / 3.0
+_SEQ_ALIGN = 64
+# MiniMax H3 packs video latents with a fixed [1, 2, 2] (t, h, w) patch.
+_PATCH_T = 1
+_PATCH_H = 2
+_PATCH_W = 2
+
+
+def _keyframe_cond_frame_indices(
+    *,
+    include_keyframe_cond: bool,
+    keyframe_frame_indices: list[int] | tuple[int, ...] | None,
+) -> list[int]:
+    if not include_keyframe_cond:
+        if keyframe_frame_indices is not None:
+            raise ValueError("keyframe_frame_indices must be omitted when keyframe cond is not included")
+        return []
+    if keyframe_frame_indices is None:
+        raise ValueError("strict fl2va packed layout requires keyframe_frame_indices")
+    if any(isinstance(value, bool) or not isinstance(value, int) for value in keyframe_frame_indices):
+        raise ValueError("strict fl2va packed layout requires integer keyframe_frame_indices")
+    out = list(keyframe_frame_indices)
+    if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+        raise ValueError(
+            "strict fl2va packed layout requires keyframe_frame_indices in "
+            f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}"
+        )
+    return out
+
+
+def _resolve_keyframe_frame_indices(
+    frame_indices: Sequence[int],
+    *,
+    frame_count: int | None,
+) -> list[int]:
+    if frame_indices and frame_count is None:
+        raise ValueError("frame_count is required when keyframe_frame_indices are provided")
+    if frame_count is None:
+        return []
+    if frame_count <= 0:
+        raise ValueError("frame_count must be positive")
+    seen: dict[int, int] = {}
+    resolved: list[int] = []
+    for block_index, semantic_index in enumerate(frame_indices):
+        if semantic_index == -1:
+            resolved_index = frame_count - 1
+        elif 0 <= semantic_index < frame_count:
+            resolved_index = semantic_index
+        else:
+            raise ValueError(f"keyframe frame index {semantic_index} must be -1 or in [0, {frame_count})")
+        previous = seen.get(resolved_index)
+        if previous is not None:
+            raise ValueError(
+                f"keyframe frame index at block {block_index} resolves to "
+                f"{resolved_index}, already bound by block {previous}"
+            )
+        seen[resolved_index] = block_index
+        resolved.append(resolved_index)
+    return resolved
+
+
+def _temporal_position_span(temporal_length: int) -> float:
+    """Temporal position span for patch_t=1, in fp64.
+
+    NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums
+    via numpy (pairwise summation), matching the fl2va anchor computation,
+    while ``_video_t_span`` sums sequentially, matching the ref2va
+    t-origin accumulation. The two orders diverge in the last ulp
+    from n=16 onward, so each path must keep its own summation order.
+    """
+    spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE
+    for token_index in range(_T_GROUP):
+        spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index]
+    return float(spans.sum())
+
+
+def minimax_h3_packed_sequence(
+    *,
+    text_len: int,
+    latent_t: int,
+    latent_h: int,
+    latent_w: int,
+    audio_t: int,
+    audio_channel: int = 2,
+    include_keyframe_cond: bool,
+    keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
+    frame_count: int | None = None,
+) -> dict[str, torch.Tensor]:
+    """Build the packed-sequence structural fields for one CFG branch.
+
+    The used length is padded up to a multiple of 64.
+    """
+    ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
+    frame_rows = ph * pw
+    cond_frame_indices = _keyframe_cond_frame_indices(
+        include_keyframe_cond=include_keyframe_cond,
+        keyframe_frame_indices=keyframe_frame_indices,
+    )
+    resolved_cond_frame_indices = _resolve_keyframe_frame_indices(
+        cond_frame_indices,
+        frame_count=frame_count,
+    )
+    cond_rows = len(cond_frame_indices) * frame_rows
+    video_rows = latent_t * frame_rows
+    audio_rows = audio_t * audio_channel
+    used = text_len + cond_rows + audio_rows + video_rows
+    seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN
+
+    text_sl = slice(0, text_len)
+    cond_sl = slice(text_len, text_len + cond_rows)
+    audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows)
+    video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
+    pad_sl = slice(video_sl.stop, seq_len)
+
+    input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64)
+    input_ids[text_sl] = MINIMAX_H3_TEXT_ID
+    if cond_rows:
+        input_ids[cond_sl] = MINIMAX_H3_IMGVID_COND_ID
+    input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID
+    input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID
+    input_ids[video_sl] = MINIMAX_H3_VIDEO_ID
+    input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID
+    input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID
+
+    image_mask = torch.zeros(seq_len, dtype=torch.bool)
+    image_mask[cond_sl] = True
+    image_mask[video_sl] = True
+    audio_mask = torch.zeros(seq_len, dtype=torch.bool)
+    audio_mask[audio_sl] = True
+
+    img_pos = torch.cat(
+        [
+            torch.arange(cond_sl.start, cond_sl.stop),
+            torch.arange(video_sl.start, video_sl.stop),
+        ]
+    )
+    update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
+    update_mask[cond_rows:] = True
+    audio_pos = torch.arange(audio_sl.start, audio_sl.stop)
+    text_pos = torch.arange(0, text_len)
+
+    g = torch.zeros(seq_len, 3, dtype=torch.float64)
+    g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
+
+    t_grid = _video_t_grid(latent_t, float(text_len))
+    sqrt_area = np.sqrt(latent_h * latent_w)
+    h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area)
+    w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area)
+    hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
+    frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
+    video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64)
+    video_g[:, :, 0] = t_grid[:, None]
+    video_g[:, :, 1:] = frame[None]
+    for block_index, pixel_index in enumerate(resolved_cond_frame_indices):
+        sl = slice(
+            cond_sl.start + block_index * frame_rows,
+            cond_sl.start + (block_index + 1) * frame_rows,
+        )
+        if pixel_index == 0:
+            cond_t = float(text_len)
+        elif frame_count is not None and pixel_index == frame_count - 1:
+            cond_t = float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE
+        else:
+            raise ValueError(
+                f"fl2va packed layout only supports first/last keyframe anchors, got resolved frame index {pixel_index}"
+            )
+        g[sl, 0] = cond_t
+        g[sl, 1:] = frame
+    g[video_sl] = video_g.reshape(-1, 3)
+    audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64)
+    g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
+    g[audio_sl, 2] = torch.cat(
+        [
+            torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64),
+            torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64),
+        ]
+    )
+
+    token_tags = torch.full((seq_len,), -1, dtype=torch.long)  # PADDING
+    token_tags[text_sl] = 1  # TEXT (fl2va image-segment override happens upstream)
+    token_tags[audio_sl] = 2  # AUDIO
+    token_tags[img_pos] = 0  # VIDEO
+
+    cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
+    doc = torch.zeros(seq_len, dtype=torch.int32)
+    doc[pad_sl] = 1
+    return {
+        "seq_len": torch.tensor(seq_len),
+        "input_ids": input_ids,
+        "image_mask": image_mask,
+        "audio_mask": audio_mask,
+        "img_pos": img_pos,
+        "audio_pos": audio_pos,
+        "text_pos": text_pos,
+        "update_mask": update_mask,
+        "img_position_ids": g,
+        "token_tags": token_tags,
+        "cu_seqlens": cu,
+        "document_id": doc,
+    }
+
+
+def _positive_int(
+    block: Mapping[str, object],
+    key: str,
+    path: str,
+    *,
+    allow_zero: bool = False,
+) -> int:
+    value = block.get(key)
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ValueError(f"{path}.{key} must be an integer")
+    if value < 0 or (value == 0 and not allow_zero):
+        predicate = "non-negative" if allow_zero else "positive"
+        raise ValueError(f"{path}.{key} must be {predicate}")
+    return int(value)
+
+
+def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor:
+    ratio = dim / sqrt_area
+    left = (1.0 - ratio) * 1.0 / 2.0
+    right = left + ratio * 1.0
+    grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP
+    return torch.from_numpy(grid).to(torch.float64)
+
+
+def _video_t_grid(n: int, origin: float) -> torch.Tensor:
+    spans = torch.tensor(
+        [_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)],
+        dtype=torch.float64,
+    )
+    return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
+
+
+def _video_t_span(n: int) -> float:
+    # Sequential fp64 summation on purpose — see _temporal_position_span for
+    # why the two span implementations must not be unified.
+    return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n))
+
+
+def _range_for_slice(sl: slice) -> torch.Tensor:
+    return torch.arange(sl.start, sl.stop, dtype=torch.long)
+
+
+def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor:
+    if parts:
+        return torch.cat(parts)
+    return torch.empty(0, dtype=torch.long)
+
+
+def minimax_h3_packed_sequence_ref2va_blocks(
+    *,
+    text_len: int,
+    latent_t: int,
+    latent_h: int,
+    latent_w: int,
+    audio_t: int,
+    ref_blocks: Sequence[Mapping[str, object]],
+    audio_channel: int = 2,
+    seq_len: int | None = None,
+) -> dict[str, torch.Tensor]:
+    """General ref2va-family packed layout.
+
+    ``ref_blocks`` are consumed in request/plan order:
+    - ``{"kind": "image", "latent_h": H, "latent_w": W}``
+    - ``{"kind": "audio", "ref_audio_t": T}``
+    - ``{"kind": "video"|"video_audio", "ref_audio_t": T,
+       "latent_t": RT, "latent_h": RH, "latent_w": RW}``
+
+    Video-bearing blocks pack their audio rows immediately before their video
+    rows; both share the same temporal origin and advance by the longer of the
+    audio and video spans. Standalone audio advances the target origin by its
+    own T, and image blocks advance it by one integer slot.
+    """
+    if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)):
+        raise ValueError("ref_blocks must be a sequence")
+
+    parsed: list[dict[str, object]] = []
+    ref_visual_rows = 0
+    ref_audio_rows = 0
+    for index, raw in enumerate(ref_blocks):
+        path = f"ref_blocks[{index}]"
+        if not isinstance(raw, Mapping):
+            raise ValueError(f"{path} must be an object")
+        kind = raw.get("kind", raw.get("type"))
+        if not isinstance(kind, str) or not kind:
+            raise ValueError(f"{path}.kind must be a non-empty string")
+        if kind == "image":
+            rh = _positive_int(raw, "latent_h", path)
+            rw = _positive_int(raw, "latent_w", path)
+            rows = (rh // _PATCH_H) * (rw // _PATCH_W)
+            item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows}
+            ref_visual_rows += rows
+        elif kind == "audio":
+            rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
+            rows = rt * audio_channel
+            item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows}
+            ref_audio_rows += rows
+        elif kind in ("video", "video_audio"):
+            rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
+            vt = _positive_int(raw, "latent_t", path)
+            vh = _positive_int(raw, "latent_h", path)
+            vw = _positive_int(raw, "latent_w", path)
+            frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W)
+            audio_rows = rt * audio_channel
+            video_rows = vt * frame_rows
+            item = {
+                "kind": kind,
+                "ref_audio_t": rt,
+                "latent_t": vt,
+                "latent_h": vh,
+                "latent_w": vw,
+                "frame_rows": frame_rows,
+                "audio_rows": audio_rows,
+                "video_rows": video_rows,
+            }
+            ref_audio_rows += audio_rows
+            ref_visual_rows += video_rows
+        else:
+            raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}")
+        parsed.append(item)
+
+    ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
+    frame_rows = ph * pw
+    video_rows = latent_t * frame_rows
+    audio_rows = audio_t * audio_channel
+    ref_rows = ref_visual_rows + ref_audio_rows
+    used = text_len + ref_rows + audio_rows + video_rows
+    if seq_len is None:
+        seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN
+    if seq_len < used:
+        raise ValueError(f"seq_len {seq_len} < used rows {used}")
+
+    text_sl = slice(0, text_len)
+    cursor = text_len
+    block_slices: list[dict[str, object]] = []
+    for item in parsed:
+        kind = str(item["kind"])
+        if kind == "image":
+            rows = int(item["rows"])
+            visual_sl = slice(cursor, cursor + rows)
+            cursor = visual_sl.stop
+            block_slices.append({**item, "visual_sl": visual_sl})
+        elif kind == "audio":
+            rows = int(item["audio_rows"])
+            audio_sl = slice(cursor, cursor + rows)
+            cursor = audio_sl.stop
+            block_slices.append({**item, "audio_sl": audio_sl})
+        else:
+            a_rows = int(item["audio_rows"])
+            v_rows = int(item["video_rows"])
+            audio_sl = slice(cursor, cursor + a_rows)
+            visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows)
+            cursor = visual_sl.stop
+            block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl})
+
+    audio_sl = slice(cursor, cursor + audio_rows)
+    video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
+    pad_sl = slice(video_sl.stop, seq_len)
+
+    input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64)
+    input_ids[text_sl] = MINIMAX_H3_TEXT_ID
+    image_mask = torch.zeros(seq_len, dtype=torch.bool)
+    audio_mask = torch.zeros(seq_len, dtype=torch.bool)
+
+    ref_img_pos_parts: list[torch.Tensor] = []
+    ref_audio_pos_parts: list[torch.Tensor] = []
+    g = torch.zeros(seq_len, 3, dtype=torch.float64)
+    g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
+
+    target_area = np.sqrt(latent_h * latent_w)
+    h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area)
+    w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area)
+    hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
+    target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
+
+    t_cursor = float(text_len)
+    for item in block_slices:
+        kind = str(item["kind"])
+        if kind == "image":
+            visual_sl = item["visual_sl"]
+            assert isinstance(visual_sl, slice)
+            input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID
+            image_mask[visual_sl] = True
+            ref_img_pos_parts.append(_range_for_slice(visual_sl))
+            rh = int(item["latent_h"])
+            rw = int(item["latent_w"])
+            area = np.sqrt(rh * rw)
+            ref_hh, ref_ww = torch.meshgrid(
+                _axis_from_sqrt_area(rh, _PATCH_H, area),
+                _axis_from_sqrt_area(rw, _PATCH_W, area),
+                indexing="ij",
+            )
+            g[visual_sl, 0] = t_cursor
+            g[visual_sl, 1] = ref_hh.reshape(-1)
+            g[visual_sl, 2] = ref_ww.reshape(-1)
+            t_cursor += 1.0
+        elif kind == "audio":
+            audio_ref_sl = item["audio_sl"]
+            assert isinstance(audio_ref_sl, slice)
+            ref_t = int(item["ref_audio_t"])
+            input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID
+            audio_mask[audio_ref_sl] = True
+            ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
+            ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
+            g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
+            if ref_t:
+                g[audio_ref_sl, 2] = torch.cat(
+                    [
+                        torch.full((ref_t,), float(w_grid[0]), dtype=torch.float64),
+                        torch.full(
+                            (int(item["audio_rows"]) - ref_t,),
+                            float(w_grid[-1]),
+                            dtype=torch.float64,
+                        ),
+                    ]
+                )
+            t_cursor += float(ref_t)
+        else:
+            audio_ref_sl = item["audio_sl"]
+            visual_sl = item["visual_sl"]
+            assert isinstance(audio_ref_sl, slice)
+            assert isinstance(visual_sl, slice)
+            ref_t = int(item["ref_audio_t"])
+            vt = int(item["latent_t"])
+            vh = int(item["latent_h"])
+            vw = int(item["latent_w"])
+            input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID
+            input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID
+            audio_mask[audio_ref_sl] = True
+            image_mask[visual_sl] = True
+            ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
+            ref_img_pos_parts.append(_range_for_slice(visual_sl))
+
+            ref_area = np.sqrt(vh * vw)
+            rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area)
+            rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area)
+            rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij")
+
+            ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
+            g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
+            if ref_t:
+                g[audio_ref_sl, 2] = torch.cat(
+                    [
+                        torch.full((ref_t,), float(rv_w_grid[0]), dtype=torch.float64),
+                        torch.full(
+                            (int(item["audio_rows"]) - ref_t,),
+                            float(rv_w_grid[-1]),
+                            dtype=torch.float64,
+                        ),
+                    ]
+                )
+
+            rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1)
+            rv_g = torch.empty(vt, int(item["frame_rows"]), 3, dtype=torch.float64)
+            rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None]
+            rv_g[:, :, 1:] = rv_frame[None]
+            g[visual_sl] = rv_g.reshape(-1, 3)
+            t_cursor += max(float(ref_t), _video_t_span(vt))
+
+    input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID
+    input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID
+    input_ids[video_sl] = MINIMAX_H3_VIDEO_ID
+    input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID
+    input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID
+    audio_mask[audio_sl] = True
+    image_mask[video_sl] = True
+
+    audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64)
+    g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
+    g[audio_sl, 2] = torch.cat(
+        [
+            torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64),
+            torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64),
+        ]
+    )
+
+    video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64)
+    video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None]
+    video_g[:, :, 1:] = target_frame[None]
+    g[video_sl] = video_g.reshape(-1, 3)
+
+    target_img_pos = _range_for_slice(video_sl)
+    target_audio_pos = _range_for_slice(audio_sl)
+    img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos])
+    audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos])
+
+    update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
+    update_mask[ref_visual_rows:] = True
+    audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool)
+    audio_update_mask[ref_audio_rows:] = True
+    text_pos = torch.arange(0, text_len)
+
+    token_tags = torch.full((seq_len,), -1, dtype=torch.long)  # PADDING
+    token_tags[text_sl] = 1  # TEXT
+    if ref_audio_rows:
+        token_tags[_cat_ranges(ref_audio_pos_parts)] = 2  # AUDIO (reference)
+    token_tags[audio_sl] = 2  # AUDIO (target)
+    token_tags[img_pos] = 0  # VIDEO (refs + target)
+
+    cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
+    doc = torch.zeros(seq_len, dtype=torch.int32)
+    doc[pad_sl] = 1
+    return {
+        "seq_len": torch.tensor(seq_len),
+        "input_ids": input_ids,
+        "image_mask": image_mask,
+        "audio_mask": audio_mask,
+        "img_pos": img_pos,
+        "audio_pos": audio_pos,
+        "text_pos": text_pos,
+        "update_mask": update_mask,
+        "audio_update_mask": audio_update_mask,
+        "img_position_ids": g,
+        "token_tags": token_tags,
+        "cu_seqlens": cu,
+        "document_id": doc,
+    }
+
+
+__all__ = [
+    "MINIMAX_H3_AUDIO_FIRST_ID",
+    "MINIMAX_H3_AUDIO_ID",
+    "MINIMAX_H3_AUDIO_REF_COND_ID",
+    "MINIMAX_H3_IMGVID_COND_ID",
+    "MINIMAX_H3_PAD_ID",
+    "MINIMAX_H3_TEXT_ID",
+    "MINIMAX_H3_VIDEO_FIRST_ID",
+    "MINIMAX_H3_VIDEO_ID",
+    "MINIMAX_H3_VIDEO_LAST_ID",
+    "minimax_h3_packed_sequence",
+    "minimax_h3_packed_sequence_ref2va_blocks",
+]

vllm_omni/diffusion/models/minimax_h3/packed_tokens.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def _int_tuple(value: Sequence[int], name: str, length: int) -> tuple[int, ...]:
+    if len(value) != length:
+        raise ValueError(f"{name} must have length {length}, got {list(value)!r}")
+    out = tuple(int(item) for item in value)
+    if any(item <= 0 for item in out):
+        raise ValueError(f"{name} values must be positive, got {list(value)!r}")
+    return out
+
+
+def _rank(tensor: torch.Tensor, name: str, rank: int) -> None:
+    if tensor.ndim != rank:
+        raise ValueError(f"{name} must be rank {rank}, got shape={list(tensor.shape)}")
+
+
+def minimax_h3_patchify_video_latent(
+    latent: torch.Tensor,
+    *,
+    patch_size: Sequence[int],
+) -> torch.Tensor:
+    """Pack video latent [B,C,T,H,W] into DiT token rows."""
+
+    _rank(latent, "video latent", 5)
+    pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
+    batch, channel, full_t, full_h, full_w = (int(dim) for dim in latent.shape)
+    if full_t % pt or full_h % ph or full_w % pw:
+        raise ValueError(
+            "video latent spatial/time dims must be divisible by patch_size: "
+            f"shape={list(latent.shape)}, patch_size={[pt, ph, pw]}"
+        )
+    t, h, w = full_t // pt, full_h // ph, full_w // pw
+    packed = latent.reshape(batch, channel, t, pt, h, ph, w, pw)
+    packed = torch.einsum("nctrhpwq->nthwcrpq", packed)
+    return packed.reshape(batch * t * h * w, channel * pt * ph * pw).contiguous()
+
+
+def minimax_h3_unpatchify_video_tokens(
+    rows: torch.Tensor,
+    *,
+    latent_shape: Sequence[int],
+    patch_size: Sequence[int],
+) -> torch.Tensor:
+    """Unpack DiT video token rows into latent [B,C,T,H,W]."""
+
+    _rank(rows, "video token rows", 2)
+    t, h, w, channel = _int_tuple(latent_shape, "latent_shape", 4)
+    pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
+    expected_dim = pt * ph * pw * channel
+    if int(rows.shape[-1]) != expected_dim:
+        raise ValueError(
+            f"video token dim {int(rows.shape[-1])} != patch volume * channel "
+            f"{expected_dim} for latent_shape={list(latent_shape)}, "
+            f"patch_size={[pt, ph, pw]}"
+        )
+    rows_per_sample = t * h * w
+    if int(rows.shape[0]) % rows_per_sample:
+        raise ValueError(
+            f"video rows {int(rows.shape[0])} must be divisible by t*h*w "
+            f"{rows_per_sample} for latent_shape={list(latent_shape)}"
+        )
+    packed = rows.reshape(-1, t, h, w, channel, pt, ph, pw)
+    latent = torch.einsum("nthwcrpq->nctrhpwq", packed)
+    return latent.reshape(-1, channel, t * pt, h * ph, w * pw).contiguous()
+
+
+def minimax_h3_pack_audio_latent(latent: torch.Tensor) -> torch.Tensor:
+    """Pack audio VAE latent [audio_channel, latent_dim, T] into rows."""
+
+    _rank(latent, "audio latent", 3)
+    audio_channel, latent_dim, steps = (int(dim) for dim in latent.shape)
+    return (
+        latent.permute(0, 2, 1)
+        .reshape(
+            audio_channel * steps,
+            latent_dim,
+        )
+        .contiguous()
+    )
+
+
+def minimax_h3_unpack_audio_tokens(
+    rows: torch.Tensor,
+    *,
+    audio_t: int,
+    audio_channel: int,
+) -> torch.Tensor:
+    """Unpack DiT audio token rows into audio VAE latent [C,latent_dim,T]."""
+
+    _rank(rows, "audio token rows", 2)
+    audio_t = int(audio_t)
+    audio_channel = int(audio_channel)
+    if audio_t <= 0 or audio_channel <= 0:
+        raise ValueError(f"audio_t and audio_channel must be positive, got {audio_t=} {audio_channel=}")
+    if int(rows.shape[0]) != audio_t:
+        raise ValueError(f"audio rows {int(rows.shape[0])} != audio_t {audio_t}")
+    if audio_t % audio_channel:
+        raise ValueError(f"audio_t must be divisible by audio_channel, got {audio_t=} {audio_channel=}")
+    native = rows.reshape(audio_channel, audio_t // audio_channel, int(rows.shape[-1]))
+    return native.permute(0, 2, 1).contiguous()
+
+
+__all__ = [
+    "minimax_h3_pack_audio_latent",
+    "minimax_h3_patchify_video_latent",
+    "minimax_h3_unpack_audio_tokens",
+    "minimax_h3_unpatchify_video_tokens",
+]

vllm_omni/diffusion/models/minimax_h3/presentation.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 Qwen presentation building.
+
+Builds the positive presentation token stream:
+- fl2va: '<Picture 1>: ' label + vision block (<|vision_start|> +
+  N*<|image_pad|> + <|vision_end|>) + prompt text.
+- t2va: prompt text only (no vision block).
+Prompt text passes through verbatim (no stripping or rewriting).
+
+All presentation variants are emitted through the shared ``_Presentation``
+accumulator so ids and AdaLN token tags cannot drift apart.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+IMAGE_PAD = "<|image_pad|>"
+VIDEO_PAD = "<|video_pad|>"
+
+_TEXT_TAG = 1
+_VIDEO_TAG = 0
+
+
+def _text_ids(tokenizer: Any, text: str) -> list[int]:
+    return list(tokenizer(text, add_special_tokens=False)["input_ids"])
+
+
+def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]:
+    return (
+        [tokenizer.convert_tokens_to_ids(VISION_START)]
+        + [tokenizer.convert_tokens_to_ids(pad_token)] * int(count)
+        + [tokenizer.convert_tokens_to_ids(VISION_END)]
+    )
+
+
+class _Presentation:
+    """Accumulates aligned (ids, token_tags) presentation segments."""
+
+    def __init__(self) -> None:
+        self.ids: list[int] = []
+        self.tags: list[int] = []
+
+    def text(self, token_ids: list[int]) -> None:
+        self.ids += token_ids
+        self.tags += [_TEXT_TAG] * len(token_ids)
+
+    def vision(self, token_ids: list[int]) -> None:
+        self.ids += token_ids
+        self.tags += [_VIDEO_TAG] * len(token_ids)
+
+    def build(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return (
+            torch.tensor(self.ids, dtype=torch.long),
+            torch.tensor(self.tags, dtype=torch.long),
+        )
+
+
+def _timestamped_video_blocks(
+    presentation: _Presentation,
+    tokenizer: Any,
+    *,
+    counts: Sequence[int],
+    timestamps: Sequence[float],
+    context: str,
+) -> None:
+    """Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision."""
+
+    counts = [int(value) for value in counts]
+    timestamps = [float(value) for value in timestamps]
+    if not counts or len(counts) != len(timestamps):
+        raise ValueError(f"{context}video block token counts and timestamps must align")
+    for count, timestamp in zip(counts, timestamps):
+        if count <= 0:
+            raise ValueError(f"{context}video block token count must be positive")
+        presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>"))
+        presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count))
+
+
+def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor:
+    """t2va presentation: verbatim prompt, no special tokens."""
+    if not prompt:
+        raise ValueError("prompt must be non-empty")
+    return torch.tensor(_text_ids(tokenizer, prompt), dtype=torch.long)
+
+
+def _multi_image_presentation(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    image_token_counts: list[int],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not image_token_counts:
+        raise ValueError("image_token_counts must be non-empty")
+    presentation = _Presentation()
+    for index, count in enumerate(image_token_counts, start=1):
+        if int(count) <= 0:
+            raise ValueError("image_token_count must be positive")
+        presentation.text(_text_ids(tokenizer, f"<Picture {index}>: "))
+        presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+    presentation.text(_text_ids(tokenizer, prompt))
+    return presentation.build()
+
+
+def minimax_h3_multi_image_presentation_ids(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    image_token_counts: list[int],
+) -> torch.Tensor:
+    """Positive fl2va presentation with one or more labeled image slots."""
+    ids, _ = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts)
+    return ids
+
+
+def minimax_h3_multi_image_presentation_token_tags(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    image_token_counts: list[int],
+) -> torch.Tensor:
+    """AdaLN token tags aligned with minimax_h3_multi_image_presentation_ids."""
+    _, tags = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts)
+    return tags
+
+
+def minimax_h3_ref2va_presentation(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    condition_labels: list[tuple[str, int]],
+    image_token_count: int | list[int] | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """ref2va positive presentation:
+
+    per condition in request order — image i: ``<Picture i>: `` label followed
+    by the vision block; audio j: ``<Audio j>: `` label only (audio content
+    never enters Qwen) — then the verbatim prompt. Returns (ids, token_tags)
+    with the vision block tagged VIDEO(0) and everything else TEXT(1).
+
+    condition_labels: [("image", 1), ("audio", 1), ...] with 1-based ordinals
+    per type.
+    """
+    return minimax_h3_ref2va_video_presentation(
+        tokenizer,
+        prompt=prompt,
+        condition_labels=condition_labels,
+        image_token_count=image_token_count,
+        video_block_token_counts=None,
+        video_block_timestamps=None,
+    )
+
+
+def _as_int_list(value: int | Sequence[int] | None, *, name: str) -> list[int]:
+    if value is None:
+        return []
+    if isinstance(value, int):
+        return [int(value)]
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise ValueError(f"{name} must be an int or a sequence of ints")
+    return [int(item) for item in value]
+
+
+def _as_nested_int_list(
+    value: Sequence[int] | Sequence[Sequence[int]] | None,
+    *,
+    name: str,
+) -> list[list[int]]:
+    if value is None:
+        return []
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise ValueError(f"{name} must be a sequence")
+    if len(value) == 0:
+        return []
+    first = value[0]
+    if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+        out: list[list[int]] = []
+        for group in value:
+            if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
+                raise ValueError(f"{name} must not mix nested and flat entries")
+            out.append([int(item) for item in group])
+        return out
+    return [[int(item) for item in value]]
+
+
+def _as_nested_float_list(
+    value: Sequence[float] | Sequence[Sequence[float]] | None,
+    *,
+    name: str,
+) -> list[list[float]]:
+    if value is None:
+        return []
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise ValueError(f"{name} must be a sequence")
+    if len(value) == 0:
+        return []
+    first = value[0]
+    if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+        out: list[list[float]] = []
+        for group in value:
+            if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
+                raise ValueError(f"{name} must not mix nested and flat entries")
+            out.append([float(item) for item in group])
+        return out
+    return [[float(item) for item in value]]
+
+
+def minimax_h3_ref2va_video_presentation(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    condition_labels: list[tuple[str, int]],
+    image_token_count: int | list[int] | None,
+    video_block_token_counts: list[int] | list[list[int]] | None,
+    video_block_timestamps: list[float] | list[list[float]] | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """ref2va (optionally with video refs) positive presentation:
+
+    per condition in request order —
+    - image i:  ``<Picture i>: `` label + one image vision block;
+    - audio j:  ``<Audio j>: `` label only (audio content never enters Qwen);
+    - video k:  ``<Video k>: `` label, then per temporal block a timestamp
+      text ``<{t:.1f} seconds>`` followed by a VIDEO vision block
+      (<|vision_start|> + <|video_pad|> x n + <|vision_end|>). Timestamps are
+      the mean of each merged frame pair (Qwen3VL temporal merge 2; odd frame
+      counts repeat the last frame), emitting the
+      ``<0.2 seconds>`` ..
+      ``<4.0 seconds>`` sequence — note Python bankers-rounding at .1f.
+    then the verbatim prompt. Vision blocks are tagged VIDEO(0), everything
+    else TEXT(1).
+    """
+    if not prompt:
+        raise ValueError("prompt must be non-empty")
+    presentation = _Presentation()
+    image_token_counts = _as_int_list(image_token_count, name="image_token_count")
+    video_counts_by_ref = _as_nested_int_list(
+        video_block_token_counts,
+        name="video_block_token_counts",
+    )
+    video_timestamps_by_ref = _as_nested_float_list(
+        video_block_timestamps,
+        name="video_block_timestamps",
+    )
+    if len(video_counts_by_ref) != len(video_timestamps_by_ref):
+        raise ValueError("video block token counts and timestamps must align")
+    image_seen = 0
+    video_seen = 0
+    for cond_type, ordinal in condition_labels:
+        if cond_type == "image":
+            image_seen += 1
+            if image_seen > len(image_token_counts):
+                raise ValueError("image_token_count required for an image reference")
+            count = int(image_token_counts[image_seen - 1])
+            if count <= 0:
+                raise ValueError("image_token_count required for an image reference")
+            presentation.text(_text_ids(tokenizer, f"<Picture {ordinal}>: "))
+            presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+        elif cond_type == "audio":
+            presentation.text(_text_ids(tokenizer, f"<Audio {ordinal}>: "))
+        elif cond_type == "video":
+            video_seen += 1
+            if video_seen > len(video_counts_by_ref):
+                raise ValueError("video reference requires block token counts and timestamps")
+            counts = video_counts_by_ref[video_seen - 1]
+            timestamps = video_timestamps_by_ref[video_seen - 1]
+            if not counts or not timestamps:
+                raise ValueError("video reference requires block token counts and timestamps")
+            presentation.text(_text_ids(tokenizer, f"<Video {ordinal}>: "))
+            _timestamped_video_blocks(
+                presentation,
+                tokenizer,
+                counts=counts,
+                timestamps=timestamps,
+                context="",
+            )
+        else:
+            raise ValueError(f"unsupported ref2va condition type {cond_type!r}")
+    if image_seen != len(image_token_counts):
+        raise ValueError("unused image_token_count entries")
+    if video_seen != len(video_counts_by_ref):
+        raise ValueError("unused video block token count entries")
+    presentation.text(_text_ids(tokenizer, prompt))
+    return presentation.build()
+
+
+__all__ = [
+    "minimax_h3_multi_image_presentation_ids",
+    "minimax_h3_multi_image_presentation_token_tags",
+    "minimax_h3_ref2va_presentation",
+    "minimax_h3_ref2va_video_presentation",
+    "minimax_h3_text_only_ids",
+]

vllm_omni/diffusion/models/minimax_h3/reference_video.py

@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 Ref2VA reference-video preparation."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import subprocess
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+MINIMAX_H3_FPS = 24.0
+MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0
+MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2
+MINIMAX_H3_BASE_SHORT_EDGE = 768
+MINIMAX_H3_MAX_PIXELS = 768 * 1344
+MINIMAX_H3_CANVAS_MULTIPLE = 32
+
+
+def _nearest_multiple(value: float, multiple: int) -> int:
+    return max(multiple, int(round(float(value) / multiple)) * multiple)
+
+
+def _probe_video(path: str) -> dict[str, Any]:
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-count_frames",
+            "-show_entries",
+            ("stream=width,height,r_frame_rate,nb_read_frames,nb_frames,sample_aspect_ratio:stream_tags=rotate"),
+            "-of",
+            "json",
+            path,
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    streams = json.loads(result.stdout).get("streams") or []
+    if not streams:
+        raise ValueError(f"media has no video stream: {path}")
+    stream = streams[0]
+    numerator, denominator = str(stream["r_frame_rate"]).split("/")
+    raw_count = stream.get("nb_read_frames") or stream.get("nb_frames")
+    if raw_count in (None, "", "N/A"):
+        raise ValueError(f"cannot determine video frame count: {path}")
+    return {
+        "width": int(stream["width"]),
+        "height": int(stream["height"]),
+        "fps": float(numerator) / float(denominator),
+        "frame_count": int(raw_count),
+    }
+
+
+def _has_audio(path: str) -> bool:
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "a:0",
+            "-show_entries",
+            "stream=index",
+            "-of",
+            "json",
+            path,
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return bool(json.loads(result.stdout).get("streams") or [])
+
+
+def _reference_video_shape(width: int, height: int) -> tuple[int, int]:
+    ratio = float(width) / float(height)
+    if not 0.25 <= ratio <= 4.0:
+        raise ValueError(f"reference video aspect ratio must be in [1:4, 4:1], got {width}x{height}")
+    if ratio >= 1.0:
+        target_width = MINIMAX_H3_BASE_SHORT_EDGE * ratio
+        target_height = float(MINIMAX_H3_BASE_SHORT_EDGE)
+    else:
+        target_width = float(MINIMAX_H3_BASE_SHORT_EDGE)
+        target_height = MINIMAX_H3_BASE_SHORT_EDGE / ratio
+    area = target_width * target_height
+    if area > MINIMAX_H3_MAX_PIXELS:
+        scale = math.sqrt(MINIMAX_H3_MAX_PIXELS / area)
+        target_width *= scale
+        target_height *= scale
+    return (
+        _nearest_multiple(target_width, MINIMAX_H3_CANVAS_MULTIPLE),
+        _nearest_multiple(target_height, MINIMAX_H3_CANVAS_MULTIPLE),
+    )
+
+
+def _transcode_reference_video(
+    source: str,
+    *,
+    target_width: int,
+    target_height: int,
+    target_frame_count: int,
+    workdir: str,
+) -> str:
+    output = str(Path(workdir) / "prepared.mp4")
+    subprocess.run(
+        [
+            "ffmpeg",
+            "-y",
+            "-loglevel",
+            "error",
+            "-i",
+            source,
+            "-map",
+            "0:v:0",
+            "-an",
+            "-vf",
+            (f"fps={MINIMAX_H3_FPS:g},scale={target_width}:{target_height}:flags=lanczos,setsar=1"),
+            "-frames:v",
+            str(target_frame_count),
+            "-metadata:s:v:0",
+            "rotate=0",
+            "-c:v",
+            "libx264",
+            "-pix_fmt",
+            "yuv420p",
+            output,
+        ],
+        check=True,
+    )
+    return output
+
+
+def prepare_reference_videos(
+    values: Any,
+    *,
+    target_frame_count: int,
+    workdir: str,
+) -> list[dict[str, Any]]:
+    if isinstance(values, (str, os.PathLike)):
+        values = [values]
+    if not isinstance(values, (list, tuple)) or not values:
+        raise ValueError("MiniMax H3 Ref2VA video input must be a path or a non-empty list of paths")
+    prepared: list[dict[str, Any]] = []
+    for index, value in enumerate(values):
+        if not isinstance(value, (str, os.PathLike)):
+            raise TypeError(
+                f"MiniMax H3 multi-video Ref2VA currently requires file paths, got item {index}: {type(value)!r}"
+            )
+        source = str(value)
+        meta = _probe_video(source)
+        width, height = _reference_video_shape(meta["width"], meta["height"])
+        item_workdir = Path(workdir) / f"video_{index}"
+        item_workdir.mkdir(parents=True)
+        prepared_path = _transcode_reference_video(
+            source,
+            target_width=width,
+            target_height=height,
+            target_frame_count=target_frame_count,
+            workdir=str(item_workdir),
+        )
+        prepared.append(
+            {
+                "original_path": source,
+                "prepared_path": prepared_path,
+                "input_has_audio": _has_audio(source),
+                "width": width,
+                "height": height,
+            }
+        )
+    return prepared
+
+
+def load_video_frames(path: str) -> np.ndarray:
+    try:
+        import decord
+    except ImportError:
+        import av
+
+        with av.open(path) as container:
+            frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
+        if not frames:
+            raise ValueError(f"video has no frames: {path}")
+        return np.stack(frames)
+
+    reader = decord.VideoReader(path)
+    if len(reader) <= 0:
+        raise ValueError(f"video has no frames: {path}")
+    frames = reader.get_batch(list(range(len(reader))))
+    return frames.asnumpy() if hasattr(frames, "asnumpy") else np.asarray(frames)
+
+
+def sample_reference_video_frames(
+    prepared_path: str,
+    *,
+    workdir: str,
+) -> dict[str, Any]:
+    from PIL import Image
+
+    meta = _probe_video(prepared_path)
+    ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS
+    indices: list[int] = []
+    cursor = 0.0
+    while True:
+        frame_index = int(round(cursor))
+        if frame_index >= meta["frame_count"]:
+            break
+        if not indices or frame_index > indices[-1]:
+            indices.append(frame_index)
+        cursor += ratio
+    if not indices:
+        raise ValueError(f"no frames sampled from {prepared_path}")
+
+    frame_dir = Path(workdir)
+    frame_dir.mkdir(parents=True, exist_ok=True)
+    frames = []
+    for output_index, source_index in enumerate(indices, start=1):
+        output = frame_dir / f"frame_{output_index:06d}.png"
+        subprocess.run(
+            [
+                "ffmpeg",
+                "-y",
+                "-loglevel",
+                "error",
+                "-i",
+                prepared_path,
+                "-vf",
+                f"select=eq(n\\,{source_index})",
+                "-vsync",
+                "vfr",
+                "-frames:v",
+                "1",
+                str(output),
+            ],
+            check=True,
+        )
+        frames.append(np.asarray(Image.open(output).convert("RGB")))
+
+    timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))]
+    timestamps += [timestamps[-1]] * ((-len(timestamps)) % MINIMAX_H3_QWEN_TEMPORAL_PATCH)
+    block_timestamps = [
+        (timestamps[index] + timestamps[index + MINIMAX_H3_QWEN_TEMPORAL_PATCH - 1]) / 2
+        for index in range(
+            0,
+            len(timestamps),
+            MINIMAX_H3_QWEN_TEMPORAL_PATCH,
+        )
+    ]
+    return {
+        "frames": frames,
+        "block_timestamps": block_timestamps,
+    }
+
+
+def _soundfile_to_waveform(path: str) -> tuple[torch.Tensor, int]:
+    import soundfile as sf
+
+    data, sample_rate = sf.read(path, dtype="float32", always_2d=True)
+    # soundfile is (samples, channels); torchaudio.load is (channels, samples).
+    waveform = torch.from_numpy(data).t().contiguous()
+    return waveform, int(sample_rate)
+
+
+def load_audio_file(path: str) -> tuple[torch.Tensor, int]:
+    """Load an audio file as ``(waveform[C, T] float32, sample_rate)``.
+
+    torchaudio 2.6+ routes ``load`` through TorchCodec, whose wheels do not load
+    on aarch64 with a CPU-only torch (they link CUDA libraries that are absent).
+    Fall back to soundfile; if the container is not libsndfile-readable (e.g.
+    mp3/m4a/mp4), demux to wav with ffmpeg first so any ffmpeg-supported input
+    works at its native sample rate.
+    """
+    try:
+        import torchaudio
+
+        return torchaudio.load(path)
+    except (ImportError, RuntimeError, OSError):
+        pass
+    try:
+        return _soundfile_to_waveform(path)
+    except Exception:
+        import tempfile
+
+        with tempfile.TemporaryDirectory(prefix="minimax_h3_audio_") as tmpdir:
+            wav = str(Path(tmpdir) / "audio.wav")
+            subprocess.run(
+                [
+                    "ffmpeg",
+                    "-y",
+                    "-loglevel",
+                    "error",
+                    "-i",
+                    path,
+                    "-vn",
+                    "-f",
+                    "wav",
+                    wav,
+                ],
+                check=True,
+            )
+            return _soundfile_to_waveform(wav)
+
+
+def load_video_audio(path: str) -> tuple[torch.Tensor, int]:
+    import tempfile
+
+    with tempfile.TemporaryDirectory(prefix="minimax_h3_video_audio_") as tmpdir:
+        output = str(Path(tmpdir) / "audio.wav")
+        subprocess.run(
+            [
+                "ffmpeg",
+                "-y",
+                "-loglevel",
+                "error",
+                "-i",
+                path,
+                "-vn",
+                "-ac",
+                "2",
+                "-ar",
+                "44100",
+                "-f",
+                "wav",
+                output,
+            ],
+            check=True,
+        )
+        return load_audio_file(output)
+
+
+__all__ = [
+    "load_audio_file",
+    "load_video_audio",
+    "load_video_frames",
+    "prepare_reference_videos",
+    "sample_reference_video_frames",
+]

vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+
+
+def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None:
+    if not bool(torch.isfinite(tensor).all().item()):
+        raise ValueError(f"{name} must be finite")
+
+
+def _validate_unit_timestep(timestep: torch.Tensor, name: str) -> None:
+    if not isinstance(timestep, torch.Tensor):
+        raise ValueError(f"{name} must be a torch.Tensor")
+    if not torch.is_floating_point(timestep):
+        raise ValueError(f"{name} must be a floating point tensor")
+    _require_finite_tensor(timestep, name)
+    out_of_range = (timestep < 0) | (timestep > 1)
+    if bool(out_of_range.any().item()):
+        raise ValueError(f"{name} must be in [0, 1]")
+
+
+def _validate_sigma(value: float, name: str) -> float:
+    sigma = float(value)
+    if not math.isfinite(sigma):
+        raise ValueError(f"{name} must be finite")
+    if sigma < 0.0:
+        raise ValueError(f"{name} must be non-negative")
+    return sigma
+
+
+def _validate_timestep_sigma_pair(
+    timestep: torch.Tensor,
+    sigma_curr: float,
+    name: str,
+) -> float:
+    _validate_unit_timestep(timestep, f"{name}_timestep")
+    sigma = _validate_sigma(sigma_curr, f"{name}_sigma_curr")
+    expected = 1.0 - timestep.detach().to(dtype=torch.float32)
+    actual = torch.full_like(expected, sigma)
+    if not torch.allclose(actual, expected, rtol=1e-5, atol=1e-5):
+        raise ValueError(f"{name}_sigma_curr must equal 1 - {name}_timestep")
+    return sigma
+
+
+def minimax_h3_rf_v_to_x0(
+    xt: torch.Tensor,
+    v: torch.Tensor,
+    timestep: torch.Tensor,
+) -> torch.Tensor:
+    if xt.shape != v.shape:
+        raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}")
+    if not torch.is_floating_point(xt):
+        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_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")
+    return x0
+
+
+def minimax_h3_euler_eta0_step(
+    state: torch.Tensor,
+    denoised: torch.Tensor,
+    *,
+    sigma_curr: float,
+    sigma_next: float,
+) -> torch.Tensor:
+    if state.shape != denoised.shape:
+        raise ValueError(f"state and denoised shapes must match, got {state.shape} vs {denoised.shape}")
+    if not torch.is_floating_point(state):
+        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")
+    sigma_curr = _validate_sigma(sigma_curr, "sigma_curr")
+    sigma_next = _validate_sigma(sigma_next, "sigma_next")
+    if sigma_curr == 0.0:
+        if sigma_next != 0.0:
+            raise ValueError("sigma_next must be 0 when sigma_curr is 0")
+        return state
+    compute_dtype = torch.float32
+    if state.dtype not in (torch.float16, torch.bfloat16):
+        compute_dtype = state.dtype
+    sigma_curr_t = state.new_tensor(sigma_curr, dtype=compute_dtype)
+    sigma_next_t = state.new_tensor(sigma_next, dtype=compute_dtype)
+    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")
+    return out
+
+
+class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
+    def __init__(self, **config: Any) -> None:
+        if config:
+            raise ValueError(f"{type(self).__name__} does not accept config fields: {sorted(config)}")
+
+    def set_shift(self, _flow_shift: float) -> None:
+        """Ignore flow shift, matching the previous loader-specific path."""
+
+    def step_denoising(
+        self,
+        *,
+        input_visual_latent: torch.Tensor,
+        input_audio_latent: torch.Tensor,
+        timestep: torch.Tensor,
+        noise_pred_visual: torch.Tensor,
+        noise_pred_audio: torch.Tensor,
+        sigma_curr: float,
+        sigma_next: float,
+        video_timestep: torch.Tensor | None = None,
+        audio_timestep: torch.Tensor | None = None,
+        video_sigma_curr: float | None = None,
+        video_sigma_next: float | None = None,
+        audio_sigma_curr: float | None = None,
+        audio_sigma_next: float | None = None,
+    ) -> dict[str, torch.Tensor]:
+        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
+        visual_sigma_next = sigma_next if video_sigma_next is None else video_sigma_next
+        audio_sigma_curr = sigma_curr if audio_sigma_curr is None else audio_sigma_curr
+        audio_sigma_next = sigma_next if audio_sigma_next is None else audio_sigma_next
+        visual_sigma_curr = _validate_timestep_sigma_pair(
+            visual_timestep,
+            visual_sigma_curr,
+            "video",
+        )
+        audio_sigma_curr = _validate_timestep_sigma_pair(
+            audio_timestep,
+            audio_sigma_curr,
+            "audio",
+        )
+
+        denoised_visual = minimax_h3_rf_v_to_x0(
+            input_visual_latent,
+            noise_pred_visual,
+            visual_timestep,
+        )
+        denoised_audio = minimax_h3_rf_v_to_x0(
+            input_audio_latent,
+            noise_pred_audio,
+            audio_timestep,
+        )
+        return {
+            "output_visual_latent": minimax_h3_euler_eta0_step(
+                input_visual_latent,
+                denoised_visual,
+                sigma_curr=visual_sigma_curr,
+                sigma_next=visual_sigma_next,
+            ),
+            "output_audio_latent": minimax_h3_euler_eta0_step(
+                input_audio_latent,
+                denoised_audio,
+                sigma_curr=audio_sigma_curr,
+                sigma_next=audio_sigma_next,
+            ),
+        }
+
+
+EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter
+
+__all__ = [
+    "MiniMaxH3EulerAncestralEta0SchedulerAdapter",
+    "minimax_h3_euler_eta0_step",
+    "minimax_h3_rf_v_to_x0",
+]

vllm_omni/diffusion/models/minimax_h3/time_request.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+
+def _align_frame_count(frame_count: int) -> int:
+    """Snap ``frame_count`` up to the MiniMax H3 17n+5 frame boundary."""
+    if frame_count <= 0:
+        return 1
+    current = int(frame_count)
+    while current % 17 != 5:
+        current += 1
+    return current
+
+
+def _video_latent_t(frame_count: int) -> int:
+    if frame_count <= 5:
+        return 2
+    return ((int(frame_count) - 5) // 17) * 5 + 2
+
+
+def _frame_count_from_video_latent_t(out_t: int) -> int:
+    if out_t == 1:
+        return 1
+    if out_t < 2 or (out_t - 2) % 5 != 0:
+        raise ValueError("MiniMax H3 video latent T must be 1 or match 5n+2")
+    return 17 * ((int(out_t) - 2) // 5) + 5
+
+
+def _audio_latent_t(duration_seconds: float) -> int:
+    # Rounding happens at the 40 Hz audio latent boundary.
+    return int(round(float(duration_seconds) * 40.0))
+
+
+def _time_shift_sigmas(
+    *,
+    num_steps: int = 50,
+    shift_scale: float = 6.0,
+) -> list[float]:
+    if shift_scale <= 0:
+        raise ValueError("MiniMax H3 shift_scale must be > 0")
+    if num_steps <= 0:
+        raise ValueError("MiniMax H3 num_steps must be > 0")
+
+    import torch
+
+    # The rectified-flow sigma range is fixed at [1.0, 0.0].
+    base = torch.linspace(
+        1.0,
+        0.0,
+        int(num_steps),
+        device="cpu",
+        dtype=torch.float32,
+    )
+    shifted = float(shift_scale) * base / (1 + (float(shift_scale) - 1) * base)
+    shifted, _ = torch.unique_consecutive(shifted, return_counts=True)
+    # A one-point request is still exactly one point.  Normal serving uses
+    # multiple points, but preserving the requested cardinality keeps
+    # ``num_inference_steps`` the sole schedule-size control.
+    if num_steps > 1 and shifted[-1].item() > 0.0:
+        shifted = torch.cat([shifted, torch.tensor([0.0], dtype=shifted.dtype)])
+    return [float(value) for value in shifted.tolist()]
+
+
+class MiniMaxH3ShapePlanner:
+    """Compute MiniMax H3 frame, latent-shape, and time-shift values."""
+
+    def align_frame_count(self, frame_count: int) -> int:
+        return _align_frame_count(frame_count)
+
+    def video_latent_t(self, frame_count: int) -> int:
+        return _video_latent_t(frame_count)
+
+    def frame_count_from_video_latent_t(self, out_t: int) -> int:
+        return _frame_count_from_video_latent_t(out_t)
+
+    def audio_latent_t(self, duration_seconds: float) -> int:
+        return _audio_latent_t(duration_seconds)
+
+    def time_shift_sigmas(
+        self,
+        *,
+        num_steps: int = 50,
+        shift_scale: float = 6.0,
+    ) -> list[float]:
+        return _time_shift_sigmas(
+            num_steps=num_steps,
+            shift_scale=shift_scale,
+        )
+
+
+MINIMAX_H3_SHAPE_PLANNER = MiniMaxH3ShapePlanner()
+
+
+def minimax_h3_align_frame_count(frame_count: int) -> int:
+    return MINIMAX_H3_SHAPE_PLANNER.align_frame_count(frame_count)
+
+
+def minimax_h3_frame_count_from_video_latent_t(out_t: int) -> int:
+    return MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(out_t)
+
+
+def minimax_h3_time_shift_sigmas(
+    *,
+    num_steps: int = 50,
+    shift_scale: float = 6.0,
+) -> list[float]:
+    return MINIMAX_H3_SHAPE_PLANNER.time_shift_sigmas(
+        num_steps=num_steps,
+        shift_scale=shift_scale,
+    )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -0,0 +1,360 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 remote-code VAE adapters and exact latent contracts."""
+
+from __future__ import annotations
+
+import importlib
+import json
+from contextlib import AbstractContextManager
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from PIL import Image
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedVaeMixin,
+)
+from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+
+from .packed_tokens import minimax_h3_patchify_video_latent
+
+MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
+MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
+MINIMAX_H3_AUDIO_CHANNELS = 2
+
+
+def _load_component_config(component_path: str) -> dict[str, Any]:
+    config_path = Path(component_path) / "config.json"
+    config = json.loads(config_path.read_text(encoding="utf-8"))
+    channels = int(config["latent_channels"])
+    for key in ("latents_mean", "latents_std"):
+        values = config.get(key)
+        if not isinstance(values, list) or len(values) != channels:
+            raise ValueError(f"{config_path}: {key} must contain {channels} values")
+    return config
+
+
+def _load_remote_component(
+    component_path: str,
+    config: dict[str, Any],
+) -> nn.Module:
+    auto_map = config.get("auto_map") or {}
+    class_reference = auto_map.get("AutoModel")
+    if not isinstance(class_reference, str):
+        raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel")
+    component_cls = get_class_from_dynamic_module(
+        class_reference,
+        component_path,
+    )
+    return component_cls.from_pretrained(component_path)
+
+
+class _AudioVAEDeterminismContext(AbstractContextManager):
+    def __enter__(self):
+        backends = torch.backends
+        self._saved = (
+            backends.cuda.matmul.allow_tf32,
+            backends.cudnn.allow_tf32,
+            backends.cudnn.benchmark,
+            backends.cudnn.deterministic,
+            backends.cudnn.enabled,
+            backends.cuda.flash_sdp_enabled(),
+            backends.cuda.mem_efficient_sdp_enabled(),
+            backends.cuda.math_sdp_enabled(),
+        )
+        backends.cuda.matmul.allow_tf32 = False
+        backends.cudnn.allow_tf32 = False
+        backends.cudnn.benchmark = False
+        backends.cudnn.deterministic = True
+        backends.cudnn.enabled = False
+        backends.cuda.enable_flash_sdp(False)
+        backends.cuda.enable_mem_efficient_sdp(False)
+        backends.cuda.enable_math_sdp(True)
+        return self
+
+    def __exit__(self, exc_type, exc, traceback):
+        backends = torch.backends
+        (
+            backends.cuda.matmul.allow_tf32,
+            backends.cudnn.allow_tf32,
+            backends.cudnn.benchmark,
+            backends.cudnn.deterministic,
+            backends.cudnn.enabled,
+            flash,
+            memory_efficient,
+            math_sdp,
+        ) = self._saved
+        backends.cuda.enable_flash_sdp(flash)
+        backends.cuda.enable_mem_efficient_sdp(memory_efficient)
+        backends.cuda.enable_math_sdp(math_sdp)
+        return False
+
+
+class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
+    """Adapter around the checkpoint's native parallel-tiled video VAE."""
+
+    def __init__(
+        self,
+        component_path: str,
+        *,
+        device: torch.device,
+    ) -> None:
+        super().__init__()
+        self.config_dict = _load_component_config(component_path)
+        self.remote = _load_remote_component(
+            component_path,
+            self.config_dict,
+        )
+        # Match the reference loader contract: video VAE weights stay FP32.
+        # Keyframe encoding is numerically sensitive to first casting the
+        # checkpoint through FP16; decode still runs under FP16 autocast.
+        self.remote.eval().to(device=device, dtype=torch.float32)
+        self.model = self.remote.model
+        self.use_tiling = True
+        self.use_slicing = False
+        self.parallel_size = 1
+
+    def set_parallel_size(
+        self,
+        parallel_size: int,
+        mode: str = "tile",
+    ) -> None:
+        if mode != "tile":
+            raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
+        group = get_dit_group()
+        world_size = dist.get_world_size(group)
+        rank = dist.get_rank(group)
+        parallel_size = int(parallel_size)
+        if parallel_size not in (1, world_size):
+            raise ValueError(
+                "MiniMax H3 native VAE patch parallelism currently requires "
+                "vae_patch_parallel_size=1 or the full DiT group size "
+                f"({world_size}), got {parallel_size}"
+            )
+        self.parallel_size = parallel_size
+        enabled = parallel_size > 1
+
+        package = self.remote.__class__.__module__.rsplit(".", 1)[0]
+        parallel_module = importlib.import_module(f"{package}.parallel")
+        state = parallel_module.get_parallel_state()
+        state.clear()
+        state.update(
+            group_size=parallel_size,
+            group_rank=rank if enabled else 0,
+            local_process_group=group if enabled else None,
+            sp_size=parallel_size,
+            sp_rank=rank if enabled else 0,
+            sp_enabled=enabled,
+            sp_process_group=group if enabled else None,
+            tp_size=1,
+            tp_rank=0,
+        )
+        self.model.parallel_tiling = enabled
+
+    def is_distributed_enabled(self) -> bool:
+        return self.parallel_size > 1 and dist.is_initialized()
+
+    @torch.inference_mode()
+    def encode_image(self, image: Image.Image) -> torch.Tensor:
+        previous_parallel = self.model.parallel_tiling
+        self.model.parallel_tiling = False
+        parameter = next(self.parameters())
+        previous_dtype = parameter.dtype
+        if previous_dtype != torch.float32:
+            self.to(torch.float32)
+        devices = [parameter.device] if parameter.device.type == "cuda" else []
+        try:
+            with torch.random.fork_rng(devices=devices):
+                torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
+                for device in devices:
+                    with torch.cuda.device(device):
+                        torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
+                latent = self.model.encode_images(
+                    image,
+                    use_fp16_latent=True,
+                )[0]
+        finally:
+            self.model.parallel_tiling = previous_parallel
+            if previous_dtype != torch.float32:
+                self.to(previous_dtype)
+
+        # Match the reference contract exactly: normalization and patchify
+        # happen on CPU in FP32 after the sampled encode. The condition noise
+        # path is sensitive enough that doing these elementwise operations on
+        # CUDA can noticeably change the final conditioned video.
+        latent = latent.float().cpu()
+        if latent.ndim == 4:
+            latent = latent[None]
+        channels = int(self.config_dict["latent_channels"])
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+        ).view(1, channels, 1, 1, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+        ).view(1, channels, 1, 1, 1)
+        return minimax_h3_patchify_video_latent(
+            (latent - mean) / std,
+            patch_size=(1, 2, 2),
+        ).float()
+
+    @torch.inference_mode()
+    def encode_video(
+        self,
+        frames: Any,
+    ) -> tuple[torch.Tensor, tuple[int, int, int]]:
+        parameter = next(self.parameters())
+        previous_dtype = parameter.dtype
+        if previous_dtype != torch.float32:
+            self.to(torch.float32)
+        devices = [parameter.device] if parameter.device.type == "cuda" else []
+        try:
+            with torch.random.fork_rng(devices=devices):
+                torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
+                for device in devices:
+                    with torch.cuda.device(device):
+                        torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
+                latent = self.model.encode_videos(
+                    frames,
+                    use_fp16_latent=True,
+                )[0]
+        finally:
+            if previous_dtype != torch.float32:
+                self.to(previous_dtype)
+
+        latent = latent.float().cpu()
+        if latent.ndim == 4:
+            latent = latent[None]
+        channels = int(self.config_dict["latent_channels"])
+        if latent.ndim != 5 or int(latent.shape[1]) != channels:
+            raise ValueError(f"unexpected reference video latent shape {tuple(latent.shape)}")
+        shape = (
+            int(latent.shape[2]),
+            int(latent.shape[3]),
+            int(latent.shape[4]),
+        )
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+        ).view(1, channels, 1, 1, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+        ).view(1, channels, 1, 1, 1)
+        rows = minimax_h3_patchify_video_latent(
+            (latent - mean) / std,
+            patch_size=(1, 2, 2),
+        ).float()
+        return rows, shape
+
+    @torch.inference_mode()
+    def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        channels = int(self.config_dict["latent_channels"])
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        decoded = self.model.decode_base(latent * std + mean)
+        frames = self.model.processor.revert_tensor(decoded)
+        if frames.ndim == 4:
+            frames = frames.unsqueeze(0).transpose(1, 2)
+        if frames.ndim != 5:
+            raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
+        return frames.float()
+
+
+class MiniMaxH3AudioVAE(nn.Module):
+    def __init__(
+        self,
+        component_path: str,
+        *,
+        device: torch.device,
+    ) -> None:
+        super().__init__()
+        self.config_dict = _load_component_config(component_path)
+        self.remote = _load_remote_component(
+            component_path,
+            self.config_dict,
+        )
+        # The checkpoint's audio VAE contract is FP32 for both reference
+        # encoding and waveform decoding.
+        self.remote.eval().to(device=device, dtype=torch.float32)
+        self.model = self.remote.model
+        self.sample_rate = int(self.config_dict["sample_rate"])
+
+    @torch.inference_mode()
+    def encode_waveform(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+    ) -> tuple[torch.Tensor, int]:
+        import torchaudio
+
+        waveform = waveform.float()
+        if waveform.ndim == 1:
+            waveform = waveform[None]
+        if int(sample_rate) != MINIMAX_H3_AUDIO_SAMPLE_RATE:
+            waveform = torchaudio.transforms.Resample(
+                int(sample_rate),
+                MINIMAX_H3_AUDIO_SAMPLE_RATE,
+            )(waveform)
+        if waveform.shape[0] < MINIMAX_H3_AUDIO_CHANNELS:
+            waveform = waveform.repeat(
+                MINIMAX_H3_AUDIO_CHANNELS,
+                1,
+            )
+        waveform = waveform[:MINIMAX_H3_AUDIO_CHANNELS]
+        device = next(self.model.parameters()).device
+        waveform = waveform.to(device)
+
+        with _AudioVAEDeterminismContext():
+            audio = self.model.preprocess(
+                waveform.unsqueeze(1),
+                MINIMAX_H3_AUDIO_SAMPLE_RATE,
+            )
+            latent = self.model.encoder(audio)
+            if bool(getattr(self.model, "attn_proj", False)):
+                latent = self.model.pre_block(latent.transpose(1, 2)).transpose(1, 2)
+            latent = self.model.mean_proj(latent).float().cpu()
+
+        channels = int(self.config_dict["latent_channels"])
+        if latent.shape[-1] != channels:
+            if latent.shape[1] != channels:
+                raise ValueError(f"cannot canonicalize audio latent {tuple(latent.shape)}")
+            latent = latent.transpose(1, 2).contiguous()
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+        ).view(1, 1, channels)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+        ).view(1, 1, channels)
+        rows = ((latent - mean) / std).reshape(-1, channels)
+        return rows.float(), int(latent.shape[1])
+
+    @torch.inference_mode()
+    def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        channels = int(self.config_dict["latent_channels"])
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1)
+        waveform = self.remote.decode(latent * std + mean)
+        if waveform.ndim != 3 or waveform.shape[1] != 1:
+            raise ValueError(f"unexpected decoded audio shape {tuple(waveform.shape)}")
+        return waveform.permute(1, 0, 2).contiguous().float()
+
+
+__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]

vllm_omni/diffusion/registry.py

@@ -86,6 +86,11 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2",
         "LTX2I2VDMD2Pipeline",
     ),
+    "MiniMaxH3Pipeline": (
+        "minimax_h3",
+        "pipeline_minimax_h3",
+        "MiniMaxH3Pipeline",
+    ),
     "StableAudioPipeline": (
         "stable_audio",
         "pipeline_stable_audio",
@@ -502,6 +507,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LTX2DistilledPipeline": "get_ltx2_post_process_func",
     "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
     "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
+    "MiniMaxH3Pipeline": "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/distributed/omni_connectors/kv_transfer_manager.py

@@ -69,6 +69,12 @@ class _TransferTopoConfig:
     sp_rank: int
     sp_group: Any
     world: Any
+    # The stage's tensor-model-parallel group (vLLM parallel state). Only the
+    # ranks in this group execute the same request in lockstep — DP replicas
+    # and PP stages in ``world`` do not — so per-request consensus exchanges
+    # must use this group, never ``world``.
+    tp_group: Any = None
+    tp_size: int = 1
 
     @property
     def is_leader(self) -> bool:
@@ -642,6 +648,16 @@ class OmniKVTransferManager:
         else:
             role = ReceiveRole.LEADER if world_rank == 0 else ReceiveRole.FOLLOWER
 
+        # The stage's TP group comes from vLLM parallel state (the TP-sharded
+        # LM path); it may legitimately be uninitialized on non-TP stages.
+        try:
+            from vllm.distributed.parallel_state import get_tp_group
+
+            tp_group = get_tp_group()
+            tp_size = tp_group.world_size
+        except Exception:
+            tp_group, tp_size = None, 1  # vLLM TP not initialized
+
         return _TransferTopoConfig(
             role=role,
             tp_active=tp_active,
@@ -652,6 +668,8 @@ class OmniKVTransferManager:
             sp_rank=sp_rank,
             sp_group=sp_group,
             world=world,
+            tp_group=tp_group,
+            tp_size=tp_size,
         )
 
     def _resolve_sender_info(
@@ -764,6 +782,20 @@ class OmniKVTransferManager:
 
         return kv_payload
 
+    @staticmethod
+    def _clear_request_kv_payload(req: Any) -> None:
+        """Drop every request-side KV field a receive may have applied.
+
+        Enumerates fields via ``_collect_request_kv_payload`` so clearing stays
+        symmetric with what apply/broadcast set — the pipeline then takes the
+        same no-KV path as a rank that never received.
+        """
+        for key in OmniKVTransferManager._collect_request_kv_payload(req):
+            if key.startswith("sp."):
+                setattr(req.sampling_params, key[3:], None)
+            else:
+                setattr(req, key, None)
+
     @staticmethod
     def _apply_request_kv_payload(
         req: Any,
@@ -1753,6 +1785,64 @@ class OmniKVTransferManager:
         combined = combined.to(device)
         return self._unpack_kv_payload(combined)
 
+    def _tp_local_receive_consensus(self, req: Any, received: bool) -> bool:
+        """Identical-KV-state verdict across pure-TP ranks (LOCAL role).
+
+        In pure TP every rank fetches its own KV shard with an independent
+        timeout and ``distribute_kv_cache`` is a no-op, so a single rank's
+        silent miss sends only that rank into the pipeline's no-KV prefill —
+        it then enters TP collectives its peers never join, and the NCCL
+        watchdog kills the worker (#5627). Each rank reports a signature of
+        its applied KV state — the primary-receive flag plus the set of
+        applied payload roles, so partially-received CFG companion KV
+        (``cfg_text``/``cfg_img``) diverging between ranks is caught too. If
+        the signatures differ, every rank drops its KV so all take the same
+        fallback path together; identical-but-partial state is kept, since
+        symmetric absence is consistent (e.g. Bagel's empty unconditional
+        branch).
+
+        The exchange runs over the stage's tensor-parallel group only: those
+        are the ranks executing this request in lockstep. ``world`` would
+        also contain DP replicas / PP stages working on different requests —
+        mixing their flags corrupts the verdict, and a replica skipping the
+        exchange (dummy request) while another enters it would deadlock.
+        Within the TP group the skip guards below depend only on topology,
+        config, and the request id — identical on every TP rank — so ranks
+        always agree on whether the exchange happens.
+        """
+        pt = self.topo_config
+        group = pt.tp_group
+        if not (pt.is_local and pt.tp_active) or group is None or pt.tp_size <= 1:
+            return received
+        if not self.config.need_recv_cache:
+            return received
+        request_id = self._resolve_request_id(req)
+        if request_id is None or OmniDiffusionRequest.is_dummy_run_request_id(request_id):
+            return received
+
+        signature = (bool(received), tuple(sorted(self._collect_request_kv_payload(req))))
+        if group.rank_in_group == 0:
+            signatures = [signature]
+            signatures.extend(group.recv_object(src=src) for src in range(1, pt.tp_size))
+            consistent = all(sig == signature for sig in signatures)
+            group.broadcast_object(consistent, src=0)
+        else:
+            group.send_object(signature, dst=0)
+            consistent = bool(group.broadcast_object(None, src=0))
+
+        if not consistent:
+            logger.warning(
+                "KV state for %s diverged across TP ranks (this rank: received=%s, roles=%s); "
+                "dropping KV on every rank so all take the no-KV fallback together "
+                "instead of desyncing TP collectives",
+                request_id,
+                signature[0],
+                list(signature[1]),
+            )
+            self._clear_request_kv_payload(req)
+            return False
+        return received
+
     def receive_multi_kv_cache_distributed(
         self,
         req: Any,
@@ -1763,6 +1853,7 @@ class OmniKVTransferManager:
         received = False
         if not self.topo_config.is_follower:
             received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
+        received = self._tp_local_receive_consensus(req, received)
         kv_payload = self.distribute_kv_cache(req, target_device, received=received)
         if kv_payload is not None:
             self._apply_request_kv_payload(req, kv_payload, target_device)
@@ -1791,6 +1882,7 @@ class OmniKVTransferManager:
         if not received and not self.topo_config.is_follower and not payload_consumed:
             logger.debug("KV prefetch miss for %s; falling back to sync receive", self._resolve_request_id(req))
             received = self.receive_multi_kv_cache(req, None, target_device)
+        received = self._tp_local_receive_consensus(req, received)
         kv_payload = self.distribute_kv_cache(req, target_device, received=received)
         if kv_payload is not None:
             self._apply_request_kv_payload(req, kv_payload, target_device)

vllm_omni/engine/arg_utils.py

@@ -490,6 +490,7 @@ class OrchestratorArgs:
     cfg_parallel_size: int = 1
     vae_patch_parallel_size: int = 1
     vae_parallel_mode: str = "tile"
+    text_encoder_tp_size: int = 1
     default_sampling_params: str | None = None
     max_generated_image_size: int | None = None
     tts_max_instructions_length: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -213,6 +213,12 @@ class AsyncOmniEngine:
         # beforehand. The stage CLI exposes the same deploy YAML through
         # stage_configs_path.
         deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
+        # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
+        # specified" so stage-config resolution can defer to the deploy yaml's
+        # per-stage value (see ``with_trust_remote_code_override``). The
+        # restriction path below loads the top-level HF config via vLLM's
+        # ``get_config``, which needs a real bool, so collapse ``None`` to the
+        # default ``False`` here (#5495).
         pipeline_config = StageConfigFactory.get_pipeline_config(
             model=model,
             trust_remote_code=bool(trust_remote_code),
@@ -959,6 +965,7 @@ class AsyncOmniEngine:
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
             vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
             vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
+            text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
             enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", False)
             hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
@@ -979,6 +986,7 @@ class AsyncOmniEngine:
                 cfg_parallel_size=cfg_parallel_size,
                 vae_patch_parallel_size=vae_patch_parallel_size,
                 vae_parallel_mode=vae_parallel_mode,
+                text_encoder_tp_size=text_encoder_tp_size,
                 use_hsdp=use_hsdp,
                 hsdp_shard_size=hsdp_shard_size,
                 hsdp_replicate_size=hsdp_replicate_size,

vllm_omni/entrypoints/cli/serve.py

@@ -713,6 +713,15 @@ class OmniServeCommand(CLISubcommand):
             "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. "
             "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.",
         )
+        omni_config_group.add_argument(
+            "--text-encoder-tp-size",
+            type=int,
+            default=1,
+            help="Tensor-parallel degree for the diffusion text encoder. "
+            "Shards the Qwen3-VL encoder across the first N DiT ranks, "
+            "removing the rank-0 encoder memory hotspot in no-offload runs. "
+            "Equivalent to setting DiffusionParallelConfig.text_encoder_tp_size.",
+        )
         omni_config_group.add_argument(
             "--vae-parallel-mode",
             type=str,

vllm_omni/entrypoints/openai/api_server.py

@@ -11,6 +11,7 @@ import os
 
 # Image generation API imports
 import random
+import tempfile
 import time
 from argparse import Namespace
 from collections.abc import AsyncIterator
@@ -2856,6 +2857,18 @@ async def _cleanup_video(video_id: str):
         logger.warning("Failed to cleanup partial video file '%s'", video_id)
 
 
+def _cleanup_video_references(
+    reference_video: ReferenceVideo | None,
+    reference_audio: ReferenceAudio | None,
+) -> None:
+    if reference_video is not None:
+        for path in reference_video.cleanup_paths:
+            if os.path.exists(path):
+                os.unlink(path)
+    if reference_audio is not None and os.path.exists(reference_audio.path):
+        os.unlink(reference_audio.path)
+
+
 async def _run_video_generation_job(
     handler: OmniOpenAIServingVideo,
     request: VideoGenerationRequest,
@@ -2936,17 +2949,37 @@ async def _run_video_generation_job(
         await VIDEO_STORE.pop(video_id)
         raise
     finally:
-        if reference_audio is not None and os.path.exists(reference_audio.path):
-            os.unlink(reference_audio.path)
+        _cleanup_video_references(reference_video, reference_audio)
 
 
 VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0))
 
 
+async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[str]:
+    paths: list[str] = []
+    try:
+        for upload in uploads:
+            suffix = Path(upload.filename or "").suffix.lower()
+            if suffix not in {".mkv", ".mov", ".mp4", ".webm"}:
+                suffix = ".mp4"
+            fd, path = tempfile.mkstemp(prefix="vllm_omni_video_reference_", suffix=suffix)
+            paths.append(path)
+            with os.fdopen(fd, "wb") as output:
+                while chunk := await upload.read(1024 * 1024):
+                    output.write(chunk)
+    except Exception:
+        for path in paths:
+            if os.path.exists(path):
+                os.unlink(path)
+        raise
+    return paths
+
+
 async def _parse_video_form(
     raw_request: Request,
     prompt: str = Form(...),
     input_reference: UploadFile | None = File(default=None),
+    input_references: list[UploadFile] | None = File(default=None),
     image_reference: str | None = Form(default=None),
     video_reference: str | None = Form(default=None),
     audio_reference: str | None = Form(default=None),
@@ -2987,11 +3020,19 @@ async def _parse_video_form(
 
     Used by both ``POST /v1/videos`` (async) and ``POST /v1/videos/sync``.
     """
+    input_references = input_references or []
     input_reference_bytes = await input_reference.read() if input_reference is not None else None
     parsed_image_reference = _parse_form_json(image_reference)
     parsed_video_reference = _parse_form_json(video_reference)
     parsed_audio_reference = _parse_form_json(audio_reference)
 
+    if input_references and any(
+        item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes)
+    ):
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="Provide input_references alone, without input_reference, image_reference, or video_reference.",
+        )
     provided_references = sum(
         item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes)
     )
@@ -3063,32 +3104,39 @@ async def _parse_video_form(
         )
 
     decode_spec = ReferenceVideoDecodeSpec()
-    if parsed_video_reference is not None or input_reference_bytes is not None:
+    if not input_references and (parsed_video_reference is not None or input_reference_bytes is not None):
         stage_configs = (
             handler.stage_configs
             or app_stage_configs
             or getattr(getattr(handler, "_engine_client", None), "stage_configs", None)
         )
         decode_spec = _reference_video_decode_spec(request, stage_configs)
-    try:
-        media_data = await decode_input_reference(
-            request.image_reference,
-            request.video_reference,
-            input_reference_bytes,
-            max_video_frames=decode_spec.max_frames,
-            video_keep=decode_spec.keep,
-        )
-    except InvalidInputReferenceError as exc:
-        raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
+    reference_image = None
+    reference_video = None
+    if input_references:
+        paths = await _persist_uploaded_video_references(input_references)
+        reference_video = ReferenceVideo(data=paths, cleanup_paths=tuple(paths))
+    else:
+        try:
+            media_data = await decode_input_reference(
+                request.image_reference,
+                request.video_reference,
+                input_reference_bytes,
+                max_video_frames=decode_spec.max_frames,
+                video_keep=decode_spec.keep,
+            )
+        except InvalidInputReferenceError as exc:
+            raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
 
-    reference_image = ReferenceImage(data=media_data) if isinstance(media_data, Image.Image) else None
-    reference_video = ReferenceVideo(data=media_data) if isinstance(media_data, list) else None
+        reference_image = ReferenceImage(data=media_data) if isinstance(media_data, Image.Image) else None
+        reference_video = ReferenceVideo(data=media_data) if isinstance(media_data, list) else None
 
     reference_audio: ReferenceAudio | None = None
     if request.audio_reference is not None:
         try:
             audio_path = await decode_audio_url(request.audio_reference.audio_url)
         except InvalidInputReferenceError as exc:
+            _cleanup_video_references(reference_video, None)
             raise HTTPException(400, detail=str(exc)) from exc
         reference_audio = ReferenceAudio(path=audio_path)
 
@@ -3201,8 +3249,7 @@ async def create_video_sync(
             detail=f"Video generation failed: {str(exc)}",
         ) from exc
     finally:
-        if reference_audio is not None and os.path.exists(reference_audio.path):
-            os.unlink(reference_audio.path)
+        _cleanup_video_references(reference_video, reference_audio)
     inference_time_s = time.perf_counter() - started_at
 
     return Response(

vllm_omni/entrypoints/openai/serving_video.py

@@ -49,7 +49,8 @@ class ReferenceImage:
 class ReferenceVideo:
     """Reference video frames for video-conditioned generation."""
 
-    data: list[Image.Image]
+    data: list[Image.Image] | list[str]
+    cleanup_paths: tuple[str, ...] = ()
 
 
 @dataclass