Fix README.md typo

Signed-off-by: Samit <285365963@qq.com>

.claude/skills/precheck-pr/SKILL.md

@@ -0,0 +1,97 @@
+---
+name: precheck-pr
+description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
+---
+
+# PR Pre-Check
+
+Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
+
+## Mode Selection
+
+| Mode | When | Time |
+|------|------|------|
+| **Quick** | About to push, final sanity check | ~3 min |
+| **Full** | Ready for review, want maintainer-level scan | ~10 min |
+
+Default to quick if unsure. Run full before marking a PR "ready for review."
+
+## Workflow
+
+### Step 1: Detect Base Branch
+
+```bash
+BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
+         || git merge-base HEAD main 2>/dev/null \
+         || echo origin/main)
+echo "diffing against ${BASE_SHA}"
+git diff --name-only ${BASE_SHA}...HEAD
+```
+
+### Step 2: Validate PR Title
+
+Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
+
+| Prefix | Applies to |
+|--------|-----------|
+| `[Bugfix]` | Bug fixes |
+| `[CI/Build]` | Build or CI improvements |
+| `[Doc]` | Documentation changes |
+| `[Model]` | New/improved models (include model name) |
+| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
+| `[Kernel]` | CUDA/kernel changes |
+| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
+| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
+| `[Misc]` | Other changes (use sparingly) |
+
+✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
+⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
+
+### Step 3: Categorize the PR
+
+| Diff contains | PR type |
+|---------------|---------|
+| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
+| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
+| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
+| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
+| Everything else | **General** |
+
+If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
+
+### Step 4: Run Checklist
+
+Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
+
+### Step 5: Print Report
+
+```
+Pre-check report for <branch>
+
+  Mode: quick | full
+  Type: <new-model | diffusion-model | bug-fix | perf | general>
+
+  Dimension          Result
+  ─────────────────  ──────
+  PR title format    ✓
+  PR desc integrity  ✓
+  Registry/config    ✓
+  Dead code          ⚠ 2 warnings
+  Accuracy           ✓
+  Benchmark          ✗ missing software versions
+
+  Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
+```
+
+**Severity:**
+
+| Mark | Meaning |
+|------|---------|
+| ✗ | Blocking — fix before opening PR |
+| ⚠ | Warning — consider fixing |
+| ✓ | Pass |
+| — | Skipped (not applicable) |
+
+### Stop Here
+
+Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.

.claude/skills/precheck-pr/references/checklists.md

@@ -0,0 +1,135 @@
+# Pre-Check Checklists
+
+Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
+
+---
+
+## All PRs: Title Format (check first)
+
+PR title must follow the project convention documented in `docs/contributing/README.md`:
+
+| Prefix | Applies to |
+|--------|-----------|
+| `[Bugfix]` | Bug fixes |
+| `[CI/Build]` | Build or CI improvements |
+| `[Doc]` | Documentation changes |
+| `[Model]` | New/improved models (include model name) |
+| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
+| `[Kernel]` | CUDA/kernel changes |
+| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
+| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
+| `[Misc]` | Other changes (use sparingly) |
+
+**Examples:**
+```
+✓ [Model] Add MiniCPM-o 4.5 support
+✓ [Bugfix] Fix Ovis image text encoder dtype
+✓ [Core] Qwen3-Omni performance optimization
+✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
+✗ Fix Ovis image text encoder dtype              ← missing prefix
+✗ [bugfix] Fix Ovis image text encoder           ← wrong case
+✗ [WIP] Add model                                ← WIP in title
+```
+
+- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
+
+---
+
+## Bug Fix PRs
+
+### Quick
+
+- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
+- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
+- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
+- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
+- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
+- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
+- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
+- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
+
+---
+
+## Performance PRs
+
+### Quick
+
+- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
+- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
+- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
+- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
+- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
+- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
+- [ ] **VRAM measured:** peak GPU memory before and after
+- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
+- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
+- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
+- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
+
+---
+
+## New Model PRs
+
+### Quick
+
+- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
+- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
+- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
+- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
+- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Dead code scan (5 patterns):**
+  1. Dead `forward()` — defined but never called through the pipeline path
+  2. Dead factory/builder functions — instantiated nowhere
+  3. Dead wrappers — thin wrappers with no callers
+  4. Dead branch guards — `if version > X` where X is always satisfied
+  5. Unused parameters — `__init__` args never accessed
+- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
+- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
+- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
+- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
+- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
+- [ ] **Benchmark script checked in + exact command line + pytest summary**
+
+---
+
+## Diffusion Model PRs
+
+Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
+
+- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
+- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
+- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
+- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
+- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
+- [ ] **E2E test exists:** at least one test exercises the full generate path
+
+---
+
+## General PRs
+
+### Quick
+
+- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
+
+### Full (adds)
+
+- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
+- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
+- [ ] **Import hygiene:** no redundant imports added
+- [ ] **No unrelated changes:** diff scope matches PR description

.claude/skills/readme.md

@@ -25,7 +25,12 @@ include:
 - `add-tts-model`: covers integration of new TTS models and related serving
   workflows
 - `generate-release-note`: helps prepare release notes for repository changes
+- `precheck-pr`: self-check a branch before creating a PR — validates PR title
+  format, catches dead code, verifies accuracy/perf claims, and confirms merge
+  readiness
 - `review-pr`: provides a structured workflow for reviewing pull requests
+- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
+  latest vllm-ascend NPUModelRunner
 
 ## Maintenance Guidelines
 

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,4 +1,6 @@
 <!-- markdownlint-disable -->
+**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
+
 PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
 
 ## Purpose
@@ -18,4 +20,5 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
-**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
+(anything written below this line will be removed by GitHub Actions)

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!

docs/api/README.md

@@ -92,7 +92,6 @@ Core scheduling and caching components.
 Configuration classes.
 
 - [vllm_omni.config.model.OmniModelConfig][]
-- [vllm_omni.config.stage_config.ModelPipeline][]
 - [vllm_omni.config.stage_config.StageConfig][]
 - [vllm_omni.config.stage_config.StageConfigFactory][]
 - [vllm_omni.config.stage_config.StageType][]

docs/contributing/README.md

@@ -116,6 +116,20 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
     If the PR spans more than one category, please include all relevant prefixes.
 
+### Pre-Check Before Submitting
+
+Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
+
+```
+/precheck-pr
+```
+
+The skill offers two modes:
+- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
+- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
+
+The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
+
 ### Local Test
 Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
 

docs/design/module/dit_module.md

@@ -305,8 +305,8 @@ class GPUWorker:
         )
 
         # Load model
-        model_loader = DiffusersPipelineLoader(load_config)
-        self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
+        model_loader = DiffusersPipelineLoader(load_config, self.od_config)
+        self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
 
         # Setup cache backend
         from vllm_omni.diffusion.cache.selector import get_cache_backend

docs/getting_started/installation/gpu/cuda.inc.md

@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
 uv pip install vllm==0.22.0 --torch-backend=auto
 ```
 
-The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -70,8 +70,8 @@ git checkout v0.22.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

docs/getting_started/quickstart.md

@@ -22,7 +22,7 @@ source .venv/bin/activate
 uv pip install vllm==0.22.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
+uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -33,7 +33,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -50,6 +50,8 @@ th {
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
 | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -62,6 +64,10 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
+| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
+| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
+| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -69,11 +75,13 @@ th {
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
 | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
+| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
+| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
 

docs/serving/speech_api.md

@@ -116,7 +116,7 @@ Content-Type: application/json
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
-**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
 
 #### Voice Clone Parameters (Base task)
 

docs/serving/videos_api.md

@@ -67,8 +67,9 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
-| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
+| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
+| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
+| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
 | `width` | integer | model default | Output video width |
 | `height` | integer | model default | Output video height |
 | `num_frames` | integer | model default | Number of generated frames |
@@ -80,6 +81,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
 | `flow_shift` | number | null | Scheduler flow-shift value |
 | `true_cfg_scale` | number | null | True CFG scale when supported by the model |
 | `seed` | integer | null | Random seed for reproducibility |
+| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
+| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
 | `negative_prompt` | string | null | Text describing what to avoid in the generated video |
 | `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
 | `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
@@ -123,6 +126,41 @@ curl -s http://localhost:8091/v1/videos \
   -F "fps=16"
 ```
 
+### Video-to-Video
+
+For models that support video conditioning, upload the reference video with
+`input_reference`:
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=continue this motion with consistent subjects and lighting" \
+  -F "input_reference=@input.mp4;type=video/mp4" \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16"
+```
+
+You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
+through `video_reference`. Do not send `video_reference` together with
+`input_reference` or `image_reference`.
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=continue this motion with consistent subjects and lighting" \
+  -F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16"
+```
+
+JSON references currently support `image_url`/`video_url`; `file_id` references
+are not implemented yet. Models may expose additional V2V controls through
+`extra_params`. For example, Cosmos3 supports
+`condition_frame_indexes_vision` and `condition_video_keep` to select which
+decoded reference frames are used as clean conditioning.
+
 ### Synchronous Generation
 
 ```bash
@@ -146,7 +184,10 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
 
 ## Model-Specific Examples
 
-For complete text-to-video and image-to-video walkthroughs, see:
+For complete text-to-video, image-to-video, and model-specific video-to-video
+walkthroughs, see:
 
 - [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
 - [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
+- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
+  for model-specific video-to-video examples and conditioning controls

docs/user_guide/diffusion_features.md

@@ -148,7 +148,7 @@ The following tables show which models support each feature:
 | **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/quantization/fp8.md

@@ -28,6 +28,50 @@ in deep DiT blocks.
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
 
+### Faster FP8 GEMM on Blackwell (quack)
+
+On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
+applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
+DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
+lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
+that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
+
+```bash
+# CUDA 12.9
+pip install vllm-omni[quack]
+
+# CUDA 13.x
+pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
+```
+
+It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
+on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
+Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
+not installed, FP8 still works — it just keeps the unfused FlashInfer path.
+
+#### Compile cache and warmup
+
+quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
+the first time across all autotuned configs). The compiled `.o` files are cached on
+disk and reused on later runs, so this is a one-time cost — **not** per request.
+
+vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
+`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
+In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
+the image — so the first cold start does not recompile. The engine's startup dummy
+run already exercises the kernels, so with a warm cache the first real request is fast.
+
+To pre-warm specific shapes (e.g. at image build time):
+
+```python
+from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
+# (M, K, N) per linear; M = number of tokens for your resolution/frame count
+warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
+```
+
+> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
+> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
+
 ## Model Type Support
 
 ### Diffusion Model (Qwen-Image, Wan2.2)

examples/offline_inference/cosmos3/inputs/v2v.json

@@ -0,0 +1,14 @@
+{
+    "prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
+    "negative_prompt": "",
+    "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
+    "height": 720,
+    "width": 1280,
+    "num_frames": 189,
+    "num_inference_steps": 35,
+    "guidance_scale": 6.0,
+    "flow_shift": 10.0,
+    "fps": 24,
+    "condition_frame_indexes_vision": [0, 1],
+    "condition_video_keep": "first"
+}

examples/online_serving/text_to_speech/fish_speech/gradio_demo.py

@@ -51,13 +51,18 @@ def build_payload(
     response_format: str = "wav",
     stream: bool = False,
 ) -> dict:
-    """Build the /v1/audio/speech request payload."""
+    """Build the /v1/audio/speech request payload.
+
+    Fish Speech has no built-in speakers; omit ``voice`` so the model
+    generates with a random timbre (controllable via text tags like
+    ``[cool]``, ``[warm]``, etc.).  When ``ref_audio`` is supplied we
+    enter voice-cloning mode instead.
+    """
     if not text or not text.strip():
         raise gr.Error("Please enter text to synthesize.")
 
     payload: dict = {
         "input": text.strip(),
-        "voice": "default",
         "response_format": "pcm" if stream else response_format,
         "stream": stream,
     }
@@ -193,7 +198,8 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
 
                 with gr.Accordion("Voice Cloning (optional)", open=False):
                     gr.Markdown(
-                        "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
+                        "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. "
+                        "Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output."
                     )
                     ref_audio = gr.Audio(
                         label="Reference Audio",

examples/online_serving/text_to_speech/fish_speech/speech_client.py

@@ -39,7 +39,6 @@ def run_tts(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "voice": "default",
         "response_format": args.response_format,
     }
 

examples/quantization/check_modelopt_fp8_export.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Verify a ModelOpt FP8 diffusers checkpoint exported by
+quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
+
+Three checks:
+  A. transformer/config.json has a sane quantization_config block.
+  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
+  C. transformer disk size is materially smaller than a BF16 baseline.
+
+Example:
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8
+
+    # Optional: compare disk size against a local or HF BF16 baseline.
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import Counter
+from pathlib import Path
+
+SUPPORTED_ALGOS = ("FP8",)
+
+
+def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
+    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
+    cfg_path = transformer_dir / "config.json"
+    if not cfg_path.exists():
+        print(f"[FAIL] {cfg_path} missing.")
+        return 1, None
+
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
+        return 1, None
+
+    print(f"[A] quantization_config from {cfg_path}:")
+    print(json.dumps(qc, indent=2))
+
+    quant_algo = qc.get("quant_algo")
+    issues = []
+    if qc.get("quant_method") != "modelopt":
+        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
+    if quant_algo not in SUPPORTED_ALGOS:
+        issues.append(
+            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
+        )
+
+    if issues:
+        print("[A] WARN — config looks incomplete:")
+        for issue in issues:
+            print(f"    - {issue}")
+        return 2, quant_algo
+    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
+    return 0, quant_algo
+
+
+def _read_safetensors_header(path: Path) -> dict:
+    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
+
+    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
+    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
+    """
+    import struct
+
+    with open(path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_len))
+    header.pop("__metadata__", None)
+    return header
+
+
+def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
+    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
+
+    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
+    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
+    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
+    """
+    if not weight_scale_shapes:
+        return "no weight_scale tensors found"
+
+    def meaningful_dims(shape: list[int]) -> int:
+        return sum(1 for d in shape if d > 1)
+
+    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
+    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
+    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
+    total = len(weight_scale_shapes)
+    if per_tensor == total:
+        return "per-tensor (all scalar scales)"
+    if per_channel == total:
+        return "per-channel (1 meaningful dim)"
+    if per_block == total:
+        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
+    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
+
+
+def _check_safetensors(transformer_dir: Path) -> int:
+    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
+    files = sorted(transformer_dir.glob("*.safetensors"))
+    if not files:
+        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
+        return 1
+
+    header_dtype_counts: Counter[str] = Counter()
+    sample_quant_weight_keys: list[str] = []
+    sample_scale_keys: list[str] = []
+    weight_scale_shapes: list[list[int]] = []
+    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
+
+    for f in files:
+        try:
+            header = _read_safetensors_header(f)
+        except Exception as exc:
+            print(f"[B] WARN — could not parse header of {f}: {exc}")
+            continue
+        for k, info in header.items():
+            dtype = info.get("dtype", "?")
+            header_dtype_counts[dtype] += 1
+            # FP8 stores weights as F8_E4M3 directly.
+            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
+                sample_quant_weight_keys.append(k)
+            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
+                sample_scale_keys.append(k)
+            if k.endswith(".weight_scale"):
+                weight_scale_shapes.append(info.get("shape", []))
+                if len(sample_weight_scale_entries) < 5:
+                    sample_weight_scale_entries.append((k, info.get("shape", [])))
+
+    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
+    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
+        marker = "  <-- FP8" if dtype.startswith("F8") else ""
+        print(f"    {dtype:10s} {count:>6d}{marker}")
+
+    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
+    if quant_count == 0:
+        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
+        return 1
+
+    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
+    if sample_quant_weight_keys:
+        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
+    if sample_scale_keys:
+        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
+    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
+    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
+
+    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
+    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
+    for key, shape in sample_weight_scale_entries[:3]:
+        print(f"      {key}: shape {shape}")
+    return 0
+
+
+def _disk_size_gib(p: Path) -> float:
+    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
+
+
+def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
+    """Returns 0 always (informational only)."""
+    quant_size = _disk_size_gib(transformer_dir)
+    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
+
+    if baseline is None:
+        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
+        return 0
+
+    baseline_path = Path(baseline)
+    if not baseline_path.exists():
+        # Try HF download.
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
+            return 0
+        print(f"    Downloading baseline transformer from HF: {baseline}")
+        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
+
+    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
+    bf16_size = _disk_size_gib(bf16_dir)
+    if bf16_size == 0:
+        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
+        return 0
+
+    # Expected reduction for FP8: ~50%.
+    min_reduction = 30
+    reduction = (1 - quant_size / bf16_size) * 100
+    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
+    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
+    if reduction < min_reduction:
+        print(
+            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
+        )
+    return 0
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
+    p.add_argument(
+        "--baseline",
+        default=None,
+        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
+    )
+    args = p.parse_args()
+
+    out_root = Path(args.output).expanduser().resolve()
+    transformer_dir = out_root / "transformer"
+    if not transformer_dir.exists():
+        print(f"[FAIL] {transformer_dir} does not exist.")
+        sys.exit(1)
+
+    print(f"Checking: {out_root}\n")
+
+    fail = 0
+    config_status, _quant_algo = _check_config(transformer_dir)
+    fail |= config_status
+    fail |= _check_safetensors(transformer_dir)
+    _check_size_vs_baseline(transformer_dir, args.baseline)
+
+    print()
+    if fail == 0:
+        print("=" * 60)
+        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
+    elif fail == 1:
+        print("=" * 60)
+        print("FAILURES detected — calibration may need to be re-run.")
+        sys.exit(1)
+    else:
+        print("=" * 60)
+        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: modulation,
+AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
+and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
+empirically degrades under FP8 (see #2920 ablation).
+
+Example:
+    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
+        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=480)
+    p.add_argument("--width", type=int, default=832)
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=33,
+        help="Frames per calibration sample. 33 matches the typical short benchmark.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=6.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
+        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
+        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
+# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
+# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
+# - proj_out: final output projection
+# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
+def _filter_func_hv15(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
+        r"time_embed|cond_type_embed|"
+        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
+        r"token_refiner).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_hv15)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try to set guidance on the pipeline's guider object up front (modern
+    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
+    # back silently — calibration uses whatever default the pipeline ships with.
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines (like HV-1.5) that take CFG via guider config.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
+    helper.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
+
+    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
+    `weight_block_size` is given, advertise block-wise weight quantization in
+    the saved metadata (so consumers know to expect multi-element scale tensors).
+    """
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "context_embedder*",
+            "context_embedder_2*",
+            "cond_type_embed*",
+            "image_embedder*",
+            "norm1.linear*",
+            "norm1_context.linear*",
+            "norm2*",
+            "norm2_context*",
+            "norm_out*",
+            "proj_out*",
+            "time_embed*",
+            "token_refiner*",
+            "x_embedder*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
+    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
+    cfg_path = output_dir / "transformer" / "config.json"
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
+    existing = cfg.get("quantization_config")
+    if isinstance(existing, dict):
+        producer = existing.get("producer")
+        if isinstance(producer, dict):
+            new_qc["producer"] = producer
+
+    cfg["quantization_config"] = new_qc
+    with cfg_path.open("w", encoding="utf-8") as f:
+        json.dump(cfg, f, indent=2)
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Save the pipeline with the (now FP8) transformer.
+
+    Copies the source directory verbatim except for `transformer/`, then
+    saves the transformer with quantizers hidden so the state dict contains
+    only the FP8 weights + scale tensors.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
+
+    transformer_out = output_dir / "transformer"
+    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
+    # nn.Module (transformer), not the diffusers Pipeline wrapper.
+    with hide_quantizers_from_state_dict(pipe.transformer):
+        pipe.transformer.save_pretrained(
+            str(transformer_out),
+            safe_serialization=True,
+            max_shard_size=max_shard_size,
+        )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+
+    print("Quantization plan:")
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    backbone = pipe.transformer
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        # Switch from per-tensor (default) to block-wise weight quantization.
+        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+    quantized = mtq.quantize(backbone, quant_config, forward_loop)
+    if quantized is not None:
+        pipe.transformer = quantized
+        backbone = quantized
+
+    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
+
+    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    exported = _force_export_quantized_weights(backbone, dtype)
+    print(f"  -> {exported} weights converted to FP8 in memory")
+    if exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
+        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
+        "    --enforce-eager"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_wan2_2_modelopt_fp8.py

@@ -0,0 +1,471 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: condition embedder
+(time/text/image), patch embedding, modulation (scale_shift_table), final
+norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
+are quantized — static calibration handles the numerics that online FP8
+couldn't (see #2920 ablation).
+
+Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
+fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
+
+Example:
+    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
+        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
+        --output ./wan22-ti2v-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
+    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=49,
+        help="Frames per calibration sample. 49 matches the typical short benchmark; "
+        "use 17 to reduce memory pressure during calibration.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=5.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
+        "sequences amplified FP8 drift in the online ablation (see #2920).",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
+        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
+        "for the HV-1.5 investigation status before relying on this.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision. Wan2.2's module naming:
+# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
+# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
+# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
+# - norm_out: AdaLayerNorm final
+# - proj_out: final nn.Linear
+# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
+def _filter_func_wan22(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(condition_embedder|patch_embedding|"
+        r"norm_out|scale_shift_table|"
+        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_wan22)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines that take CFG via guider config only.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
+    calibration helpers.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "condition_embedder*",
+            "norm_out*",
+            "output_scale_shift_prepare*",
+            "patch_embedding*",
+            "proj_out*",
+            "scale_shift_table*",
+            "timestep_proj_prepare*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
+    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
+
+    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
+    for sub in ("transformer", "transformer_2"):
+        cfg_path = output_dir / sub / "config.json"
+        if not cfg_path.exists():
+            continue
+        with cfg_path.open(encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
+        existing = cfg.get("quantization_config")
+        if isinstance(existing, dict):
+            producer = existing.get("producer")
+            if isinstance(producer, dict):
+                new_qc["producer"] = producer
+
+        cfg["quantization_config"] = new_qc
+        with cfg_path.open("w", encoding="utf-8") as f:
+            json.dump(cfg, f, indent=2)
+
+
+def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
+    """Return [(attr_name, module)] for every DiT on the pipeline.
+
+    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
+    (high-noise). TI2V-5B has only ``transformer``.
+    """
+    out = []
+    for attr in ("transformer", "transformer_2"):
+        mod = getattr(pipe, attr, None)
+        if mod is not None:
+            out.append((attr, mod))
+    return out
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
+
+    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
+    as the single-DiT TI2V-5B.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
+
+    for attr, backbone in _list_transformers(pipe):
+        transformer_out = output_dir / attr
+        # Pass the nn.Module (transformer), not the Pipeline wrapper.
+        with hide_quantizers_from_state_dict(backbone):
+            backbone.save_pretrained(
+                str(transformer_out),
+                safe_serialization=True,
+                max_shard_size=max_shard_size,
+            )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print("Quantization plan:")
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    transformers = _list_transformers(pipe)
+    if not transformers:
+        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
+    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+
+    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
+    # on every attached transformer (A14B: transformer + transformer_2) get
+    # calibrated during the same passes.
+    for attr, backbone in transformers:
+        print(f"\nQuantizing {attr}...")
+        quantized = mtq.quantize(backbone, quant_config, forward_loop)
+        if quantized is not None:
+            setattr(pipe, attr, quantized)
+        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
+
+    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
+    # inside the loop would convert transformer's weights to real FP8, and the
+    # next transformer's forward_loop (full pipeline) would then fake-quantize an
+    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
+    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    total_exported = 0
+    for attr, backbone in _list_transformers(pipe):
+        exported = _force_export_quantized_weights(backbone, dtype)
+        print(f"  -> {exported} weights converted to FP8 in {attr}")
+        total_exported += exported
+
+    if total_exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer(s)...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
+        "    --output outputs/wan22_modelopt_fp8.mp4"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -90,6 +90,12 @@ minicpmo = [
     "stepaudio2-minicpmo",
 ]
 
+# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
+# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
+quack = [
+    "quack-kernels>=0.3.11",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

recipes/README.md

@@ -36,8 +36,8 @@ recipes/
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
-| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy  | 1x H200 141GB / B300 |
-| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
+| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
+| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
@@ -50,6 +50,7 @@ recipes/
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
+| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 |

recipes/StabilityAI/Stable-Diffusion-3.5.md

@@ -0,0 +1,139 @@
+# Stable Diffusion 3.5
+
+> Text-to-image serving with Stability AI's SD 3.5 family
+
+## Summary
+
+- Vendor: Stability AI
+- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible images API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a lightweight, fast text-to-image model that
+fits on a single consumer or workstation GPU. SD 3.5-medium requires only
+~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of
+the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces
+noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per
+image.
+
+## References
+
+- Model (medium): <https://huggingface.co/stabilityai/stable-diffusion-3.5-medium>
+- Model (large): <https://huggingface.co/stabilityai/stable-diffusion-3.5-large>
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
+- User guide:
+  [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
+
+## Hardware Support
+
+This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB,
+one for each model variant. Extend it with more hardware sections as community
+validation lands.
+
+## GPU
+
+### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium
+
+#### Environment
+
+- OS: Ubuntu 22.04
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
+
+#### Command
+
+```bash
+vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091
+```
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "stabilityai/stable-diffusion-3.5-medium",
+    "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
+    "negative_prompt": "blurry, out of focus",
+    "guidance_scale": 4.5,
+    "n": 1,
+    "size": "1024x1024"
+  }' | python3 -c "
+import json, base64, sys
+data = json.load(sys.stdin)
+with open('output.png', 'wb') as f:
+    f.write(base64.b64decode(data['data'][0]['b64_json']))
+print('saved output.png')
+"
+```
+
+Or use the example client:
+
+```bash
+python examples/online_serving/text_to_image/openai_chat_client.py \
+  --server http://localhost:8091 \
+  --prompt "a cute cat sitting on a windowsill" \
+  --output output.png
+```
+
+#### Notes
+
+- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled).
+- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`.
+- Model loading time: ~55 seconds (including weight download on first run).
+- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
+
+### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large
+
+#### Environment
+
+- OS: Ubuntu 22.04
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
+
+#### Command
+
+```bash
+vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091
+```
+
+> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license
+> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)
+> and set `HF_TOKEN` before running.
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "stabilityai/stable-diffusion-3.5-large",
+    "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
+    "negative_prompt": "blurry, out of focus",
+    "guidance_scale": 4.5,
+    "n": 1,
+    "size": "1024x1024"
+  }' | python3 -c "
+import json, base64, sys
+data = json.load(sys.stdin)
+with open('output.png', 'wb') as f:
+    f.write(base64.b64decode(data['data'][0]['b64_json']))
+print('saved output.png')
+"
+```
+
+#### Notes
+
+- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`.
+- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048.
+- Image quality is noticeably better than medium, with sharper details.
+- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
+- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).

recipes/cosmos3/Cosmos3-Nano.md

@@ -1,34 +1,43 @@
 # Cosmos3-Nano
 
-> Text-to-image, text-to-video, and image-to-video serving
+> Text-to-image, text-to-video, image-to-video, and video-to-video serving
 
 ## Summary
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Nano`
-- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy
+- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
 - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
 - Maintainer: Community
 
 ## When to use this recipe
 
 Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
-A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
-the mode is selected per request:
+A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the
+mode is selected per request:
 
 - **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
-- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
+- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
 - **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
   upload, or `image_reference` JSON).
+- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
+  video upload, or `video_reference` JSON). Cosmos3 conditions on selected
+  reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
+  and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
+  generation.
 - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
   T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
   the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
 - **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
-  - `forward_dynamics` — given a first frame **and** an action trajectory, roll out
-    the resulting video. Synchronous: `POST /v1/videos/sync`.
-  - `policy` — given a first frame and a language instruction, **predict** the action
-    trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
-    read the predicted action from the top-level `action` field
+  - `forward_dynamics` — given a first frame or video **and** an action trajectory,
+    roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
+  - `policy` — given a first frame or video and a language instruction,
+    **predict** the action trajectory (and a rollout video). Use the async
+    `POST /v1/videos` endpoint and read the predicted action from the top-level
+    `action` field.
+  - `inverse_dynamics` — given a video, **recover** the action trajectory. Use
+    the async `POST /v1/videos` endpoint and read the recovered action from
+    the top-level `action` field
     (`{data, shape, dtype, raw_action_dim, domain_id}`).
 
   Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
@@ -38,8 +47,9 @@ the mode is selected per request:
   **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
   (`domain_name=droid_lerobot`).
 
-  `inverse_dynamics` (recover the action from a given video) is supported by the
-  pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
+  Action requests can use `input_reference` or `video_reference` for video input.
+  `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
+  requires a video reference.
 
 ## References
 
@@ -142,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
   -o cosmos3_i2v.mp4
 
+# Video-to-video -> /v1/videos/sync with an uploaded reference video.
+# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
+# temporal VAE stride this decodes only the first 5 input frames.
+# The model works best when the prompt describes the actual situation happening in the video.
+# Generic prompts may create sub-standard generations.
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
+  -F "seed=2222" \
+  -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
+  -o cosmos3_v2v.mp4
+
+# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
+# video_reference with input_reference or image_reference.
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
+  -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
+  -o cosmos3_v2v_from_url.mp4
 
 # Text-to-video-with-sound
 curl -sS -X POST http://localhost:8000/v1/videos/sync \
@@ -190,6 +230,21 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
 # poll until status == completed, then:
 curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
 curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
+
+# Action — inverse dynamics (video -> recovered action trajectory).
+# Asynchronous: use the job metadata to read the recovered action.
+VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
+  -H "Accept: application/json" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=Recover the robot action trajectory from this clip." \
+  -F "input_reference=@motion_clip.mp4;type=video/mp4" \
+  -F "size=640x480" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
+  -F "seed=0" | jq -r '.id')
+# poll until status == completed, then:
+curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
+curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
 ```
 
 #### Notes
@@ -206,24 +261,27 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic
 - **Determinism:** identical seed reproduces identical output on the same
   hardware; outputs are not bit-identical across different GPU types.
 - **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
-  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
-  189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
+  1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
 - **Key flags / params:** `--no-guardrails` (server) or
   `extra_params={"guardrails":false}` (per request) toggles safety. The
   per-request flag only takes effect when the server was launched **with**
   guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
   `use_resolution_template` / `use_duration_template` are off by default and only
   needed when not using upsampled prompts that already encode resolution/duration.
+  For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
+  frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
+  API decodes the first or last needed reference frames (`"first"` by default).
 - **Known limitations:**
   - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
     `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
     the server fails at pipeline build with a gated-repo / safety-checker error.
   - A guardrail-blocked prompt currently returns HTTP 500
     (`"Guardrail blocked prompt"`).
-  - Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
-    `/v1/videos`, returns the predicted action under the top-level `action`
-    field) are supported online. **Online inference of inverse dynamics will be
-    added in a follow-up MR.**
+  - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
+    online. Use async `POST /v1/videos` when you need the predicted/recovered
+    action payload under the top-level `action` field; sync `/v1/videos/sync`
+    returns raw MP4 bytes and does not expose action metadata in the response body.
 
 ### 1x GPU (Offline generation)
 
@@ -287,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
 
 - Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
   `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
-  `num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
-  frames are returned in `outputs[0].request_output.images` as an
+  `num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
+  `multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
+  For video, frames are returned in `outputs[0].request_output.images` as an
   `(B, F, H, W, 3)` array.
 - The offline entry must be guarded by `if __name__ == "__main__":` — the engine
   spawns workers with the `spawn` start method.

recipes/cosmos3/Cosmos3-Super.md

@@ -1,12 +1,12 @@
 # Cosmos3-Super
 
-> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
+> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
 
 ## Summary
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
-- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
+- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
 - Mode: Online serving with the OpenAI-compatible image/video APIs
 - Maintainer: Community
 
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
 Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
 formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
 recommended parallelism differ. Mode is selected per request (T2I →
-`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
-for audio).
+`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
+`generate_sound=true` for audio).
 
 ## References
 
@@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
 #### Verification
 
 Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
-T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
-num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
+T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
+fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
+max_sequence_length=4096`.
 
 ```bash
 curl http://localhost:8000/v1/models
@@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
   -o cosmos3_super_i2v.mp4
 
+# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
+  -F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
+  -o cosmos3_super_v2v.mp4
+
 # T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
 curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
@@ -103,10 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
   - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
 - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
-- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
-  semantics as Nano, including the **action** modality: `forward_dynamics`
-  (sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
-  the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
-  Online inference of inverse dynamics will be added in a follow-up MR. Verified on
-  the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
-  action (`[16, 10]`) and the rollout video reliably.
+- Same generation defaults, supported sizes, V2V reference-video controls
+  (`condition_frame_indexes_vision`, `condition_video_keep`), and
+  `generate_sound`/`sound_duration`
+  semantics as Nano, including the **action** modality: `forward_dynamics`,
+  `policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
+  shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
+  under the top-level `action` field. Verified on the 64B Super under
+  `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
+  and the rollout video reliably.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,19 +1,46 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+"""
+Tests for the DiffusersPipelineLoader.
+"""
+
 from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
+from huggingface_hub import snapshot_download
+from vllm.config.load import LoadConfig
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
+model_path = "hf-internal-testing/tiny-helios-modular-pipe"
+
+
+@pytest.fixture(scope="module")
+def prefetch_helios_model():
+    """Downloads the tiny helios model prior to running a test."""
+    snapshot_download(model_path)
+
+
+@pytest.fixture(scope="function")
+def mock_tp_group(mocker):
+    """Mocks the tensor parallel group; this is needed to initialize the Helios model."""
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+    mock_group = mocker.MagicMock()
+    mock_group.world_size = 1
+    mock_group.rank_in_group = 0
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
 
 class _DummyPipelineModel(nn.Module):
     def __init__(self, *, source_prefix: str):
@@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module):
 
 
 def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
 
@@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
         quantization_config=None,
     )
 
-    loader = object.__new__(DiffusersPipelineLoader)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda _od_config: False  # type: ignore[assignment]
+    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
 
     model = loader.load_model(
-        od_config,
         load_device="cpu",
         load_format="custom_pipeline",
         custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
@@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert model.captured_config is od_config
     assert model.seen_config_during_init is od_config
     assert get_current_diffusion_config_or_none() is None
+
+
+def test_get_all_weights(prefetch_helios_model, mock_tp_group):
+    """Ensure that get all weights on a tiny model resolves to nonempty weights."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    pipeline = HeliosPipeline(od_config=od_config)
+
+    weights = list(loader.get_all_weights(pipeline))
+    assert len(weights) > 0
+
+
+def test_load_model(prefetch_helios_model, mock_tp_group):
+    """Ensure that load model creates an instance of the expected pipeline class."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    model = loader.load_model(load_device="cpu")
+    assert isinstance(model, HeliosPipeline)

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader:
         ignore_patterns=["*.tmp"],
     )
     loader.od_config = None
+    loader.quant_config = None
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
     return loader
@@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader:
 
 def test_is_gguf_quantization_accepts_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config={
-            "method": "gguf",
-            "gguf_model": "weights.gguf",
-        }
-    )
+    loader.quant_config = {
+        "method": "gguf",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_is_gguf_quantization_rejects_non_gguf_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config={
-            "method": "fp8",
-            "gguf_model": "weights.gguf",
-        }
-    )
+    loader.quant_config = {
+        "method": "fp8",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is False
+    assert loader._is_gguf_quantization() is False
 
 
 def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(quantization_config={"method": "gguf"})
+    loader.quant_config = {"method": "gguf"}
 
     with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization(od_config)
+        loader._is_gguf_quantization()
 
 
 def test_is_gguf_quantization_accepts_object_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config=SimpleNamespace(
-            get_name=lambda: "gguf",
-            gguf_model="weights.gguf",
-        )
+    loader.quant_config = SimpleNamespace(
+        get_name=lambda: "gguf",
+        gguf_model="weights.gguf",
     )
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_is_gguf_quantization_uses_fallback_object_without_get_name():
     loader = _make_loader()
-    od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_get_model_loadable_names_collects_parameters_and_buffers():
@@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
 
 def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
+    loader.od_config = SimpleNamespace(revision="main")
     source = DiffusersPipelineLoader.ComponentSource(
         model_or_path="dummy",
         subfolder="transformer",
         revision=None,
         prefix="transformer.",
     )
-    od_config = SimpleNamespace(
-        revision="main",
-        quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
-    )
 
     class _Adapter:
         def weights_iterator(self):
@@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
     monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
     monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
 
-    weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
+    weights = list(loader._get_gguf_weights_iterator(source, object()))
 
     assert weights[0][0] == "transformer.weight"
     assert torch.equal(weights[0][1], torch.ones((2, 2)))
@@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
 def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
     model = _DummyModel()
-    od_config = SimpleNamespace()
 
     monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
     monkeypatch.setattr(
         loader,
         "_get_gguf_weights_iterator",
-        lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
+        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
     )
 
     def _hf_iter(source):
@@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert model.calls[0] == ["transformer.weight"]
@@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
 def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
     model = _DummyModel()
-    od_config = SimpleNamespace()
 
     monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
     monkeypatch.setattr(
         loader,
         "_get_gguf_weights_iterator",
-        lambda source, model, od_config: iter(
+        lambda source, model: iter(
             [
                 ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
                 ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
@@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert hf_calls == ["vae"]

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -374,6 +374,19 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
     assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
     assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
 
+    frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
+    v2v = SimpleNamespace(
+        prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
+        sampling_params=SimpleNamespace(
+            height=16,
+            width=32,
+            extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
+        ),
+    )
+    additional = preprocess(v2v).prompts[0]["additional_information"]
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
+    assert additional["condition_frame_indexes_vision"] == [0, 1]
+
 
 def test_postprocess_handles_image_video_audio_and_validation() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
@@ -464,6 +477,19 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
     assert image_latent.shape == (1, 2, 1, 2, 3)
 
+    pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
+    v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
+        torch.zeros(1, 3, 5, 16, 24),
+        16,
+        24,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+        [0, 1],
+    )
+    torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
+    assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
+    assert v2v_condition.shape == (1, 2, 3, 2, 3)
+
     pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
     pipeline._sound_tokenizer = SimpleNamespace(
         sample_rate=10,
@@ -670,6 +696,34 @@ class TestForwardRouting:
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
 
+        video_tensor = torch.zeros(1, 3, 5, 16, 16)
+        v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
+        v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
+        pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
+            torch.zeros(1, 2, 2, 1, 1),
+            v2v_mask,
+            v2v_condition,
+        )
+        pipeline.forward(
+            SimpleNamespace(
+                prompts=[
+                    {
+                        "prompt": "continue",
+                        "modalities": ["video"],
+                        "additional_information": {
+                            "preprocessed_video": video_tensor,
+                            "condition_frame_indexes_vision": [0],
+                        },
+                    }
+                ],
+                sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
+            )
+        )
+        assert captured["flow_shifts"][-1] == 10.0
+        assert captured["format"]["negative_prompt"] == ""
+        assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
+        assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
+
         pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
         sound_latents = torch.zeros(1, 3, 4)
         pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)

tests/diffusion/models/hunyuan_video/__init__.py

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

tests/diffusion/models/hunyuan_video/test_hunyuan_video_quant_config_propagation.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation.
+
+Tests cover:
+- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel
+- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module
+import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TestHunyuanVideoQuantConfigPropagation:
+    """Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines."""
+
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device",
+        return_value="cpu",
+    )
+    @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders")
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs",
+        return_value={},
+    )
+    def test_t2v_quant_config_passed_through(
+        self,
+        mock_get_kwargs,
+        mock_scheduler,
+        mock_t5_encoder,
+        mock_auto_config,
+        mock_byt5_tok,
+        mock_qwen2_tok,
+        mock_from_pretrained,
+        mock_prefetch,
+        mock_get_device,
+    ):
+        captured_kwargs = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured_kwargs.update(kwargs)
+
+        with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
+            fake_qc = MagicMock()
+            od_config = SimpleNamespace(
+                model="fake-t2v-model",
+                dtype=torch.bfloat16,
+                flow_shift=None,
+                tf_model_config=SimpleNamespace(),
+                quantization_config=fake_qc,
+                enable_diffusion_pipeline_profiler=False,
+            )
+
+            t2v_module.HunyuanVideo15Pipeline(od_config=od_config)
+
+            assert captured_kwargs.get("quant_config") is fake_qc
+
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device",
+        return_value="cpu",
+    )
+    @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders")
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs",
+        return_value={},
+    )
+    def test_i2v_quant_config_passed_through(
+        self,
+        mock_get_kwargs,
+        mock_scheduler,
+        mock_siglip_processor,
+        mock_t5_encoder,
+        mock_auto_config,
+        mock_byt5_tok,
+        mock_qwen2_tok,
+        mock_from_pretrained,
+        mock_prefetch,
+        mock_get_device,
+    ):
+        captured_kwargs = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured_kwargs.update(kwargs)
+
+        with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
+            fake_qc = MagicMock()
+            od_config = SimpleNamespace(
+                model="fake-i2v-model",
+                dtype=torch.bfloat16,
+                flow_shift=None,
+                tf_model_config=SimpleNamespace(),
+                quantization_config=fake_qc,
+                enable_diffusion_pipeline_profiler=False,
+            )
+
+            i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config)
+
+            assert captured_kwargs.get("quant_config") is fake_qc

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -72,6 +72,31 @@ class TestPipelineIndependence:
 
         assert issubclass(LTX23Pipeline, CFGParallelMixin)
 
+    def test_ltx23_pipeline_declares_offload_components(self):
+        """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+        pipe = object.__new__(LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer = torch.nn.Linear(1, 1)
+        pipe.text_encoder = torch.nn.Linear(1, 1)
+        pipe.connectors = torch.nn.Linear(1, 1)
+        pipe.vae = torch.nn.Linear(1, 1)
+        pipe.audio_vae = torch.nn.Linear(1, 1)
+        pipe.vocoder = torch.nn.Linear(1, 1)
+
+        modules = ModuleDiscovery.discover(pipe)
+
+        assert LTX23Pipeline._dit_modules == ["transformer"]
+        assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
+        assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
+        assert LTX23Pipeline._resident_modules == ["vocoder"]
+        assert modules.dit_names == ["transformer"]
+        assert modules.encoder_names == ["text_encoder", "connectors"]
+        assert modules.resident_names == ["vocoder"]
+        assert len(modules.vaes) == 2
+
 
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from torch import nn
+
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
+
+
+def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
+
+    assert "weight" not in dict(norm.named_parameters())
+    assert "weight" in dict(norm.named_buffers())
+    assert "weight" not in norm.state_dict()
+
+
+def test_ltx_rms_norm_affine_weight_remains_parameter():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
+
+    assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
+    assert "weight" not in dict(norm.named_buffers())
+    assert "weight" in norm.state_dict()

tests/diffusion/test_worker_wrapper_base.py

@@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture):
     config.cache_backend = None
     config.cache_config = None
     config.model = "test-model"
-    config.diffusion_load_format = None
+    config.diffusion_load_format = "default"
     config.dtype = "float32"
     config.max_cpu_loras = 0
     config.lora_path = None

tests/e2e/accuracy/test_videogen_modelopt_quant.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs.
+
+Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``:
+for each calibrated FP8 checkpoint the test generates one video with BF16 and one
+with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the
+middle frame, then asserts both ``CLIP >= absolute_floor`` and
+``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model
+paths are local) so it never runs in generic CI.
+
+To run::
+
+    VIDEOGEN_RUN_QUANT_ACCURACY=1 \\
+    WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\
+    WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\
+    HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+    HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\
+        pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+# ---------------------------------------------------------------------------
+# Tunables (override via env to suit local hardware / calibration budget)
+# ---------------------------------------------------------------------------
+SEED = 42
+HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480"))
+WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832"))
+NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17"))
+NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20"))
+
+WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+
+# CLIP gates: absolute floor + max drop from the BF16 reference.
+CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0"))
+CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0"))
+
+RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY"
+_TRUE = {"1", "true", "yes", "on"}
+
+
+# ---------------------------------------------------------------------------
+@dataclass(frozen=True)
+class _Spec:
+    """A model under test (BF16 + its FP8 variant)."""
+
+    family: str
+    bf16_env: str
+    fp8_env: str
+    prompt: str
+    guidance_scale: float
+    guidance_scale_2: float | None = None
+    boundary_ratio: float | None = None
+
+
+WAN_SPEC = _Spec(
+    family="wan22_a14b_t2v",
+    bf16_env="WAN22_A14B_BF16_MODEL",
+    fp8_env="WAN22_A14B_FP8_MODEL",
+    prompt=WAN_PROMPT,
+    guidance_scale=3.5,
+    guidance_scale_2=4.0,
+    boundary_ratio=0.875,
+)
+
+HV_SPEC = _Spec(
+    family="hunyuanvideo15_t2v",
+    bf16_env="HV15_BF16_MODEL",
+    fp8_env="HV15_FP8_MODEL",
+    prompt=HV_PROMPT,
+    guidance_scale=6.0,
+)
+
+
+@dataclass(frozen=True)
+class _QuantCase:
+    name: str
+    spec: _Spec
+    clip_drop: float
+
+
+def _cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8),
+        _QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8),
+    ]
+    params: list[pytest.ParameterSet] = []
+    opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE
+    for case in cases:
+        marks: list[Any] = []
+        if not opted_in:
+            marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy."))
+        if not os.environ.get(case.spec.bf16_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}."))
+        if not os.environ.get(case.spec.fp8_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
+# ---------------------------------------------------------------------------
+def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image:
+    if isinstance(frames, np.ndarray):
+        idx = frames.shape[0] // 2
+        arr = frames[idx]
+        if arr.dtype != np.uint8:
+            arr = (arr * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(arr).convert("RGB")
+    if not frames:
+        raise AssertionError("Empty video output")
+    mid = frames[len(frames) // 2]
+    if isinstance(mid, np.ndarray):
+        if mid.dtype != np.uint8:
+            mid = (mid * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(mid).convert("RGB")
+    return mid.convert("RGB")
+
+
+def _extract_frames(result) -> list[Image.Image]:
+    """Pull a frame list out of an Omni text-to-video result."""
+    from vllm_omni.entrypoints.omni import OmniRequestOutput
+
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
+        result = result.request_output
+    images = getattr(result, "images", None)
+    if not images:
+        raise AssertionError("Pipeline output had no images")
+    if isinstance(images, list) and images and isinstance(images[0], tuple):
+        # (frames, fps)
+        images = images[0][0]
+    return list(images)
+
+
+def _generate_video(
+    *,
+    model: str,
+    spec: _Spec,
+    quantization: str | None,
+    output_path: Path,
+) -> tuple[Image.Image, float]:
+    """Generate one video and return its middle frame + elapsed wall time."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    kwargs: dict[str, Any] = dict(
+        model=model,
+        parallel_config=DiffusionParallelConfig(),
+        enforce_eager=True,
+        vae_use_tiling=True,
+    )
+    if quantization:
+        kwargs["quantization"] = quantization
+    if spec.boundary_ratio is not None:
+        kwargs["boundary_ratio"] = spec.boundary_ratio
+    omni = Omni(**kwargs)
+
+    try:
+        gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED)
+        sampling = dict(
+            height=HEIGHT,
+            width=WIDTH,
+            generator=gen,
+            guidance_scale=spec.guidance_scale,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            num_frames=NUM_FRAMES,
+            seed=SEED,
+        )
+        if spec.guidance_scale_2 is not None:
+            sampling["guidance_scale_2"] = spec.guidance_scale_2
+
+        import time as _time
+
+        t0 = _time.perf_counter()
+        out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling))
+        elapsed = _time.perf_counter() - t0
+
+        frames = _extract_frames(out)
+        mid = _middle_frame(frames)
+        mid.save(output_path)
+        return mid, elapsed
+    finally:
+        del omni
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+
+
+@pytest.mark.parametrize("case", _cases())
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 1,
+    reason="Needs at least 1 GPU.",
+)
+def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None:
+    """FP8 DiT should preserve prompt-aligned video quality vs BF16."""
+    out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant")
+
+    with tempfile.TemporaryDirectory():
+        bf16_model = os.environ[case.spec.bf16_env]
+        quant_model = os.environ[case.spec.fp8_env]
+
+        bf16_frame, bf16_time = _generate_video(
+            model=bf16_model,
+            spec=case.spec,
+            quantization=None,
+            output_path=out_dir / "bf16.png",
+        )
+        quant_frame, quant_time = _generate_video(
+            model=quant_model,
+            spec=case.spec,
+            quantization="fp8",
+            output_path=out_dir / f"{case.name}.png",
+        )
+
+    clip = CLIPScorer()
+    bf16_clip = clip.score(bf16_frame, case.spec.prompt)
+    quant_clip = clip.score(quant_frame, case.spec.prompt)
+    clip_drop = bf16_clip - quant_clip
+    speedup = (bf16_time / quant_time) if quant_time else float("nan")
+
+    metrics = {
+        "case": case.name,
+        "family": case.spec.family,
+        "quantization": "fp8",
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": case.spec.prompt,
+        "height": HEIGHT,
+        "width": WIDTH,
+        "num_frames": NUM_FRAMES,
+        "num_inference_steps": NUM_INFERENCE_STEPS,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "speedup_vs_bf16": speedup,
+        "bf16_clip_score": bf16_clip,
+        "quant_clip_score": quant_clip,
+        "clip_score_drop": clip_drop,
+    }
+    metrics_path = out_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})")
+    print(f"           CLIP BF16={bf16_clip:.2f}  FP8={quant_clip:.2f}  drop={clip_drop:.2f}")
+    print(f"           metrics={metrics_path}")
+
+    assert quant_clip >= CLIP_ABSOLUTE_FLOOR, (
+        f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}"
+    )
+    assert clip_drop <= case.clip_drop, (
+        f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}"
+    )

tests/entrypoints/openai_api/test_serving_speech.py

@@ -33,6 +33,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     SpeechBatchItem,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
+    _TTS_LANGUAGES,
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
@@ -1257,6 +1258,67 @@ class TestTTSMethods:
         # Verify speakers are normalized to lowercase
         assert server.supported_speakers == {"ryan", "vivian", "aiden"}
 
+    def test_load_supported_languages_from_config(self, speech_server):
+        """Languages/dialects from codec_language_id are loaded title-cased; 'Auto' is added."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                talker_config=SimpleNamespace(
+                    codec_language_id={"chinese": 2055, "english": 2050, "beijing_dialect": 2074}
+                )
+            )
+        )
+        assert speech_server._load_supported_languages() == {
+            "Chinese",
+            "English",
+            "Beijing_Dialect",
+            "Auto",
+        }
+
+    def test_load_supported_languages_from_dict_config(self, speech_server):
+        """talker_config provided as a plain dict is handled the same as an object."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}})
+        )
+        assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"}
+
+    def test_validate_language_custom_dialect_accepted(self, speech_server):
+        """A language present in the model config passes validation, case-insensitively."""
+        speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"}
+        for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"):
+            req = OpenAICreateSpeechRequest(input="Hello", language=language)
+            result = speech_server._validate_tts_request(req)
+            assert result is None or "Invalid language" not in result
+            # Language is normalized to the title-cased config form.
+            assert req.language == language.title()
+
+    def test_validate_language_unknown_rejected(self, speech_server):
+        """A language not in the configured set is rejected."""
+        speech_server.supported_languages = {"Chinese", "English", "Auto"}
+        for language in ("Klingon", "klingon"):
+            req = OpenAICreateSpeechRequest(input="Hello", language=language)
+            assert "Invalid language" in speech_server._validate_tts_request(req)
+
+    def test_load_supported_languages_default_when_no_config(self, speech_server):
+        """Empty/missing codec_language_id on a Qwen3-TTS model falls back to the default list."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={}))
+        )
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
+    def test_load_supported_languages_default_on_config_error(self, speech_server):
+        """If the model config cannot be read, fall back to the default list."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client = SimpleNamespace()  # no model_config -> AttributeError
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
+    def test_load_supported_languages_default_for_non_qwen(self, speech_server):
+        """Non-qwen3_tts model types get the default language set."""
+        speech_server._tts_model_type = None
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
     def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture):
         """Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only)."""
         speech_server.uploaded_speakers = {

tests/entrypoints/openai_api/test_video_server.py

@@ -13,12 +13,14 @@ import threading
 import time
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture
 
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -84,7 +86,8 @@ class BlockingVideoHandler:
         if self.stage_configs is None:
             self.stage_configs = stage_configs
 
-    async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
+    async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
+        del request, reference_id, reference_image, reference_video
         self.started.set()
         try:
             await asyncio.Future()
@@ -130,6 +133,30 @@ def _make_test_image_data_url(size=(64, 64)) -> str:
     return f"data:image/png;base64,{encoded}"
 
 
+def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes:
+    width, height = size
+    frames = np.zeros((num_frames, height, width, 3), dtype=np.uint8)
+    for idx in range(num_frames):
+        frames[idx, :, :, 0] = idx * 40
+        frames[idx, :, :, 1] = 128
+        frames[idx, :, :, 2] = 255 - idx * 40
+    return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"})
+
+
+def _make_test_video_data_url(size=(32, 24), num_frames=3) -> str:
+    encoded = base64.b64encode(_make_test_video_bytes(size, num_frames)).decode("utf-8")
+    return f"data:video/mp4;base64,{encoded}"
+
+
+def _cosmos3_stage_configs():
+    return [
+        SimpleNamespace(
+            stage_type="diffusion",
+            engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"),
+        )
+    ]
+
+
 def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
     deadline = time.time() + timeout_s
     last_payload = None
@@ -332,6 +359,104 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
     assert input_image.size == (40, 24)
 
 
+def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
+    video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "Continue this motion."},
+        files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    prompt = engine.captured_prompt
+    assert "multi_modal_data" in prompt
+    assert "video" in prompt["multi_modal_data"]
+    input_video = prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 3
+    assert all(isinstance(frame, Image.Image) for frame in input_video)
+    assert input_video[0].size == (32, 24)
+
+
+def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "Continue this motion.",
+            "video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_video = engine.captured_prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 2
+    assert input_video[0].size == (32, 24)
+
+
+def test_decode_video_bytes_can_keep_last_frames():
+    from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
+
+    frames = _decode_video_bytes(
+        _make_test_video_bytes((32, 24), num_frames=6),
+        source="input_reference",
+        max_frames=2,
+        keep="last",
+    )
+
+    assert len(frames) == 2
+    red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
+    assert red_means[0] > 100
+    assert red_means[1] > red_means[0]
+
+
+def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
+    request = VideoGenerationRequest(
+        prompt="Continue this motion.",
+        num_frames=189,
+        extra_params={"condition_frame_indexes_vision": [0, 2]},
+    )
+
+    spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs())
+    assert spec.max_frames == 9
+    assert spec.keep == "first"
+
+
+def test_cosmos3_reference_video_limit_preserves_action_frames():
+    request = VideoGenerationRequest(
+        prompt="Predict the action.",
+        num_frames=17,
+        extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16},
+    )
+
+    assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17
+
+
+def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
+    request = VideoGenerationRequest(
+        prompt="Continue this motion.",
+        num_frames=5,
+        extra_params={"condition_frame_indexes_vision": [0, 20]},
+    )
+
+    assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5
+
+
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     fps_values = []
 
@@ -794,7 +919,20 @@ def test_rejects_input_reference_and_image_reference_together(test_client):
         files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
     )
     assert response.status_code == 400
-    assert "either input_reference or image_reference" in response.json()["detail"].lower()
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
+
+
+def test_rejects_image_reference_and_video_reference_together(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "bad refs",
+            "image_reference": '{"image_url": "https://example.com/cat.png"}',
+            "video_reference": '{"video_url": "https://example.com/cat.mp4"}',
+        },
+    )
+    assert response.status_code == 400
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
 
 
 def test_invalid_seconds_returns_422(test_client):
@@ -856,6 +994,18 @@ def test_unsupported_image_reference_file_id_returns_400(test_client):
     assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
 
 
+def test_unsupported_video_reference_file_id_returns_400(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "unsupported ref",
+            "video_reference": '{"file_id": "file-123"}',
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Invalid video_reference: file_id is not supported yet."
+
+
 def test_invalid_uploaded_input_reference_returns_400(test_client):
     response = test_client.post(
         "/v1/videos",
@@ -863,7 +1013,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
         files={"input_reference": ("input.png", b"not-an-image", "image/png")},
     )
     assert response.status_code == 400
-    assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
+    assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image or video."
 
 
 def test_video_request_validation():
@@ -880,6 +1030,8 @@ def test_video_request_validation():
 
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", video_reference={"file_id": "file-1", "video_url": "https://example.com"})
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
     with pytest.raises(ValueError):
@@ -1265,6 +1417,23 @@ def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture):
     assert response.content == b"ref-video"
 
 
+def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture):
+    video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
+    _mock_encode_video_bytes(mocker, b"v2v-video-data")
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "Continue this motion."},
+        files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
+    )
+
+    assert response.status_code == 200
+    assert response.content == b"v2v-video-data"
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_video = engine.captured_prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 3
+    assert input_video[0].size == (32, 24)
+
+
 def test_sync_missing_handler_returns_503():
     app = FastAPI()
     app.include_router(router)
@@ -1297,7 +1466,7 @@ def test_sync_rejects_both_references(test_client):
         files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
     )
     assert response.status_code == 400
-    assert "either input_reference or image_reference" in response.json()["detail"].lower()
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
 
 
 def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):

tests/entrypoints/test_omni_sleep_mode.py

@@ -532,7 +532,10 @@ async def test_pure_diffusion_scenario(tp_size: int):
     try:
         await engine.sleep(level=1)
         await engine.wake_up()
-        async for _ in engine.generate("test", sampling_params=SamplingParams()):
+        async for _ in engine.generate(
+            "test",
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
+        ):
             pass
         logger.info("Pure diffusion OK (TP=%s)", tp_size)
     finally:

tests/test_config_factory.py

@@ -17,7 +17,6 @@ from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
     DeployConfig,
-    ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
@@ -351,93 +350,6 @@ class TestStageConfig:
         assert "parallel_config" not in omega_config.engine_args
 
 
-class TestModelPipeline:
-    """Tests for ModelPipeline class."""
-
-    def test_valid_linear_dag(self):
-        """Test validation of a valid linear DAG."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
-            StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-    def test_valid_branching_dag(self):
-        """Test validation of a valid branching DAG."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
-            StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-    def test_missing_entry_point(self):
-        """Test that missing entry point is detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
-            StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("entry point" in e.lower() for e in errors)
-
-    def test_missing_dependency(self):
-        """Test that missing stage reference is detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="output", input_sources=[99]),  # Invalid
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("non-existent" in e.lower() for e in errors)
-
-    def test_duplicate_stage_ids(self):
-        """Test that duplicate stage IDs are detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
-            StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]),  # Duplicate
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("duplicate" in e.lower() for e in errors)
-
-    def test_self_reference(self):
-        """Test that self-references are detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]),  # Self
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("itself" in e.lower() for e in errors)
-
-    def test_get_stage_by_id(self):
-        """Test getting stage by ID."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-
-        stage = pipeline.get_stage(1)
-        assert stage is not None
-        assert stage.model_stage == "talker"
-
-        missing = pipeline.get_stage(99)
-        assert missing is None
-
-    def test_empty_pipeline(self):
-        """Test validation of empty pipeline."""
-        pipeline = ModelPipeline(model_type="test", stages=[])
-        errors = pipeline.validate_pipeline()
-        assert any("no stages" in e.lower() for e in errors)
-
-
 class TestStageConfigFactory:
     """Tests for StageConfigFactory class."""
 
@@ -598,304 +510,6 @@ class TestStageResolutionHelpers:
         assert overridden == []
 
 
-class TestPipelineYamlParsing:
-    """Tests for pipeline YAML file parsing (@ZJY0516)."""
-
-    def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
-        """Test parsing the qwen3_omni_moe pipeline YAML."""
-        yaml_content = """\
-model_type: qwen3_omni_moe
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    model_stage: thinker
-    stage_type: llm
-    input_sources: []
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    hf_config_name: thinker_config
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-
-  - stage_id: 1
-    model_stage: talker
-    stage_type: llm
-    input_sources: [0]
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    hf_config_name: talker_config
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-
-  - stage_id: 2
-    model_stage: code2wav
-    stage_type: llm
-    input_sources: [1]
-    worker_type: generation
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-    hf_config_name: thinker_config
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-"""
-        yaml_file = tmp_path / "qwen3_omni_moe.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
-
-        assert pipeline.model_type == "qwen3_omni_moe"
-        assert len(pipeline.stages) == 3
-
-        # Stage 0: thinker
-        s0 = pipeline.stages[0]
-        assert s0.stage_id == 0
-        assert s0.model_stage == "thinker"
-        assert s0.stage_type == StageType.LLM
-        assert s0.input_sources == []
-        assert s0.worker_type == "ar"
-        assert s0.final_output is True
-        assert s0.final_output_type == "text"
-        assert s0.is_comprehension is True
-
-        # Stage 1: talker
-        s1 = pipeline.stages[1]
-        assert s1.stage_id == 1
-        assert s1.input_sources == [0]
-        assert s1.custom_process_input_func == (
-            "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
-        )
-        assert s1.final_output is False
-
-        # Stage 2: code2wav
-        s2 = pipeline.stages[2]
-        assert s2.stage_id == 2
-        assert s2.input_sources == [1]
-        assert s2.worker_type == "generation"
-        assert s2.final_output_type == "audio"
-
-    def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
-        """Test backward compatibility with engine_input_source field."""
-        yaml_content = """\
-model_type: legacy_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-  - stage_id: 1
-    model_stage: downstream
-    stage_type: llm
-    engine_input_source: [0]
-"""
-        yaml_file = tmp_path / "legacy.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
-        assert pipeline.stages[1].input_sources == [0]
-
-    def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
-        """Test parsing pipeline with optional connectors and edges."""
-        yaml_content = """\
-model_type: test_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-
-connectors:
-  type: ray
-
-edges:
-  - from: 0
-    to: 1
-"""
-        yaml_file = tmp_path / "with_connectors.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
-        assert pipeline.connectors == {"type": "ray"}
-        assert pipeline.edges == [{"from": 0, "to": 1}]
-
-    def test_parsed_pipeline_passes_validation(self, tmp_path):
-        """Test that a well-formed YAML produces a valid pipeline."""
-        yaml_content = """\
-model_type: valid_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-    final_output: true
-    final_output_type: text
-  - stage_id: 1
-    model_stage: next
-    stage_type: llm
-    input_sources: [0]
-"""
-        yaml_file = tmp_path / "valid.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected validation errors: {errors}"
-
-    def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
-        """Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
-        yaml_content = """\
-model_type: legacy_batch
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-    runtime:
-      devices: "0"
-      max_batch_size: 16
-    engine_args:
-      model_arch: SomeModel
-"""
-        yaml_file = tmp_path / "legacy_batch.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
-        s0 = pipeline.stages[0]
-        assert "max_batch_size" not in s0.yaml_runtime
-        assert s0.yaml_engine_args.get("max_num_seqs") == 16
-
-    def test_parse_diffusion_stage_type(self, tmp_path):
-        """Test parsing a diffusion stage type from YAML."""
-        yaml_content = """\
-model_type: diff_model
-
-stages:
-  - stage_id: 0
-    model_stage: dit
-    stage_type: diffusion
-    input_sources: []
-    final_output: true
-    final_output_type: image
-"""
-        yaml_file = tmp_path / "diffusion.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
-        assert pipeline.stages[0].stage_type == StageType.DIFFUSION
-
-    def test_parse_mimo_audio_yaml(self, tmp_path):
-        """Test parsing the MiMo Audio pipeline YAML."""
-        yaml_content = """\
-model_type: mimo_audio
-
-stages:
-  - stage_id: 0
-    model_stage: fused_thinker_talker
-    stage_type: llm
-    input_sources: []
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    model_stage: code2wav
-    stage_type: llm
-    input_sources: [0]
-    worker_type: generation
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
-    final_output: true
-    final_output_type: audio
-"""
-        yaml_file = tmp_path / "mimo_audio.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
-
-        assert pipeline.model_type == "mimo_audio"
-        assert len(pipeline.stages) == 2
-
-        # Stage 0: fused_thinker_talker
-        s0 = pipeline.stages[0]
-        assert s0.stage_id == 0
-        assert s0.model_stage == "fused_thinker_talker"
-        assert s0.input_sources == []
-        assert s0.worker_type == "ar"
-        assert s0.is_comprehension is True
-        assert s0.final_output is True
-        assert s0.final_output_type == "text"
-
-        # Stage 1: code2wav
-        s1 = pipeline.stages[1]
-        assert s1.stage_id == 1
-        assert s1.input_sources == [0]
-        assert s1.worker_type == "generation"
-        assert s1.custom_process_input_func == (
-            "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
-        )
-        assert s1.final_output is True
-        assert s1.final_output_type == "audio"
-
-        # Pipeline validation
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-
-class TestAsyncChunk:
-    """Tests for async_chunk pipeline flag."""
-
-    def test_model_pipeline_async_chunk_default(self):
-        """Test ModelPipeline defaults async_chunk to False."""
-        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        assert pipeline.async_chunk is False
-
-    def test_model_pipeline_async_chunk_set(self):
-        """Test ModelPipeline with async_chunk=True."""
-        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
-        pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
-        assert pipeline.async_chunk is True
-
-    def test_parse_async_chunk_from_yaml(self, tmp_path):
-        """Test that async_chunk is parsed from pipeline YAML."""
-        yaml_content = """\
-model_type: test_async
-async_chunk: true
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-"""
-        yaml_file = tmp_path / "async.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
-        assert pipeline.async_chunk is True
-
-    def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
-        """Test that missing async_chunk defaults to False."""
-        yaml_content = """\
-model_type: test_no_async
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-"""
-        yaml_file = tmp_path / "no_async.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
-        assert pipeline.async_chunk is False
-
-
 class TestPipelineDiscovery:
     """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
 

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -410,6 +410,8 @@ class TestAudioFeedback:
             model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})()
 
         t = FakeTalker()
+        t._decode_last_codes = torch.zeros(256, 8, dtype=torch.long)
+        t._decode_has_codes = torch.zeros(256, dtype=torch.bool)
         t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
         return t
 
@@ -427,6 +429,8 @@ class TestAudioFeedback:
         t = self._make_feedback_talker()
         audio_id = t._audio_continuation_id
         t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)}
+        t._decode_last_codes[1] = torch.zeros(8, dtype=torch.long)
+        t._decode_has_codes[1] = True
         input_ids = torch.tensor([1, audio_id, 3])
         hidden = torch.zeros(3, 16)
         result = t._apply_audio_feedback(hidden, input_ids)
@@ -440,7 +444,7 @@ class TestAudioFeedback:
         """Audio position without state should not be replaced."""
         t = self._make_feedback_talker()
         audio_id = t._audio_continuation_id
-        # No state for row 0
+        # No state for row 0 — _decode_has_codes[0] defaults to False
         input_ids = torch.tensor([audio_id])
         hidden = torch.zeros(1, 16)
         result = t._apply_audio_feedback(hidden, input_ids)
@@ -908,3 +912,300 @@ class TestPromptBuilder:
         # Should not contain ref_audio or ref_text token IDs
         assert 151703 not in ids  # <|ref_audio|>
         assert 151704 not in ids  # <|ref_text|>
+
+
+# ---- Regression: ramp-down done-before-emit off-by-one ----
+
+
+def _simulate_delay_state_machine(
+    sampled_codes_per_step: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    """Reproduce the delay-pattern state machine from sample().
+
+    Feeds ``sampled_codes_per_step`` (each [Q]) through the exact same
+    logic as HiggsAudioV3TalkerForConditionalGeneration.sample() and
+    returns the list of emitted frames (non -1 codes).
+    """
+    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+        BOC_ID,
+        EOC_ID,
+    )
+
+    num_codebooks = 8
+    bos = BOC_ID
+    eos_stream = EOC_ID
+
+    state: dict = {
+        "num_delay": 0,
+        "num_remaining_delays": None,
+        "audio_out_ids": None,
+        "last_codes": torch.full((num_codebooks,), bos, dtype=torch.long),
+        "should_terminate": False,
+    }
+
+    emitted: list[torch.Tensor] = []
+    for step_codes in sampled_codes_per_step:
+        if state["should_terminate"]:
+            break
+
+        num_delay = state["num_delay"]
+        num_remaining_delays = state["num_remaining_delays"]
+        this_codes = step_codes.clone()
+
+        if num_delay + 1 < num_codebooks:
+            this_codes[num_delay + 1 :] = bos
+            num_delay += 1
+
+        if num_remaining_delays is not None:
+            this_codes[: num_codebooks - num_remaining_delays] = eos_stream
+            num_remaining_delays -= 1
+        else:
+            eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
+            if eos_positions.numel() > 0:
+                last_eos_idx = int(eos_positions[-1].item())
+                this_codes[: last_eos_idx + 1] = eos_stream
+                num_remaining_delays = num_codebooks - last_eos_idx - 1
+
+        if num_remaining_delays is not None and num_remaining_delays <= 0:
+            state["num_delay"] = 0
+            state["num_remaining_delays"] = None
+            state["should_terminate"] = True
+            continue
+
+        state["num_delay"] = num_delay
+        state["num_remaining_delays"] = num_remaining_delays
+        state["last_codes"] = this_codes.clone()
+        emitted.append(this_codes)
+
+    return emitted
+
+
+class TestRampDownOffByOne:
+    """Regression tests for the done-before-emit off-by-one bug.
+
+    When EOC fires early, the ramp-down must emit enough frames so that
+    Stage1's _revert_delay_pattern receives T >= Q = 8.
+    """
+
+    def test_eoc_at_step0_terminates(self):
+        """EOC on cb0 at step 0 terminates; T < Q is handled by Stage1 defensive try/except."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        Q = 8
+        steps = []
+        for s in range(20):
+            codes = torch.randint(0, 1024, (Q,))
+            if s == 0:
+                codes[0] = EOC_ID
+            steps.append(codes)
+
+        emitted = _simulate_delay_state_machine(steps)
+        assert len(emitted) == Q - 1, f"EOC at step 0: expected {Q - 1} emitted frames, got {len(emitted)}"
+
+    def test_eoc_at_step1_emits_enough_frames(self):
+        """EOC on cb0 at step 1 must produce >= Q emitted frames."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        Q = 8
+        steps = []
+        for s in range(20):
+            codes = torch.randint(0, 1024, (Q,))
+            if s == 1:
+                codes[0] = EOC_ID
+            steps.append(codes)
+
+        emitted = _simulate_delay_state_machine(steps)
+        assert len(emitted) >= Q, f"EOC at step 1: expected >= {Q} emitted frames, got {len(emitted)}"
+
+    def test_eoc_after_full_rampup_emits_enough(self):
+        """EOC after all codebooks are active produces T >> Q."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        Q = 8
+        eoc_step = 10
+        steps = []
+        for s in range(30):
+            codes = torch.randint(0, 1024, (Q,))
+            if s == eoc_step:
+                codes[0] = EOC_ID
+            steps.append(codes)
+
+        emitted = _simulate_delay_state_machine(steps)
+        assert len(emitted) >= Q, f"EOC at step {eoc_step}: expected >= {Q} emitted frames, got {len(emitted)}"
+
+    def test_last_rampdown_frame_not_all_negative(self):
+        """The last emitted ramp-down frame must contain real codes, not -1."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        Q = 8
+        steps = []
+        for s in range(20):
+            codes = torch.randint(0, 1024, (Q,))
+            if s == 1:
+                codes[0] = EOC_ID
+            steps.append(codes)
+
+        emitted = _simulate_delay_state_machine(steps)
+        last = emitted[-1]
+        assert (last >= 0).all(), f"Last ramp-down frame contains negative values: {last.tolist()}"
+
+    def test_emitted_frames_revertible_by_stage1(self):
+        """Emitted frames from EOC-at-step-1 must be revertible by _revert_delay_pattern."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _revert_delay_pattern,
+        )
+
+        Q = 8
+        steps = []
+        for s in range(20):
+            codes = torch.randint(0, 1024, (Q,))
+            if s == 1:
+                codes[0] = EOC_ID
+            steps.append(codes)
+
+        emitted = _simulate_delay_state_machine(steps)
+        # Stack to [T, Q], transpose to [Q, T]
+        codes_qt = torch.stack(emitted, dim=0).t().contiguous()
+        assert codes_qt.shape[0] == Q
+        assert codes_qt.shape[1] >= Q
+
+        # Must not raise
+        reverted = _revert_delay_pattern(codes_qt)
+        assert reverted.shape[0] == Q
+        assert reverted.shape[1] == codes_qt.shape[1] - Q + 1
+
+
+class TestStage1DefensiveHandling:
+    """Stage1 must not crash the server when receiving T < Q frames."""
+
+    def test_revert_delay_pattern_raises_on_insufficient_frames(self):
+        """_revert_delay_pattern correctly raises ValueError for T < Q."""
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _revert_delay_pattern,
+        )
+
+        codes = torch.zeros(8, 7)  # T=7 < Q=8
+        with pytest.raises(ValueError, match="Not enough frames"):
+            _revert_delay_pattern(codes)
+
+    def test_talker2code2wav_survives_insufficient_frames(self):
+        """talker2code2wav must not crash when audio_codes has T < Q frames."""
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            talker2code2wav,
+        )
+
+        # Build a fake talker output with only 7 frames (T=7 < Q=8)
+        class FakeOutput:
+            def __init__(self):
+                self.multimodal_output = {
+                    "codes": {"audio": torch.zeros(7, 8, dtype=torch.long)},
+                }
+
+        class FakeTalkerOutput:
+            finished = True
+            outputs = [FakeOutput()]
+
+        # Must not raise — should return gracefully
+        result = talker2code2wav([FakeTalkerOutput()])
+        assert len(result) == 1
+        # Result is an OmniTokensPrompt with empty prompt_token_ids
+        out = result[0]
+        token_ids = out.prompt_token_ids if hasattr(out, "prompt_token_ids") else out.get("prompt_token_ids", [])
+        assert token_ids == []
+
+
+# ---- Regression: buffer growth + prefill guard ----
+
+
+class TestBufferGrowthPreservesState:
+    """Bug #1: buffer growth in _apply_audio_feedback must not wipe existing state."""
+
+    def _make_feedback_talker_with_buffers(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
+        torch.nn.init.ones_(embed.weight)
+
+        class FakeTalker:
+            num_codebooks = 8
+            codebook_size = 1026
+            multimodal_embedding = embed
+
+        t = FakeTalker()
+        t._decode_last_codes = torch.zeros(4, 8, dtype=torch.long)
+        t._decode_has_codes = torch.zeros(4, dtype=torch.bool)
+        t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
+        return t
+
+    def test_buffer_growth_preserves_existing_codes(self):
+        """Growing the buffer must copy old _decode_last_codes data."""
+        t = self._make_feedback_talker_with_buffers()
+        # Set state for slot 2
+        t._decode_last_codes[2] = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800])
+        t._decode_has_codes[2] = True
+
+        # Trigger growth by calling with bs=8 > current capacity 4
+        hidden = torch.randn(8, 16)
+        t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
+
+        # After growth, slot 2 data must still be present
+        assert t._decode_has_codes[2].item() is True
+        assert t._decode_last_codes[2, 0].item() == 100
+
+    def test_buffer_growth_new_slots_are_inactive(self):
+        """New slots from buffer growth must have has_codes=False."""
+        t = self._make_feedback_talker_with_buffers()
+        t._decode_has_codes[1] = True
+
+        hidden = torch.randn(8, 16)
+        t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
+
+        # Slots 4-7 (new) must be inactive
+        assert not t._decode_has_codes[4].item()
+        assert not t._decode_has_codes[7].item()
+        # Slot 1 (old) must still be active
+        assert t._decode_has_codes[1].item()
+
+
+class TestPrefillGuard:
+    """Bug #2: _apply_audio_feedback must NOT run during prefill.
+
+    The prefill detection uses attn_metadata.max_query_len (== 1 for decode,
+    > 1 for prefill). This correctly handles concurrent decode (N requests,
+    each contributing 1 token → numel=N but max_query_len=1).
+    """
+
+    @staticmethod
+    def _is_prefill_via_max_query_len(max_query_len: int | None, input_ids: torch.Tensor) -> bool:
+        """Mirror the actual detection logic in forward()."""
+        if max_query_len is not None:
+            return int(max_query_len) > 1
+        return input_ids is not None and int(input_ids.numel()) > 1
+
+    def test_prefill_detected_via_max_query_len(self):
+        """Prefill: max_query_len > 1 → is_prefill=True."""
+        input_ids = torch.tensor([1, 2, 3, 4, 5])
+        assert self._is_prefill_via_max_query_len(5, input_ids) is True
+
+    def test_single_decode_detected_via_max_query_len(self):
+        """Single-request decode: max_query_len=1 → is_prefill=False."""
+        input_ids = torch.tensor([42])
+        assert self._is_prefill_via_max_query_len(1, input_ids) is False
+
+    def test_concurrent_decode_not_misclassified(self):
+        """Multi-request decode: numel=4 but max_query_len=1 → is_prefill=False.
+
+        This is the regression caught by reviewer: with c=4 concurrent decode
+        requests, input_ids.numel()==4 but each request contributes exactly 1
+        token, so max_query_len==1 and audio feedback must still apply.
+        """
+        input_ids = torch.tensor([42, 43, 44, 45])  # 4 concurrent decode
+        assert self._is_prefill_via_max_query_len(1, input_ids) is False
+
+    def test_fallback_when_max_query_len_unavailable(self):
+        """When attn_metadata is unavailable, fall back to numel heuristic."""
+        prefill_ids = torch.tensor([1, 2, 3])
+        assert self._is_prefill_via_max_query_len(None, prefill_ids) is True
+        decode_ids = torch.tensor([42])
+        assert self._is_prefill_via_max_query_len(None, decode_ids) is False

vllm_omni/config/__init__.py

@@ -6,7 +6,6 @@ from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.config.stage_config import (
     DeployConfig,
-    ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
@@ -30,7 +29,6 @@ __all__ = [
     "LoRAConfig",
     "StageConfig",
     "StageConfigFactory",
-    "ModelPipeline",
     "StageType",
     "StageExecutionType",
     "StagePipelineConfig",

vllm_omni/config/stage_config.py

@@ -20,13 +20,6 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
-_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
-
-
-def get_pipeline_path(model_dir: str, filename: str) -> Path:
-    return _MODELS_DIR / model_dir / filename
-
-
 logger = init_logger(__name__)
 
 
@@ -1061,77 +1054,6 @@ class StageConfig:
         return create_config(config_dict)
 
 
-@dataclass
-class ModelPipeline:
-    """Complete pipeline definition for a multi-stage model (legacy).
-
-    TODO(@lishunyang12): remove once all models migrate to PipelineConfig.
-    """
-
-    model_type: str
-    stages: list[StageConfig]
-
-    # Pipeline-wide behavior flags
-    async_chunk: bool = False
-
-    # Optional distributed configuration
-    connectors: dict[str, Any] | None = None
-    edges: list[dict[str, Any]] | None = None
-
-    def get_stage(self, stage_id: int) -> StageConfig | None:
-        """Look up a stage by its ID.
-
-        Args:
-            stage_id: The stage ID to search for.
-
-        Returns:
-            The matching StageConfig, or None if not found.
-        """
-        for stage in self.stages:
-            if stage.stage_id == stage_id:
-                return stage
-        return None
-
-    def validate_pipeline(self) -> list[str]:
-        """Validate pipeline topology at model integration time (not runtime).
-
-        Checks:
-        - All stage IDs are unique
-        - All input_sources reference valid stage IDs
-        - At least one entry point (stage with empty input_sources)
-
-        Returns:
-            List of validation error messages. Empty list if valid.
-        """
-        errors: list[str] = []
-
-        if not self.stages:
-            errors.append("Topology has no stages defined")
-            return errors
-
-        # Check for unique stage IDs
-        stage_ids = [s.stage_id for s in self.stages]
-        if len(stage_ids) != len(set(stage_ids)):
-            errors.append("Duplicate stage IDs found")
-
-        stage_id_set = set(stage_ids)
-
-        # Check input_sources reference valid stages
-        for stage in self.stages:
-            for source_id in stage.input_sources:
-                if source_id not in stage_id_set:
-                    errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
-                if source_id == stage.stage_id:
-                    errors.append(f"Stage {stage.stage_id} references itself as input source")
-
-        # Check for at least one entry point
-        entry_points = [s for s in self.stages if not s.input_sources]
-        if not entry_points:
-            errors.append("No entry point found (stage with empty input_sources)")
-
-        return errors
-
-
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -1206,33 +1128,9 @@ class StageConfigFactory:
                             continue
                     return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
 
-        # --- Legacy path: load from pipeline YAML ---
-        pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
-
-        if pipeline is None:
-            return None
-
-        errors = pipeline.validate_pipeline()
-        if errors:
-            logger.warning(f"Pipeline validation warnings for {model}: {errors}")
-
-        # Materialize the resolved pipeline-wide async_chunk value into every
-        # stage so build_engine_args_dict() can inject the stage connector
-        # spec and explicit False overrides are preserved.
-        resolved_async_chunk = cli_overrides.get("async_chunk")
-        if resolved_async_chunk is None:
-            resolved_async_chunk = bool(pipeline.async_chunk)
-        for stage in pipeline.stages:
-            stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk)
-
-        # Apply CLI overrides
-        result: list[StageConfig] = []
-        for stage in pipeline.stages:
-            # Merge global CLI overrides
-            stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
-            result.append(stage)
-
-        return result
+        # Not in the pipeline registry — let the caller fall back to the
+        # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
+        return None
 
     @classmethod
     def _create_from_registry(
@@ -1355,193 +1253,6 @@ class StageConfigFactory:
 
         return [config_dict]
 
-    @classmethod
-    def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
-        """Load a legacy ``pipeline.yaml`` for the model.
-
-        Searches ``model_executor/models/<dir>/pipeline.yaml`` by trying
-        (a) the raw ``model_type`` as the directory name, then
-        (b) ``model_type`` with hyphens replaced by underscores,
-        and finally (c) scanning every ``pipeline.yaml`` for one that
-        declares a matching ``model_type`` or ``hf_architectures``.
-
-        Returns None if no pipeline.yaml is found — caller handles the
-        ``resolve_model_config_path`` fallback via stage_configs/ YAMLs.
-        """
-        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
-        if model_type is None:
-            return None
-
-        # Direct lookups by convention
-        candidates = [model_type, model_type.replace("-", "_")]
-        for dir_name in candidates:
-            pipeline_path = get_pipeline_path(dir_name, "pipeline.yaml")
-            if pipeline_path.exists():
-                return cls._parse_pipeline_yaml(pipeline_path, model_type)
-
-        # Scan fallback: read every pipeline.yaml and match on declared fields
-        hf_archs = set(getattr(hf_config, "architectures", []) or []) if hf_config else set()
-        if _MODELS_DIR.exists():
-            for subdir in sorted(_MODELS_DIR.iterdir()):
-                if not subdir.is_dir():
-                    continue
-                pipeline_path = subdir / "pipeline.yaml"
-                if not pipeline_path.exists():
-                    continue
-                try:
-                    cfg = load_yaml_config(pipeline_path)
-                except Exception as exc:
-                    logger.debug("Skip %s: %s", pipeline_path, exc)
-                    continue
-                declared_type = getattr(cfg, "model_type", None)
-                declared_archs = set(getattr(cfg, "hf_architectures", None) or [])
-                if declared_type == model_type or (hf_archs and hf_archs.intersection(declared_archs)):
-                    return cls._parse_pipeline_yaml(pipeline_path, declared_type or model_type)
-
-        logger.debug("No pipeline.yaml found for model_type %s (archs=%s)", model_type, sorted(hf_archs))
-        return None
-
-    # Keys consumed as explicit StageConfig fields — everything else is
-    # passed through via yaml_extras.
-    _KNOWN_STAGE_KEYS: set[str] = {
-        "stage_id",
-        "model_stage",
-        "stage_type",
-        "input_sources",
-        "engine_input_source",
-        "custom_process_input_func",
-        "final_output",
-        "final_output_type",
-        "worker_type",
-        "scheduler_cls",
-        "hf_config_name",
-        "is_comprehension",
-        "engine_args",
-        "runtime",
-    }
-
-    @classmethod
-    def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
-        """Parse a pipeline YAML file.
-
-        Args:
-            path: Path to the YAML file.
-            model_type: Model type identifier.
-
-        Returns:
-            ModelPipeline object.
-        """
-        config_data = load_yaml_config(path)
-
-        stages: list[StageConfig] = []
-        for stage_data in config_data.stages:
-            # Use .get() for optional fields — idiomatic for OmegaConf DictConfig
-            stage_type_str = stage_data.get("stage_type", "llm")
-            stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
-
-            # Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
-            input_sources = stage_data.get("input_sources", None)
-            if input_sources is None:
-                input_sources = stage_data.get("engine_input_source", [])
-            if input_sources is None:
-                input_sources = []
-            input_sources = list(input_sources)
-
-            # Extract per-stage engine_args and runtime dicts
-            raw_ea = stage_data.get("engine_args", None)
-            yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
-            raw_rt = stage_data.get("runtime", None)
-            yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
-
-            # Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
-            if "max_batch_size" in yaml_runtime:
-                mbs = yaml_runtime.pop("max_batch_size")
-                yaml_engine_args.setdefault("max_num_seqs", int(mbs))
-                logger.debug(
-                    "Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
-                    stage_data.get("stage_id", "?"),
-                    mbs,
-                )
-
-            # Topology-level fields that also live inside engine_args in legacy
-            # YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
-            worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
-            scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
-            if scheduler_cls:
-                async_sched = yaml_engine_args.get("async_scheduling")
-                if async_sched is not None:
-                    logger.warning(
-                        "Stage %s: async_scheduling=%r and scheduler_cls=%r "
-                        "should not be set together. scheduler_cls will take "
-                        "precedence for which scheduler is used.",
-                        stage_data.stage_id,
-                        async_sched,
-                        scheduler_cls,
-                    )
-                else:
-                    logger.warning(
-                        "Stage %s: scheduler_cls=%r is deprecated. Use async_scheduling instead.",
-                        stage_data.stage_id,
-                        scheduler_cls,
-                    )
-            hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
-            model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
-
-            # Collect pass-through fields (default_sampling_params,
-            # output_connectors, input_connectors, tts_args, etc.)
-            yaml_extras: dict[str, Any] = {}
-            for key in stage_data:
-                if key not in cls._KNOWN_STAGE_KEYS:
-                    val = stage_data[key]
-                    try:
-                        yaml_extras[key] = to_dict(val)
-                    except ValueError:
-                        yaml_extras[key] = val
-
-            stage = StageConfig(
-                stage_id=stage_data.stage_id,
-                model_stage=model_stage or "",
-                stage_type=stage_type,
-                input_sources=input_sources,
-                custom_process_input_func=stage_data.get("custom_process_input_func", None),
-                final_output=stage_data.get("final_output", False),
-                final_output_type=stage_data.get("final_output_type", None),
-                worker_type=worker_type,
-                scheduler_cls=scheduler_cls,
-                hf_config_name=hf_config_name,
-                is_comprehension=stage_data.get("is_comprehension", False),
-                yaml_engine_args=yaml_engine_args,
-                yaml_runtime=yaml_runtime,
-                yaml_extras=yaml_extras,
-            )
-            stages.append(stage)
-
-        # Get pipeline-wide flags
-        async_chunk = config_data.get("async_chunk", False)
-
-        # Get optional connector config — check both top-level and nested
-        # under ``runtime`` (legacy stage_configs format).
-        connectors = None
-        edges = None
-        if hasattr(config_data, "connectors"):
-            connectors = to_dict(config_data.connectors)
-        if hasattr(config_data, "edges"):
-            edges = to_dict(config_data.edges)
-        if hasattr(config_data, "runtime") and config_data.runtime is not None:
-            top_runtime = config_data.runtime
-            if connectors is None and hasattr(top_runtime, "connectors"):
-                connectors = to_dict(top_runtime.connectors)
-            if edges is None and hasattr(top_runtime, "edges"):
-                edges = to_dict(top_runtime.edges)
-
-        return ModelPipeline(
-            model_type=getattr(config_data, "model_type", model_type),
-            stages=stages,
-            async_chunk=async_chunk,
-            connectors=connectors,
-            edges=edges,
-        )
-
     @classmethod
     def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
         """Auto-detect model_type from model directory.

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -81,7 +81,7 @@ class DefaultAdapter:
 
         loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
         # load_model will handle dtypes / device placement, put in .eval() mode
-        return loader.load_model(od_config=od_config, load_device=device)
+        return loader.load_model(load_device=device)
 
     @staticmethod
     def get_transformer(pipeline: Any) -> tuple[Any, str]:

vllm_omni/diffusion/data.py

@@ -610,9 +610,6 @@ class OmniDiffusionConfig:
     # Step mode settings
     step_execution: bool = False
 
-    # sleep mode
-    enable_sleep_mode: bool = False
-
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 import dataclasses
 import glob
 import os
@@ -30,8 +31,9 @@ from vllm.transformers_utils.repo_utils import file_exists
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import set_default_torch_dtype
 
+from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
@@ -100,18 +102,20 @@ class DiffusersPipelineLoader:
     counter_before_loading_weights: float = 0.0
     counter_after_loading_weights: float = 0.0
 
-    def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None):
+    def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig):
         self.load_config = load_config
         self.od_config = od_config
+        self.quant_config = od_config.quantization_config
+        self.parallel_config = od_config.parallel_config
 
     def _prepare_weights(
         self,
-        model_name_or_path: Path,
+        model_name_or_path: Path | str,
         subfolder: str | None,
         revision: str | None,
         fall_back_to_pt: bool,
         allow_patterns_overrides: list[str] | None,
-    ) -> tuple[str, list[str], bool]:
+    ) -> tuple[Path | str, list[str], bool]:
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
@@ -208,15 +212,13 @@ class DiffusersPipelineLoader:
             source.allow_patterns_overrides,
         )
 
-        od_config = self.od_config
         use_multithread = (
             use_safetensors
-            and od_config is not None
-            and getattr(od_config, "enable_multithread_weight_load", False)
+            and getattr(self.od_config, "enable_multithread_weight_load", False)
             and self.load_config.safetensors_load_strategy != "torchao"
         )
         if use_multithread:
-            num_threads = getattr(od_config, "num_weight_load_threads", 4)
+            num_threads = getattr(self.od_config, "num_weight_load_threads", 4)
             # Keep deterministic shard order before passing to vLLM helper.
             sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key)
             weights_iterator = multi_thread_safetensors_weights_iterator(
@@ -242,10 +244,7 @@ class DiffusersPipelineLoader:
         return prefixed_weights_iterator
 
     def _get_source_quant_config(self, source: "ComponentSource") -> object | None:
-        if self.od_config is None:
-            return None
-
-        quant_config = self.od_config.quantization_config
+        quant_config = self.quant_config
         if hasattr(quant_config, "resolve"):
             return quant_config.resolve(source.prefix.rstrip("."))
         return quant_config
@@ -308,7 +307,6 @@ class DiffusersPipelineLoader:
 
     def load_model(
         self,
-        od_config: OmniDiffusionConfig,
         load_device: str,
         load_format: str | None = "default",
         custom_pipeline_name: str | type[nn.Module] | None = None,
@@ -317,14 +315,13 @@ class DiffusersPipelineLoader:
         """Load a model with the given configurations."""
         if load_format is None:
             load_format = "default"
-        self.od_config = od_config
         # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
         # weights are already quantized in the checkpoint — load directly on CPU.
         # For online quantization, load on device so quantization can run on accelerator,
         # then move back to CPU afterward.
         offload_after_quant = False
-        if load_device == "cpu" and od_config.quantization_config is not None:
-            quant_cfg = od_config.quantization_config
+        if load_device == "cpu" and self.quant_config is not None and device is not None:
+            quant_cfg = self.quant_config
             is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
                 quant_cfg, "is_checkpoint_quantized", False
             )
@@ -339,48 +336,21 @@ class DiffusersPipelineLoader:
                 logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU")
 
         target_device = torch.device(load_device)
-        with set_default_torch_dtype(od_config.dtype):
-            if od_config.parallel_config.use_hsdp:
+        with set_default_torch_dtype(self.od_config.dtype):
+            if self.parallel_config.use_hsdp:
                 model = self._load_model_with_hsdp(
-                    od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
+                    target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
                 )
             else:
-                if load_format == "custom_pipeline":
-                    # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)`
-                    # internally. If we construct them under `with target_device:` (CUDA),
-                    # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc`
-                    # via the driver API and BYPASSES PyTorch's caching allocator.
-                    # That makes those bytes invisible to CuMemAllocator, so `sleep()`
-                    # cannot offload/unmap them and GPU memory stays pinned.
-                    #
-                    # Fix: build the custom pipeline on CPU first (no default device
-                    # context), then explicitly move it to the target device. The
-                    # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)`
-                    # + `copy_`, which goes through the caching allocator and is fully
-                    # tracked by CuMemAllocator.
-                    from vllm_omni.diffusion.config import set_current_diffusion_config
-
-                    model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
-                    with set_current_diffusion_config(od_config):
-                        model = model_cls(od_config=od_config)
-                    if target_device.type != "cpu":
-                        model.to(target_device)
-                else:
-                    with target_device:
-                        if load_format == "default":
-                            model = initialize_model(od_config)
-                        elif load_format == "diffusers":
-                            model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
-                        else:
-                            raise ValueError(f"Unknown load_format: {load_format}")
+                model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
                 logger.debug("Loading weights on %s ...", load_device)
                 if load_format == "diffusers":
                     # DiffusersAdapterPipeline.load_weights() calls
                     # DiffusionPipeline.from_pretrained() internally; it does
                     # not use our native customized pipeline classes.
                     cast(DiffusersAdapterPipeline, model).load_weights()
-                elif self._is_gguf_quantization(od_config):
-                    self._load_weights_with_gguf(model, od_config)
+                elif self._is_gguf_quantization():
+                    self._load_weights_with_gguf(model)
                 else:
                     # Quantization does not happen in `load_weights` but after it
                     self.load_weights(model)
@@ -403,7 +373,7 @@ class DiffusersPipelineLoader:
         """
         for _, module in model.named_modules():
             quant_method = getattr(module, "quant_method", None)
-            if isinstance(quant_method, QuantizeMethodBase):
+            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
                 # Move module to target device for processing if needed
                 module_device = next(module.parameters(), None)
                 if module_device is not None:
@@ -433,8 +403,13 @@ class DiffusersPipelineLoader:
         # that have loaded weights tracking currently.
         if loaded_weights is not None:
             weights_not_loaded = weights_to_load - loaded_weights
-            # NOTE: if the model is quantized, ignore not_loaded check for scale weights
-            weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")}
+            # NOTE: if the model is quantized, ignore not_loaded check for scale
+            # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
+            # static activation `input_scale`, which the quant method may
+            # fold/track differently than plain parameters.
+            weights_scale_not_loaded = {
+                name for name in weights_not_loaded if name.endswith(("weight_scale", "input_scale"))
+            }
             weights_not_loaded = weights_not_loaded - weights_scale_not_loaded
             if weights_not_loaded:
                 self._check_unloaded_weights(weights_not_loaded)
@@ -476,9 +451,11 @@ class DiffusersPipelineLoader:
         are logged as a warning.  Any *other* missing weight raises
         ``ValueError`` regardless of quantization.
         """
-        od_config = getattr(self, "od_config", None)
-        if od_config is None or od_config.quantization_config is None:
-            raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
+        if self.quant_config is None:
+            raise ValueError(
+                "The quantization config is None, and the following weights "
+                f"were not initialized from checkpoint: {weights_not_loaded}"
+            )
 
         expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)}
         unexpected_missing = weights_not_loaded - expected_missing
@@ -491,39 +468,33 @@ class DiffusersPipelineLoader:
         if unexpected_missing:
             raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
 
-    def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool:
-        quant_config = od_config.quantization_config
-        if quant_config is None:
-            return False
+    def _get_gguf_model_from_config(self) -> str | None:
+        """Extract gguf_model from quant_config (handles both dict and object)."""
+        if self.quant_config is None:
+            return None
+        if isinstance(self.quant_config, dict):
+            return self.quant_config.get("gguf_model")
+        return getattr(self.quant_config, "gguf_model", None)
 
-        # New API: DiffusionGGUFConfig (QuantizationConfig subclass)
+    def _is_gguf_quantization(self) -> bool:
+        """Check whether or not this pipeline loader is pointing at a GGUF config."""
         from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
 
-        if isinstance(quant_config, DiffusionGGUFConfig):
-            if quant_config.gguf_model is None:
+        maybe_gguf_model = self._get_gguf_model_from_config()
+        is_gguf = (
+            isinstance(self.quant_config, DiffusionGGUFConfig)
+            or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
+            or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
+        )
+        if is_gguf:
+            if not maybe_gguf_model:
                 raise ValueError("GGUF quantization requires gguf_model")
             return True
 
-        # Dict-style config: {"method": "gguf", "gguf_model": "..."}
-        if isinstance(quant_config, dict):
-            if quant_config.get("method") == "gguf":
-                if not quant_config.get("gguf_model"):
-                    raise ValueError("GGUF quantization requires gguf_model")
-                return True
+        if isinstance(self.quant_config, dict):
             return False
 
-        # Check by name for any config that reports as "gguf"
-        if hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf":
-            gguf_model = getattr(quant_config, "gguf_model", None)
-            if gguf_model is None:
-                raise ValueError("GGUF quantization requires gguf_model")
-            return True
-
-        # Fallback: object with gguf_model attribute but no get_name
-        if hasattr(quant_config, "gguf_model") and quant_config.gguf_model:
-            return True
-
-        return False
+        return bool(maybe_gguf_model)
 
     def _is_transformer_source(self, source: "ComponentSource") -> bool:
         if source.subfolder == "transformer":
@@ -566,25 +537,23 @@ class DiffusersPipelineLoader:
         self,
         source: "ComponentSource",
         model: nn.Module,
-        od_config: OmniDiffusionConfig,
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        quant_config = od_config.quantization_config
-        gguf_model = getattr(quant_config, "gguf_model", None)
+        gguf_model = self._get_gguf_model_from_config()
         if gguf_model is None:
             raise ValueError("GGUF quantization requires quantization_config.gguf_model")
-        gguf_file = self._resolve_gguf_model_path(gguf_model, od_config.revision)
-        adapter = get_gguf_adapter(gguf_file, model, source, od_config)
+        gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
+        adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
         weights_iter = adapter.weights_iterator()
         return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
 
-    def _load_weights_with_gguf(self, model: nn.Module, od_config: OmniDiffusionConfig) -> set[str]:
+    def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
         sources = self._get_weight_sources(model)
         loaded: set[str] = set()
         loadable_names: set[str] | None = None
 
         for source in sources:
             if self._is_transformer_source(source):
-                loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model, od_config))
+                loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
 
                 # GGUF checkpoints can be transformer-only or partially quantized.
                 # Only fall back to HF if this source still has missing loadable weights.
@@ -609,9 +578,45 @@ class DiffusersPipelineLoader:
             raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
         return loaded
 
+    def _init_from_load_format(
+        self,
+        load_format: str,
+        target_device: torch.device,
+        custom_pipeline_name: str | type[nn.Module] | None = None,
+        is_hsdp: bool = False,
+    ) -> nn.Module:
+        """Initialize the model from a specified load format."""
+        if load_format == "custom_pipeline":
+            # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)`
+            # internally. If we construct them under `with target_device:` (CUDA),
+            # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc`
+            # via the driver API and BYPASSES PyTorch's caching allocator.
+            # That makes those bytes invisible to CuMemAllocator, so `sleep()`
+            # cannot offload/unmap them and GPU memory stays pinned.
+            #
+            # Fix: build the custom pipeline on CPU first (no default device
+            # context), then explicitly move it to the target device. The
+            # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)`
+            # + `copy_`, which goes through the caching allocator and is fully
+            # tracked by CuMemAllocator.
+            model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
+            with set_current_diffusion_config(self.od_config):
+                model = model_cls(od_config=self.od_config)
+            if not is_hsdp and target_device.type != "cpu":
+                model.to(target_device)
+        else:
+            device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
+            with device_ctx:
+                if load_format == "default":
+                    model = initialize_model(self.od_config)
+                elif load_format == "diffusers":
+                    model = DiffusersAdapterPipeline(od_config=self.od_config, device=target_device)
+                else:
+                    raise ValueError(f"Unknown load_format: {load_format}")
+        return model
+
     def _load_model_with_hsdp(
         self,
-        od_config: OmniDiffusionConfig,
         target_device: torch.device,
         load_format: str = "default",
         custom_pipeline_name: str | type[nn.Module] | None = None,
@@ -624,14 +629,11 @@ class DiffusersPipelineLoader:
         Approach: Load weights first using model's load_weights (handles QKV fusion etc.),
         then apply HSDP sharding to redistribute weights across GPUs.
         """
-        from vllm_omni.diffusion.distributed.hsdp import apply_hsdp_to_model
-
-        parallel_config = od_config.parallel_config
         hsdp_config = HSDPInferenceConfig(
             enabled=True,
-            hsdp_replicate_size=parallel_config.hsdp_replicate_size,
-            hsdp_shard_size=parallel_config.hsdp_shard_size,
-            param_dtype=od_config.dtype,
+            hsdp_replicate_size=self.parallel_config.hsdp_replicate_size,
+            hsdp_shard_size=self.parallel_config.hsdp_shard_size,
+            param_dtype=self.od_config.dtype,
         )
 
         # Initialize model WITHOUT device context (weights start on CPU).
@@ -639,14 +641,9 @@ class DiffusersPipelineLoader:
         # directly on GPU, HSDP needs weights on CPU first so they can be redistributed
         # across GPUs by apply_hsdp_to_model. The model's load_weights handles weight
         # mapping (QKV fusion, etc.).
-        if load_format == "default":
-            model = initialize_model(od_config)
-        elif load_format == "custom_pipeline":
-            model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
-            from vllm_omni.diffusion.config import set_current_diffusion_config
-
-            with set_current_diffusion_config(od_config):
-                model = model_cls(od_config=od_config)
+        if load_format == "diffusers":
+            raise ValueError("HSDP is not supported with the diffusers adapter load format")
+        model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
         # Collect all transformers to shard (some models have transformer_2 for MoE)

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Cosmos3 text/image-to-video and text-to-image pipeline for vllm-omni.
+"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni.
 
-Single pipeline class supports T2V, I2V, and T2I; the mode is selected at
+Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at
 runtime by:
 
 * ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image).
@@ -10,6 +10,8 @@ runtime by:
   (text-to-video).
 * ``multi_modal_data['image']`` present on the prompt:  **I2V**
   (handled by :func:`get_cosmos3_pre_process_func`)
+* ``multi_modal_data['video']`` present on the prompt with no action mode:
+  **V2V** (handled by :func:`get_cosmos3_pre_process_func`)
 
 """
 
@@ -40,7 +42,10 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import (
+    ReferenceVideoDecodeSpec,
+    SupportImageInput,
+)
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -63,6 +68,17 @@ from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
 
 logger = init_logger(__name__)
 
+COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)
+COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first"
+# Mirrors the WAN VAE's temporal compression. Authoritative value is
+# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so
+# off-line / API code that runs before the pipeline is constructed can compute
+# pixel-frame budgets without instantiating the VAE.
+COSMOS3_VAE_TEMPORAL_COMPRESSION = 4
+COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
+    max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
+)
+COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0
 COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
 COSMOS3_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution."
 COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution."
@@ -91,16 +107,55 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
 COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
 
 
+def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
+    """Normalize Cosmos3 vision-conditioning latent frame indexes."""
+    if value is None:
+        return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION
+    if isinstance(value, str):
+        value = [item.strip() for item in value.split(",") if item.strip()]
+    elif isinstance(value, int):
+        value = [value]
+
+    if not isinstance(value, Iterable):
+        raise TypeError(
+            "Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, "
+            f"or iterable of ints; got {type(value)!r}."
+        )
+
+    indexes = tuple(sorted({int(index) for index in value}))
+    if not indexes:
+        raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.")
+    if any(index < 0 for index in indexes):
+        raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.")
+    return indexes
+
+
+def _condition_pixel_frame_count(
+    condition_frame_indexes_vision: Iterable[int],
+    temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
+) -> int:
+    return max(condition_frame_indexes_vision) * int(temporal_compression) + 1
+
+
+def _normalize_condition_video_keep(value: Any) -> str:
+    keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower()
+    if keep not in {"first", "last"}:
+        raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.")
+    return keep
+
+
 # ---------------------------------------------------------------------------
 # Post-process function (registered in registry.py)
 # ---------------------------------------------------------------------------
 def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
-    """Pre-process function for both T2V and I2V.
+    """Pre-process function for T2V, I2V, V2V, and action inputs.
 
     For T2V (no image in ``multi_modal_data``), the request is returned
     unchanged after the optional guardrails check.  For I2V (image present),
     the conditioning image is loaded, aspect-resized + center-cropped, and
     stored back on the prompt as ``additional_information.preprocessed_image``.
+    For V2V (video present without action mode), source frames are cropped to
+    the target size and stored as ``additional_information.preprocessed_video``.
     """
     from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled
 
@@ -137,7 +192,27 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
             return PIL.Image.open(value).convert("RGB")
         if isinstance(value, PIL.Image.Image):
             return value.convert("RGB")
-        raise TypeError(f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}.")
+        if isinstance(value, np.ndarray):
+            array = value
+            if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
+                array = np.transpose(array, (1, 2, 0))
+            if np.issubdtype(array.dtype, np.floating):
+                if array.min() < 0.0 or array.max() > 1.0:
+                    array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5
+                array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+            return PIL.Image.fromarray(array).convert("RGB")
+        if isinstance(value, torch.Tensor):
+            tensor = value.detach().cpu()
+            if tensor.ndim == 3 and tensor.shape[0] in (3, 4):
+                tensor = tensor.permute(1, 2, 0)
+            if tensor.is_floating_point():
+                if tensor.min().item() < 0.0 or tensor.max().item() > 1.0:
+                    tensor = tensor.clamp(-1.0, 1.0) * 0.5 + 0.5
+                tensor = (tensor.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
+            return PIL.Image.fromarray(tensor.numpy()).convert("RGB")
+        raise TypeError(
+            f"Cosmos3 preprocessing expected PIL image, numpy array, torch tensor, or path, got {type(value)!r}."
+        )
 
     def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image:
         scale = min(target_w / image.width, target_h / image.height, 1.0)
@@ -170,6 +245,61 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
         processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames]
         return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
 
+    def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
+        scale = max(target_w / image.width, target_h / image.height)
+        resize_w = int(np.ceil(scale * image.width))
+        resize_h = int(np.ceil(scale * image.height))
+        image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
+        left = (resize_w - target_w) // 2
+        top = (resize_h - target_h) // 2
+        image = image.crop((left, top, left + target_w, top + target_h))
+        return video_processor.preprocess(image, height=target_h, width=target_w)
+
+    def _video_payload_to_frames(video: Any) -> list[Any]:
+        if isinstance(video, list):
+            return video
+        if isinstance(video, torch.Tensor):
+            tensor = video.detach().cpu()
+            if tensor.ndim == 5:
+                if tensor.shape[0] != 1:
+                    raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
+                tensor = tensor[0]
+            if tensor.ndim == 4 and tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4):
+                return [tensor[:, i] for i in range(tensor.shape[1])]
+            if tensor.ndim == 4 and tensor.shape[-1] in (3, 4):
+                return [tensor[i] for i in range(tensor.shape[0])]
+        if isinstance(video, np.ndarray):
+            array = video
+            if array.ndim == 5:
+                if array.shape[0] != 1:
+                    raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
+                array = array[0]
+            if array.ndim == 4 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
+                return [array[:, i] for i in range(array.shape[1])]
+            if array.ndim == 4 and array.shape[-1] in (3, 4):
+                return [array[i] for i in range(array.shape[0])]
+        raise TypeError("Cosmos3 video input must be a non-empty list of frames or a single video tensor/array.")
+
+    def _select_video_frames(frames: list[Any], max_frames: int, keep: str) -> list[Any]:
+        if not frames:
+            raise ValueError("Cosmos3 video input must contain at least one frame.")
+        if keep == "last":
+            return frames[-max_frames:]
+        return frames[:max_frames]
+
+    def _preprocess_condition_video(
+        frames: list[Any],
+        target_h: int,
+        target_w: int,
+        max_frames: int,
+        keep: str,
+    ) -> torch.Tensor:
+        selected = _select_video_frames(frames, max_frames, keep)
+        processed = [
+            _preprocess_condition_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in selected
+        ]
+        return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
+
     def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
         action_mode = _request_action_mode(request)
         if is_guardrails_enabled(od_config, request.sampling_params):
@@ -183,16 +313,23 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
             multi_modal_data = prompt.get("multi_modal_data", {}) or {}
             raw_image = multi_modal_data.get("image")
             raw_video = multi_modal_data.get("video")
-            if raw_image is None and not (action_mode is not None and raw_video is not None):
+            if raw_image is None and raw_video is None:
                 continue
+            if raw_image is not None and raw_video is not None and action_mode is None:
+                raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
 
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
+            raw_video_frames: list[Any] | None = None
+            if raw_video is not None:
+                raw_video_frames = _video_payload_to_frames(raw_video)
+                if not raw_video_frames:
+                    raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.")
+
             if raw_image is None:
-                if not isinstance(raw_video, list) or not raw_video:
-                    raise TypeError("Cosmos3 action video input must be a non-empty list of PIL images or image paths.")
-                image = _pil_to_rgb(raw_video[0])
+                assert raw_video_frames is not None  # raw_image and raw_video can't both be None here
+                image = _pil_to_rgb(raw_video_frames[0])
             else:
                 image = _pil_to_rgb(raw_image)
 
@@ -219,23 +356,38 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
                     int(target_h),
                     int(target_w),
                 )
-            else:
-                scale = max(target_w / image.width, target_h / image.height)
-                resize_w = int(np.ceil(scale * image.width))
-                resize_h = int(np.ceil(scale * image.height))
-                image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
-                left = (resize_w - target_w) // 2
-                top = (resize_h - target_h) // 2
-                image = image.crop((left, top, left + target_w, top + target_h))
-
-                prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                    image, height=target_h, width=target_w
+            elif raw_video is None:
+                prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image(
+                    image,
+                    int(target_h),
+                    int(target_w),
                 )
-            if action_mode is not None and raw_video is not None:
-                if not isinstance(raw_video, list):
-                    raise TypeError("Cosmos3 action video input must be a list of PIL images or image paths.")
+            else:
+                assert raw_video_frames is not None
+                extra = _extra_args(request)
+                condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
+                    extra.get(
+                        "condition_frame_indexes_vision",
+                        prompt.get("condition_frame_indexes_vision"),
+                    )
+                )
+                keep = _normalize_condition_video_keep(
+                    extra.get("condition_video_keep", prompt.get("condition_video_keep"))
+                )
+                max_frames = _condition_pixel_frame_count(condition_frame_indexes_vision)
+                prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
+                    raw_video_frames,
+                    int(target_h),
+                    int(target_w),
+                    max_frames,
+                    keep,
+                )
+                prompt["additional_information"]["condition_frame_indexes_vision"] = list(
+                    condition_frame_indexes_vision
+                )
+            if action_mode is not None and raw_video_frames is not None:
                 prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video(
-                    raw_video,
+                    raw_video_frames,
                     int(target_h),
                     int(target_w),
                 )
@@ -345,13 +497,13 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
 class Cosmos3OmniDiffusersPipeline(
     nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
 ):
-    """Cosmos3 text/image-to-video / text-to-image pipeline.
+    """Cosmos3 text/image/video-to-video / text-to-image pipeline.
 
     Architecture: Mixture-of-Transformers with Qwen3-VL backbone.
     - Understanding pathway: causal self-attention on text (runs once, K/V cached)
     - Generation pathway: cross-attention on noisy visual latents (runs each step)
 
-    Supports T2V, I2V, and T2I from the same class.  Mode is selected at
+    Supports T2V, I2V, V2V, and T2I from the same class.  Mode is selected at
     runtime:
 
     * **T2I** when ``prompt["modalities"]`` contains ``"image"``.  Latent
@@ -365,12 +517,45 @@ class Cosmos3OmniDiffusersPipeline(
       Frame 0 of the initial latent is set to the VAE-encoded conditioning
       image, frame-0 noise predictions are masked to zero, and the clean
       image latent is re-injected at frame 0 after each scheduler step.
+    * **V2V** when the request supplies a preprocessed video via
+      ``multi_modal_data['video']`` without an action mode. Explicit latent
+      frame indexes are kept clean with ``noisy_frame_mask`` and re-injected
+      after each scheduler step.
     * **T2V** otherwise (default video generation).
     """
 
     support_image_input: ClassVar[bool] = True
     color_format: ClassVar[str] = "RGB"
 
+    @classmethod
+    def reference_video_decode_spec(
+        cls,
+        *,
+        num_frames: int | None = None,
+        extra_args: dict[str, Any] | None = None,
+    ) -> ReferenceVideoDecodeSpec:
+        extra_args = extra_args if isinstance(extra_args, dict) else {}
+        action_mode = normalize_action_mode(extra_args.get("action_mode"))
+        if action_mode is not None:
+            if num_frames is not None:
+                return ReferenceVideoDecodeSpec(max_frames=int(num_frames), keep="first")
+            action_chunk_size = extra_args.get("action_chunk_size")
+            if action_chunk_size is not None:
+                try:
+                    max_frames = int(action_chunk_size) + 1
+                except (TypeError, ValueError):
+                    max_frames = None
+                if max_frames is not None and max_frames > 0:
+                    return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")
+            return ReferenceVideoDecodeSpec(max_frames=None, keep="first")
+
+        condition_indexes = _normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
+        max_frames = _condition_pixel_frame_count(condition_indexes)
+        if num_frames is not None:
+            max_frames = min(max_frames, int(num_frames))
+        keep = _normalize_condition_video_keep(extra_args.get("condition_video_keep"))
+        return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep)
+
     def __init__(
         self,
         *,
@@ -1188,6 +1373,70 @@ class Cosmos3OmniDiffusersPipeline(
         velocity_mask = 1.0 - condition_mask
         return latents, velocity_mask, image_latent
 
+    def _prepare_latents_v2v(
+        self,
+        video_tensor: torch.Tensor,
+        height: int,
+        width: int,
+        num_frames: int,
+        generator: torch.Generator,
+        condition_frame_indexes_vision: Iterable[int] | int | str | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Prepare V2V latents with explicit clean conditioned latent frames."""
+        del height, width
+        if video_tensor.ndim == 4:
+            video_tensor = video_tensor.unsqueeze(0)
+        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+        if video_tensor.shape[2] < 1:
+            raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.")
+
+        C = self.transformer.latent_channel_size
+        T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
+        W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
+        indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
+        out_of_range = [index for index in indexes if index >= T_lat]
+        if out_of_range:
+            raise ValueError(
+                "Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: "
+                f"indexes={indexes}, latent_frames={T_lat}."
+            )
+
+        noise = randn_tensor(
+            (1, C, T_lat, H_lat, W_lat),
+            generator=generator,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        condition_pixel_frames = _condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
+        condition_video = video_tensor[:, :, :condition_pixel_frames]
+        if condition_video.shape[2] < condition_pixel_frames:
+            pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1)
+            condition_video = torch.cat([condition_video, pad], dim=2)
+
+        cond_prefix_latent = self._encode_video_tensor(condition_video)
+        expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat)
+        if (
+            cond_prefix_latent.shape[0] != expected_prefix[0]
+            or cond_prefix_latent.shape[1] != expected_prefix[1]
+            or cond_prefix_latent.shape[2] < expected_prefix[2]
+            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
+        ):
+            raise ValueError(
+                "Cosmos3 V2V condition latent shape mismatch: "
+                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
+            )
+
+        condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype)
+        condition_latents = torch.zeros_like(noise)
+        for index in indexes:
+            condition_mask[:, :, index, :, :] = 1.0
+            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
+        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
+        velocity_mask = 1.0 - condition_mask
+        return latents, velocity_mask, condition_latents
+
     def _prepare_latents_action_video(
         self,
         video_tensor: torch.Tensor,
@@ -1567,19 +1816,20 @@ class Cosmos3OmniDiffusersPipeline(
             prompt = prompt_data
             negative_prompt = None
             image_tensor = None
-            action_video_tensor = None
+            video_tensor = None
         else:
             prompt = prompt_data.get("prompt", "")
             negative_prompt = prompt_data.get("negative_prompt")
             additional_info = prompt_data.get("additional_information", {}) or {}
             image_tensor = additional_info.get("preprocessed_image")
-            action_video_tensor = additional_info.get("preprocessed_video")
+            video_tensor = additional_info.get("preprocessed_video")
 
         sp = req.sampling_params
         is_t2i = self._is_t2i_request(req)
         sound_enabled = self._is_sound_request(prompt_data, sp)
         action_mode = self._get_action_mode(prompt_data, sp)
         action_enabled = action_mode is not None
+        action_video_tensor = video_tensor if action_enabled else None
         if action_enabled and is_t2i:
             raise ValueError("Cosmos3 action generation is supported only for video outputs.")
         if action_enabled and sound_enabled:
@@ -1603,6 +1853,11 @@ class Cosmos3OmniDiffusersPipeline(
             )
         if negative_prompt is None:
             negative_prompt = ""
+        if image_tensor is not None and video_tensor is not None and not action_enabled:
+            raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
+        if video_tensor is not None and is_t2i:
+            raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")
+        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
 
         # T2I and T2V share the same model + forward path; only defaults
         # differ:
@@ -1626,7 +1881,7 @@ class Cosmos3OmniDiffusersPipeline(
             # Fall back to the engine-init shift, NOT None: passing None
             # to ``_set_flow_shift`` would leak a prior T2I rebuild
             # (shift=3.0) into a subsequent video request.
-            default_flow_shift = self._engine_init_flow_shift
+            default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
@@ -1775,6 +2030,23 @@ class Cosmos3OmniDiffusersPipeline(
                 generator,
             )
             image_latent = condition_latents[:, :, 0:1]
+        elif is_v2v:
+            condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
+                self._get_sp_param(
+                    sp,
+                    "condition_frame_indexes_vision",
+                    self._get_prompt_param(prompt_data, "condition_frame_indexes_vision", None),
+                )
+            )
+            latents, velocity_mask, condition_latents = self._prepare_latents_v2v(
+                video_tensor,
+                height,
+                width,
+                num_frames,
+                generator,
+                condition_frame_indexes_vision,
+            )
+            image_latent = None
         elif image_tensor is not None and not is_t2i:
             latents, velocity_mask, image_latent = self._prepare_latents_i2v(
                 image_tensor,

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
@@ -31,6 +31,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = init_logger(__name__)
 
 
@@ -332,6 +335,8 @@ class HunyuanVideo15Attention(nn.Module):
         out_bias: bool = True,
         eps: float = 1e-6,
         out_dim: int | None = None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -350,6 +355,8 @@ class HunyuanVideo15Attention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=self.heads,
             bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
         )
 
         self.to_out = nn.ModuleList(
@@ -360,6 +367,8 @@ class HunyuanVideo15Attention(nn.Module):
                     bias=out_bias,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
                 ),
                 nn.Identity(),  # placeholder for dropout (none used)
             ]
@@ -374,6 +383,8 @@ class HunyuanVideo15Attention(nn.Module):
                 head_size=self.head_dim,
                 total_num_heads=self.heads,
                 bias=added_proj_bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.add_kv_proj",
             )
 
             self.to_add_out = RowParallelLinear(
@@ -382,6 +393,8 @@ class HunyuanVideo15Attention(nn.Module):
                 bias=out_bias,
                 input_is_parallel=True,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_add_out",
             )
 
         self.rope = RotaryEmbedding(is_neox_style=False)
@@ -401,6 +414,8 @@ class HunyuanVideo15Attention(nn.Module):
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         hidden_states_mask: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # Ensure contiguous for FP8 quantized linear layers
+        hidden_states = hidden_states.contiguous()
         qkv, _ = self.to_qkv(hidden_states)
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
@@ -421,6 +436,7 @@ class HunyuanVideo15Attention(nn.Module):
             key = self.rope(key, cos, sin)
 
         if encoder_hidden_states is not None:
+            encoder_hidden_states = encoder_hidden_states.contiguous()
             encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
             add_q_size = self.add_kv_proj.num_heads * self.head_dim
             add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
@@ -492,6 +508,8 @@ class HunyuanVideo15TransformerBlock(nn.Module):
         attention_head_dim: int,
         mlp_ratio: float = 4.0,
         qk_norm: str = "rms_norm",
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         hidden_size = num_attention_heads * attention_head_dim
@@ -507,13 +525,23 @@ class HunyuanVideo15TransformerBlock(nn.Module):
             out_dim=hidden_size,
             bias=True,
             eps=1e-6,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
 
         self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+        self.ff = FeedForward(
+            dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio, quant_config=quant_config, prefix=f"{prefix}.ff"
+        )
 
         self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+        self.ff_context = FeedForward(
+            dim=hidden_size,
+            dim_out=hidden_size,
+            mult=mlp_ratio,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff_context",
+        )
 
     def forward(
         self,
@@ -601,6 +629,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
         target_size: int = 640,
         task_type: str = "i2v",
         use_meanflow: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -632,9 +661,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
         self.transformer_blocks = nn.ModuleList(
             [
                 HunyuanVideo15TransformerBlock(
-                    num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
+                    num_attention_heads,
+                    attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    qk_norm=qk_norm,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -143,7 +143,9 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
             self.scheduler._shift = od_config.flow_shift
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
-        self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = HunyuanVideo15Transformer3DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
 
         # Check if model uses meanflow (distilled variants)
         self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -186,7 +186,9 @@ class HunyuanVideo15I2VPipeline(
             self.scheduler._shift = od_config.flow_shift
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
-        self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = HunyuanVideo15Transformer3DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
 
         self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)
 

vllm_omni/diffusion/models/interface.py

@@ -2,10 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
     Any,
     ClassVar,
+    Literal,
     Protocol,
     runtime_checkable,
 )
@@ -23,6 +25,12 @@ class SupportImageInput(Protocol):
     color_format: ClassVar[str] = "RGB"  # Default color format
 
 
+@dataclass(frozen=True)
+class ReferenceVideoDecodeSpec:
+    max_frames: int | None = None
+    keep: Literal["first", "last"] = "first"
+
+
 @runtime_checkable
 class SupportAudioInput(Protocol):
     support_audio_input: ClassVar[bool] = True

vllm_omni/diffusion/models/lance/pipeline_lance.py

@@ -793,17 +793,9 @@ class LancePipeline(BagelPipeline):
             curr_rope=cfg_text_context["ropes"],
             video_shapes=[video_shape],
         )
-        cfg_img_lat = self.bagel.prepare_video_latent_cfg(
-            curr_kvlens=gen_context["kv_lens"],
-            curr_rope=gen_context["ropes"],
-            video_shapes=[video_shape],
-        )
         for k, v in cfg_text_lat.items():
             if torch.is_tensor(v):
                 cfg_text_lat[k] = v.to(self.device)
-        for k, v in cfg_img_lat.items():
-            if torch.is_tensor(v):
-                cfg_img_lat[k] = v.to(self.device)
 
         self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
 
@@ -822,7 +814,7 @@ class LancePipeline(BagelPipeline):
             latents, *_ = self.bagel.generate_image(
                 past_key_values=gen_context["past_key_values"],
                 cfg_text_past_key_values=cfg_text_context["past_key_values"],
-                cfg_img_past_key_values=gen_context["past_key_values"],  # no img CFG branch
+                cfg_img_past_key_values=None,  # no img CFG branch
                 num_timesteps=num_timesteps,
                 timestep_shift=timestep_shift,
                 cfg_text_scale=cfg_text_scale,
@@ -835,7 +827,7 @@ class LancePipeline(BagelPipeline):
                 # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
                 # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
                 # (derived from ``cfg_text_past_key_values``).
-                cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
+                cfg_img_packed_position_ids=None,
                 # ``cfg_img_*`` index/lens kwargs removed — same as above.
             )
 
@@ -1038,11 +1030,10 @@ class LancePipeline(BagelPipeline):
             if torch.is_tensor(v):
                 cfg_text_lat[k] = v.to(self.device)
 
-        # cfg_img branch: keep off (cfg_img_scale=1.0 matches upstream
-        # ``cfg_vit_scale=1.0``).  generate_image needs the cfg_img_* args
-        # but skips the branch when scale<=1.0; just point them at the
-        # cond's lat metadata so shapes/types are fine.
-        cfg_img_lat = cfg_text_lat
+        # cfg_img branch is off by default (cfg_img_scale=1.0 matches
+        # upstream ``cfg_vit_scale=1.0``).  Avoid passing synthetic metadata
+        # in that path; Bagel.generate_image only reads it when scale > 1.0.
+        cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
 
         self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
 
@@ -1071,7 +1062,9 @@ class LancePipeline(BagelPipeline):
                 # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
                 # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
                 # (derived from ``cfg_text_past_key_values``).
-                cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
+                cfg_img_packed_position_ids=(
+                    cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
+                ),
                 # ``cfg_img_*`` index/lens kwargs removed — same as above.
             )
 
@@ -1269,7 +1262,7 @@ class LancePipeline(BagelPipeline):
         for k, v in cfg_text_lat.items():
             if torch.is_tensor(v):
                 cfg_text_lat[k] = v.to(self.device)
-        cfg_img_lat = cfg_text_lat
+        cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
 
         self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
 
@@ -1358,7 +1351,9 @@ class LancePipeline(BagelPipeline):
                 # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
                 # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
                 # (derived from ``cfg_text_past_key_values``).
-                cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
+                cfg_img_packed_position_ids=(
+                    cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
+                ),
                 # ``cfg_img_*`` index/lens kwargs removed — same as above.
                 frame_condition_token_indexes=frame_condition_token_indexes,
             )
@@ -1587,9 +1582,9 @@ class LancePipeline(BagelPipeline):
             if torch.is_tensor(v):
                 cfg_text_lat[k] = v.to(self.device)
 
-        # cfg_img branch is off (cfg_img_scale=1.0); generate_image needs the
-        # arg surface but skips the branch when scale<=1.0.
-        cfg_img_lat = cfg_text_lat
+        # cfg_img branch is off by default; pass metadata only when an
+        # explicit cfg_img_scale enables the branch.
+        cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
 
         self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
 
@@ -1618,7 +1613,9 @@ class LancePipeline(BagelPipeline):
                 # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
                 # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
                 # (derived from ``cfg_text_past_key_values``).
-                cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
+                cfg_img_packed_position_ids=(
+                    cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
+                ),
                 # ``cfg_img_*`` index/lens kwargs removed — same as above.
             )
 

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -62,7 +62,17 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
         kwargs["has_weight"] = elementwise_affine
     elif not elementwise_affine:
         raise TypeError("RMSNorm backend does not support disabling affine weights.")
-    return RMSNorm(hidden_size, **kwargs)
+    norm = RMSNorm(hidden_size, **kwargs)
+    weight = getattr(norm, "weight", None)
+    if (
+        not elementwise_affine
+        and isinstance(weight, torch.Tensor)
+        and not isinstance(weight, nn.Parameter)
+        and "weight" not in dict(norm.named_buffers())
+    ):
+        delattr(norm, "weight")
+        norm.register_buffer("weight", weight, persistent=False)
+    return norm
 
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -10,7 +10,6 @@ This pipeline does NOT inherit from LTX2Pipeline because:
   versus LTX-2's per_layer_masked_mean_norm + shared projection path
 - LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz)
 - LTX-2.3 transformer requires the sigma parameter for prompt modulation
-- CPU offloading is required for the 22B transformer (~44GB VRAM)
 """
 
 from __future__ import annotations
@@ -20,7 +19,7 @@ import json
 import os
 from collections.abc import Iterable
 from contextlib import nullcontext
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -46,7 +45,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .pipeline_ltx2 import (
@@ -112,7 +113,7 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
+class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery):
     """Fully independent LTX-2.3 pipeline.
 
     Key differences from LTX2Pipeline:
@@ -120,11 +121,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
     - Connectors: uses padding_side API (not additive_mask)
     - Vocoder: uses LTX2VocoderWithBWE (48kHz output)
     - Transformer: passes sigma for prompt_adaln
-    - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU
     """
 
     # Audio is diffused jointly with video; warmup must size audio tokens.
     dummy_run_num_frames = 2
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "connectors"]
+    _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]
+    _resident_modules: ClassVar[list[str]] = ["vocoder"]
 
     def __init__(
         self,
@@ -166,7 +170,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         # --- Tokenizer (lightweight, stays wherever) ---
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
-        # --- Text encoder: load on CPU, move to GPU only during encoding ---
+        # --- Text encoder ---
         with torch.device("cpu"):
             self.text_encoder = from_pretrained_with_prefetch(
                 Gemma3ForConditionalGeneration.from_pretrained,
@@ -177,7 +181,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
                 torch_dtype=dtype,
             )
 
-        # --- Connectors: CPU (LTX-2.3 connectors include caption projection) ---
+        # --- Connectors (LTX-2.3 connectors include caption projection) ---
         self.connectors = from_pretrained_with_prefetch(
             LTX2TextConnectors.from_pretrained,
             model,
@@ -187,7 +191,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             torch_dtype=dtype,
         )
 
-        # --- VAE, Audio VAE: CPU ---
+        # --- VAE, Audio VAE ---
         self.vae = from_pretrained_with_prefetch(
             AutoencoderKLLTX2Video.from_pretrained,
             model,
@@ -220,6 +224,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
         quant_config = getattr(self.od_config, "quantization_config", None)
         self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
+        self._place_aux_components()
 
         # --- Scheduler ---
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
@@ -261,6 +266,20 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         self._num_timesteps = None
         self._current_timestep = None
 
+    def _place_aux_components(self) -> None:
+        parallel_config = getattr(self.od_config, "parallel_config", None)
+        use_managed_placement = bool(
+            getattr(self.od_config, "enable_cpu_offload", False)
+            or getattr(self.od_config, "enable_layerwise_offload", False)
+            or getattr(parallel_config, "use_hsdp", False)
+        )
+        if use_managed_placement:
+            return
+
+        modules = ModuleDiscovery.discover(self)
+        for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
+            module.to(self.device)
+
     # ------------------------------------------------------------------
     # Text Encoding (LTX-2.3 specific)
     # ------------------------------------------------------------------
@@ -303,16 +322,11 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         text_input_ids = text_inputs.input_ids.to(device)
         prompt_attention_mask = text_inputs.attention_mask.to(device)
 
-        # Move text encoder to GPU for encoding
-        self.text_encoder.to(device)
         text_encoder_outputs = self.text_encoder(
             input_ids=text_input_ids,
             attention_mask=prompt_attention_mask,
             output_hidden_states=True,
         )
-        # Move text encoder back to CPU immediately
-        self.text_encoder.to("cpu")
-        torch.accelerator.empty_cache()
 
         hidden_states = text_encoder_outputs.hidden_states
 
@@ -958,14 +972,10 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
-        self.connectors.to(device)
         tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
             prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
         )
-        self.connectors.to("cpu")
-        if torch.cuda.is_available():
-            torch.accelerator.empty_cache()
 
         positive_connector_prompt_embeds = connector_prompt_embeds
         positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds
@@ -1241,22 +1251,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
                 ]
                 latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
 
-            # Move VAE, audio_vae, vocoder to GPU for decoding
-            self.vae.to(device)
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)
-            self.vae.to("cpu")
 
-            self.audio_vae.to(device)
             audio_latents = audio_latents.to(self.audio_vae.dtype)
             generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
-            self.audio_vae.to("cpu")
 
-            self.vocoder.to(device)
             audio = self.vocoder(generated_mel_spectrograms)
-            self.vocoder.to("cpu")
-            torch.accelerator.empty_cache()
 
         return DiffusionOutput(output=(video, audio))
 

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -101,8 +101,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
     def load_model(
         self,
         memory_pool_context_fn: callable | None = None,
-        load_format: str | None = None,
-        custom_pipeline_name: str | type | None = None,
+        load_format: str = "default",
+        custom_pipeline_name: str | None = None,
     ) -> None:
         """
         Load the diffusion model, apply compilation and offloading.
@@ -138,7 +138,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         with get_memory_context():
             with DeviceMemoryProfiler() as m:
                 self.pipeline = model_loader.load_model(
-                    od_config=self.od_config,
                     load_device=load_device,
                     load_format=load_format,
                     custom_pipeline_name=custom_pipeline_name,

vllm_omni/entrypoints/chat_utils.py

@@ -1,62 +0,0 @@
-import numpy as np
-
-
-async def extract_audio_from_video_async(video_url: str) -> tuple[np.ndarray, int | float]:
-    """Extract audio from a video URL using vllm's load_audio.
-
-    Returns a (audio_array, sample_rate) tuple compatible with audio format.
-    All blocking I/O operations are run in a thread pool.
-    """
-    import asyncio
-    import os
-    import tempfile
-    from urllib.parse import urlparse
-
-    parsed_url = urlparse(video_url)
-    temp_video_file_path = None
-
-    def _download_video_sync(url: str) -> bytes:
-        from urllib.request import urlopen
-
-        return urlopen(url).read()
-
-    def _write_temp_file_sync(data: bytes, suffix: str) -> str:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
-            temp_file.write(data)
-            return temp_file.name
-
-    def _load_audio_sync(file_path: str) -> tuple[np.ndarray, int | float]:
-        from vllm.multimodal.media.audio import load_audio
-
-        return load_audio(file_path, sr=16000)
-
-    def _cleanup_file_sync(file_path: str) -> None:
-        try:
-            if os.path.exists(file_path):
-                os.unlink(file_path)
-        except OSError:
-            pass
-
-    try:
-        if parsed_url.scheme in ("http", "https"):
-            video_data = await asyncio.to_thread(_download_video_sync, video_url)
-            temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
-        elif parsed_url.scheme == "file":
-            from urllib.request import url2pathname
-
-            temp_video_file_path = url2pathname(parsed_url.path)
-        elif parsed_url.scheme == "data":
-            import base64
-
-            header, data = video_url.split(",", 1)
-            video_data = base64.b64decode(data)
-            temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
-        else:
-            # Assume local file path
-            temp_video_file_path = video_url
-
-        audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path)
-        return audio_array, sample_rate
-    finally:
-        if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"):
-            await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path)

vllm_omni/entrypoints/openai/api_server.py

@@ -91,6 +91,7 @@ from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
@@ -125,7 +126,7 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServin
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
-from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
+from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage, ReferenceVideo
 from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler
 from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
@@ -2531,6 +2532,84 @@ def _parse_form_json(value: str | None, expected_type: type | None = None) -> An
     return parsed
 
 
+def _config_get(config: Any, key: str, default: Any = None) -> Any:
+    if isinstance(config, dict):
+        return config.get(key, default)
+    if hasattr(config, "get"):
+        try:
+            return config.get(key, default)
+        except Exception:
+            pass
+    return getattr(config, key, default)
+
+
+def _stage_engine_args(stage_cfg: Any) -> Any:
+    return _config_get(stage_cfg, "engine_args", {}) or {}
+
+
+def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
+    if not stage_configs:
+        return []
+
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    model_classes: list[type] = []
+    for stage_cfg in stage_configs:
+        if get_stage_type(stage_cfg) != "diffusion":
+            continue
+        model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name")
+        if not model_class_name:
+            continue
+        model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+        if model_cls is not None:
+            model_classes.append(model_cls)
+    return model_classes
+
+
+def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec:
+    max_frames = spec.max_frames
+    if max_frames is not None:
+        try:
+            max_frames = int(max_frames)
+        except (TypeError, ValueError) as exc:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid reference video decode spec: max_frames must be an integer.",
+            ) from exc
+        if max_frames <= 0:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid reference video decode spec: max_frames must be positive.",
+            )
+
+    keep = str(spec.keep or "first").strip().lower()
+    if keep not in {"first", "last"}:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.",
+        )
+    return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep))
+
+
+def _reference_video_decode_spec(
+    req: VideoGenerationRequest,
+    stage_configs: list[Any] | None,
+) -> ReferenceVideoDecodeSpec:
+    video_params = req.resolve_video_params()
+    extra_params = req.extra_params if isinstance(req.extra_params, dict) else {}
+    for model_cls in _diffusion_model_classes(stage_configs):
+        resolver = getattr(model_cls, "reference_video_decode_spec", None)
+        if resolver is None:
+            continue
+        try:
+            spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params)
+        except (TypeError, ValueError) as exc:
+            raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
+        if spec is not None:
+            return _normalize_reference_video_decode_spec(spec)
+    return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first")
+
+
 def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse:
     resp = VideoResponse(
         model=model_name,
@@ -2597,6 +2676,7 @@ async def _run_video_generation_job(
     request: VideoGenerationRequest,
     video_id: str,
     reference_image: ReferenceImage | None = None,
+    reference_video: ReferenceVideo | None = None,
     app_state: Any | None = None,
 ) -> None:
     job = await VIDEO_STORE.get(video_id)
@@ -2609,7 +2689,10 @@ async def _run_video_generation_job(
     output_path = None
     try:
         video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
-            request, video_id, reference_image=reference_image
+            request,
+            video_id,
+            reference_image=reference_image,
+            reference_video=reference_video,
         )
 
         file_name = f"{video_id}.{job.file_extension}"
@@ -2676,6 +2759,7 @@ async def _parse_video_form(
     prompt: str = Form(...),
     input_reference: UploadFile | None = File(default=None),
     image_reference: str | None = Form(default=None),
+    video_reference: str | None = Form(default=None),
     model: str | None = Form(default=None),
     seconds: SecondStr | None = Form(default=None),
     size: SizeStr | None = Form(default=None),
@@ -2700,7 +2784,7 @@ async def _parse_video_form(
     frame_interpolation_model_path: str | None = Form(default=None),
     lora: str | None = Form(default=None),
     extra_params: str | None = Form(default=None),
-) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]:
+) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None, ReferenceVideo | None]:
     """FastAPI dependency that parses video form data, validates inputs,
     resolves the handler, and decodes any reference image.
 
@@ -2708,11 +2792,15 @@ async def _parse_video_form(
     """
     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)
 
-    if parsed_image_reference is not None and input_reference_bytes is not None:
+    provided_references = sum(
+        item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes)
+    )
+    if provided_references > 1:
         raise HTTPException(
             status_code=HTTPStatus.BAD_REQUEST.value,
-            detail="Provide either input_reference or image_reference, not both.",
+            detail="Provide only one of input_reference, image_reference, or video_reference.",
         )
 
     request_data: dict[str, Any] = {
@@ -2721,6 +2809,7 @@ async def _parse_video_form(
         "seconds": seconds,
         "size": size,
         "image_reference": parsed_image_reference,
+        "video_reference": parsed_video_reference,
         "user": user,
         "width": width,
         "height": height,
@@ -2774,13 +2863,28 @@ async def _parse_video_form(
             detail=f"Video generation setup failed: {str(e)}",
         )
 
+    decode_spec = ReferenceVideoDecodeSpec()
+    if 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:
-        image_data = await decode_input_reference(request.image_reference, input_reference_bytes)
+        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=image_data) if image_data is not None else None
-    return request, handler, effective_model_name, reference_image
+    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
+    return request, handler, effective_model_name, reference_image, reference_video
 
 
 @router.post(
@@ -2794,18 +2898,31 @@ async def _parse_video_form(
 )
 async def create_video(
     raw_request: Request,
-    ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
+    ctx: tuple[
+        VideoGenerationRequest,
+        OmniOpenAIServingVideo,
+        str,
+        ReferenceImage | None,
+        ReferenceVideo | None,
+    ] = Depends(_parse_video_form),
 ) -> VideoResponse:
     """Create an asynchronous video generation job.
 
     Accepts multipart form-data (see ``_parse_video_form`` for parameters),
     persists a queued job record, and starts generation in the background.
     """
-    request, handler, effective_model_name, reference_image = ctx
+    request, handler, effective_model_name, reference_image, reference_video = ctx
     ref = video_response_from_request(effective_model_name, request)
     await VIDEO_STORE.upsert(ref.id, ref)
     task = asyncio.create_task(
-        _run_video_generation_job(handler, request, ref.id, reference_image, app_state=raw_request.app.state)
+        _run_video_generation_job(
+            handler,
+            request,
+            ref.id,
+            reference_image,
+            reference_video,
+            app_state=raw_request.app.state,
+        )
     )
     await VIDEO_TASKS.upsert(ref.id, task)
     return ref
@@ -2822,7 +2939,13 @@ async def create_video(
 )
 async def create_video_sync(
     raw_request: Request,
-    ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
+    ctx: tuple[
+        VideoGenerationRequest,
+        OmniOpenAIServingVideo,
+        str,
+        ReferenceImage | None,
+        ReferenceVideo | None,
+    ] = Depends(_parse_video_form),
 ) -> Response:
     """Synchronous video generation endpoint.
 
@@ -2833,13 +2956,18 @@ async def create_video_sync(
     Metadata is returned via response headers ``X-Request-Id``,
     ``X-Model``, and ``X-Inference-Time-S``.
     """
-    request, handler, effective_model_name, reference_image = ctx
+    request, handler, effective_model_name, reference_image, reference_video = ctx
     request_id = f"video_sync-{random_uuid()}"
     raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
     started_at = time.perf_counter()
     try:
         video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for(
-            handler.generate_video_bytes(request, request_id, reference_image=reference_image),
+            handler.generate_video_bytes(
+                request,
+                request_id,
+                reference_image=reference_image,
+                reference_video=reference_video,
+            ),
             timeout=VIDEO_SYNC_TIMEOUT_S,
         )
     except asyncio.TimeoutError:

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -73,6 +73,19 @@ class UrlImageReference(BaseModel):
 ImageReference = UrlImageReference | FileImageReference
 
 
+class FileVideoReference(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    file_id: str
+
+
+class UrlVideoReference(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    video_url: str
+
+
+VideoReference = UrlVideoReference | FileVideoReference
+
+
 class VideoGenerationRequest(BaseModel):
     """
     OpenAI-style video generation request.
@@ -99,6 +112,10 @@ class VideoGenerationRequest(BaseModel):
         default=None,
         description="Optional JSON-safe image reference that guides generation. Provide either image_url or file_id.",
     )
+    video_reference: VideoReference | None = Field(
+        default=None,
+        description="Optional JSON-safe video reference that guides generation. Provide either video_url or file_id.",
+    )
 
     # Video params block for extensibility
     video_params: VideoParams | None = Field(default=None, description="Optional video-specific parameters")

vllm_omni/entrypoints/openai/serving_chat.py

@@ -12,8 +12,16 @@ from typing import Any, Final, cast
 import jinja2
 import torch
 from fastapi import Request
+from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
 from PIL import Image
 from pydantic import TypeAdapter
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    ChatTemplateContentFormatOption,
+    ConversationMessage,
+    get_history_tool_calls_cnt,
+    make_tool_call_id,
+)
 
 from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -32,14 +40,6 @@ except ImportError:
     soundfile = None
 
 
-from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
-from vllm.entrypoints.chat_utils import (
-    ChatCompletionMessageParam,
-    ChatTemplateContentFormatOption,
-    ConversationMessage,
-    get_history_tool_calls_cnt,
-    make_tool_call_id,
-)
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -652,19 +652,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         )
 
         tok_params = request.build_tok_params(self.model_config)
+        mm_config = self.model_config.multimodal_config
         chat_params = request.build_chat_params(
             default_template,
             default_template_content_format,
-        ).with_defaults(default_template_kwargs)
-
-        # OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm
-        # processor asserts that audio items are present alongside video items
-        # (inside render_chat_async).  We must inject audio_url items into the
-        # messages BEFORE calling render_chat_async so the mm processor can
-        # count them correctly during tokenisation.
-        mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
-        if mm_proc_kw.get("use_audio_in_video", False):
-            messages = await self._inject_audio_from_video_urls(messages)
+        ).with_defaults(
+            default_template_kwargs,
+            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
+            default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
+        )
 
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],
@@ -735,95 +731,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return conversation, [engine_prompt]
 
-    async def _inject_audio_from_video_urls(
-        self,
-        messages: list[ChatCompletionMessageParam],
-    ) -> list[ChatCompletionMessageParam]:
-        """Pre-extract audio from video URLs and inject as audio_url content items.
-
-        When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
-        processor requires that the number of audio items equals the number of
-        video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
-        The client only sends video_url items; this method adds the matching
-        audio_url items on the server side before the renderer processes them.
-        """
-        import io
-
-        from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
-
-        new_messages: list[ChatCompletionMessageParam] = []
-        for msg in messages:
-            content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
-            if not isinstance(content, list):
-                new_messages.append(msg)
-                continue
-
-            video_urls = [
-                part.get("video_url", {}).get("url")
-                for part in content
-                if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
-            ]
-
-            if not video_urls:
-                new_messages.append(msg)
-                continue
-
-            audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
-
-            audio_items: list[dict] = []
-            for audio_array, sample_rate in audios:
-                buf = io.BytesIO()
-                if soundfile is not None:
-                    soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
-                else:
-                    import struct
-
-                    import numpy as np
-
-                    audio_np = np.asarray(audio_array, dtype=np.float32)
-                    sr = int(sample_rate)
-                    num_channels = 1
-                    bits_per_sample = 32
-                    num_frames = len(audio_np)
-                    data_size = num_frames * num_channels * (bits_per_sample // 8)
-                    # Write minimal RIFF/WAV header
-                    buf.write(b"RIFF")
-                    buf.write(struct.pack("<I", 36 + data_size))
-                    buf.write(b"WAVE")
-                    buf.write(b"fmt ")
-                    buf.write(
-                        struct.pack(
-                            "<IHHIIHH",
-                            16,
-                            3,
-                            num_channels,
-                            sr,
-                            sr * num_channels * (bits_per_sample // 8),
-                            num_channels * (bits_per_sample // 8),
-                            bits_per_sample,
-                        )
-                    )
-                    buf.write(b"data")
-                    buf.write(struct.pack("<I", data_size))
-                    buf.write(audio_np.tobytes())
-
-                audio_b64 = base64.b64encode(buf.getvalue()).decode()
-                audio_items.append(
-                    {
-                        "type": "audio_url",
-                        "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
-                    }
-                )
-
-            new_content = list(content) + audio_items
-            if isinstance(msg, dict):
-                new_msg = {**msg, "content": new_content}
-            else:
-                new_msg = msg.model_copy(update={"content": new_content})
-            new_messages.append(new_msg)
-
-        return new_messages
-
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
         """Convert request dicts to stage-typed sampling params objects.
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -9,6 +9,7 @@ import re
 import struct
 import time
 from collections import OrderedDict
+from collections.abc import Mapping
 from concurrent.futures import ThreadPoolExecutor
 from http import HTTPStatus
 from pathlib import Path
@@ -102,19 +103,21 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "higgs_audio_v2",
     "higgs_audio_v3",
 }
-_TTS_LANGUAGES: set[str] = {
-    "Auto",
-    "Chinese",
-    "English",
-    "Japanese",
-    "Korean",
-    "German",
-    "French",
-    "Russian",
-    "Portuguese",
-    "Spanish",
-    "Italian",
-}
+_TTS_LANGUAGES = frozenset(
+    {
+        "Auto",
+        "Chinese",
+        "English",
+        "Japanese",
+        "Korean",
+        "German",
+        "French",
+        "Russian",
+        "Portuguese",
+        "Spanish",
+        "Italian",
+    }
+)
 _REF_AUDIO_MIN_DURATION = 1.0  # seconds
 _REF_AUDIO_MAX_DURATION = 30.0  # seconds
 _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256
@@ -434,6 +437,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Merge built-in speakers into the set initialized by _init_speaker_storage.
         self.supported_speakers |= self._load_supported_speakers()
         self.supported_speakers |= set(self.precomputed_speakers)
+
+        self.supported_languages = self._load_supported_languages()
+
         self._tts_tokenizer = None
         self._voxcpm2_tokenizer = None
         self._voxcpm2_split_map: dict[int, list[int]] = {}
@@ -694,6 +700,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         return set()
 
+    def _load_supported_languages(self) -> frozenset[str]:
+        """Load supported languages (title-cased) from the model configuration"""
+        if self._tts_model_type != "qwen3_tts":
+            return _TTS_LANGUAGES
+        try:
+            config = self.engine_client.model_config.hf_config.talker_config
+
+            if isinstance(config, dict):
+                codec_language_id = config.get("codec_language_id")
+            else:
+                codec_language_id = getattr(config, "codec_language_id", None)
+
+            if codec_language_id and isinstance(codec_language_id, Mapping):
+                return frozenset(str(language).title() for language in codec_language_id) | {"Auto"}
+
+            logger.warning("No codec_language_id found in talker_config; falling back to default languages")
+        except Exception as e:
+            logger.warning("Could not load languages from model config: %s", e)
+        return _TTS_LANGUAGES
+
     def _estimate_ref_code_len(self, ref_audio: object) -> int | None:
         """Estimate ref_code length from ref_audio waveform without running the codec.
 
@@ -1420,9 +1446,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if not request.input or not request.input.strip():
             return "Input text cannot be empty"
 
-        # Validate language
-        if request.language is not None and request.language not in _TTS_LANGUAGES:
-            return f"Invalid language '{request.language}'. Supported: {', '.join(sorted(_TTS_LANGUAGES))}"
+        # Validate language (case-insensitive; normalized to the title-cased config form)
+        if request.language is not None:
+            request.language = request.language.title()
+            if request.language not in self.supported_languages:
+                return (
+                    f"Invalid language '{request.language}'. Supported: {', '.join(sorted(self.supported_languages))}"
+                )
 
         # Validate speaker for CustomVoice task
         if task_type == "CustomVoice":

vllm_omni/entrypoints/openai/serving_video.py

@@ -39,6 +39,13 @@ class ReferenceImage:
     data: Image.Image
 
 
+@dataclass
+class ReferenceVideo:
+    """Reference video frames for video-conditioned generation."""
+
+    data: list[Image.Image]
+
+
 @dataclass
 class VideoGenerationArtifacts:
     """Normalized outputs and profiler metadata extracted from one request."""
@@ -96,6 +103,7 @@ class OmniOpenAIServingVideo:
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
+        reference_video: ReferenceVideo | None = None,
     ) -> VideoGenerationArtifacts:
         """Run the generation pipeline and extract video/audio/profiler outputs."""
         prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"])
@@ -105,6 +113,12 @@ class OmniOpenAIServingVideo:
         gen_params = self._resolve_default_sampling_params()
 
         input_image = None if reference_image is None else reference_image.data
+        input_video = None if reference_video is None else reference_video.data
+        if input_image is not None and input_video is not None:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Provide either an image reference or a video reference, not both.",
+            )
         vp = request.resolve_video_params()
         if input_image is not None and vp.width is not None and vp.height is not None:
             target_size = (vp.width, vp.height)
@@ -112,6 +126,8 @@ class OmniOpenAIServingVideo:
                 input_image = input_image.resize(target_size, Image.Resampling.LANCZOS)
         if input_image is not None:
             prompt["multi_modal_data"] = {"image": input_image}
+        elif input_video is not None:
+            prompt["multi_modal_data"] = {"video": input_video}
         if vp.width is not None and vp.height is not None:
             gen_params.width = vp.width
             gen_params.height = vp.height
@@ -211,8 +227,14 @@ class OmniOpenAIServingVideo:
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
+        reference_video: ReferenceVideo | None = None,
     ) -> VideoGenerationResponse:
-        artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
+        artifacts = await self._run_and_extract(
+            request,
+            reference_id,
+            reference_image=reference_image,
+            reference_video=reference_video,
+        )
 
         video_codec_options = {"preset": "ultrafast", "threads": "0"}
         if request.extra_params is not None and isinstance(request.extra_params, dict):
@@ -252,9 +274,15 @@ class OmniOpenAIServingVideo:
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
+        reference_video: ReferenceVideo | None = None,
     ) -> tuple[bytes, dict[str, float], float, VideoAction | None]:
         """Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
-        artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
+        artifacts = await self._run_and_extract(
+            request,
+            reference_id,
+            reference_image=reference_image,
+            reference_video=reference_video,
+        )
         if len(artifacts.videos) > 1:
             logger.warning(
                 "Video request %s generated %d outputs; returning only the first.",

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -8,8 +8,9 @@ from __future__ import annotations
 
 import base64
 import binascii
+from collections import deque
 from io import BytesIO
-from typing import Any
+from typing import Any, Literal
 
 import httpx
 import numpy as np
@@ -19,8 +20,11 @@ from PIL import Image, UnidentifiedImageError
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.protocol.videos import (
     FileImageReference,
+    FileVideoReference,
     ImageReference,
     UrlImageReference,
+    UrlVideoReference,
+    VideoReference,
 )
 
 
@@ -31,6 +35,70 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
         raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
 
 
+def _decode_video_bytes(
+    video_bytes: bytes,
+    *,
+    source: str,
+    max_frames: int | None = None,
+    keep: Literal["first", "last"] = "first",
+) -> list[Image.Image]:
+    try:
+        import av
+    except ImportError as exc:  # pragma: no cover - av is a serving dependency via media_utils
+        raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc
+
+    if keep not in {"first", "last"}:
+        raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.")
+    if max_frames is not None and max_frames <= 0:
+        raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.")
+
+    frames: list[Image.Image] = []
+    tail_frames: deque[Image.Image] | None = (
+        deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None
+    )
+    try:
+        with av.open(BytesIO(video_bytes)) as container:
+            for frame in container.decode(video=0):
+                image = frame.to_image().convert("RGB")
+                if tail_frames is not None:
+                    tail_frames.append(image)
+                else:
+                    frames.append(image)
+                if keep == "first" and max_frames is not None and len(frames) >= max_frames:
+                    break
+    except Exception as exc:
+        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc
+
+    if tail_frames is not None:
+        frames = list(tail_frames)
+    if not frames:
+        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.")
+    return frames
+
+
+def _decode_media_bytes(
+    media_bytes: bytes,
+    *,
+    source: str,
+    max_video_frames: int | None = None,
+    video_keep: Literal["first", "last"] = "first",
+) -> Image.Image | list[Image.Image]:
+    try:
+        return _decode_image_bytes(media_bytes, source=source)
+    except InvalidInputReferenceError:
+        try:
+            return _decode_video_bytes(
+                media_bytes,
+                source=source,
+                max_frames=max_video_frames,
+                keep=video_keep,
+            )
+        except InvalidInputReferenceError as video_exc:
+            raise InvalidInputReferenceError(
+                f"Invalid {source}: provided content is not a valid image or video."
+            ) from video_exc
+
+
 def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
     if input_reference:
         if input_reference.startswith("data:image"):
@@ -64,23 +132,96 @@ async def decode_image_url(image_url: str) -> Image.Image:
     raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
 
 
+def _decode_base64_video(
+    video_reference: str,
+    *,
+    source: str,
+    max_frames: int | None = None,
+    keep: Literal["first", "last"] = "first",
+) -> list[Image.Image]:
+    if video_reference:
+        if video_reference.startswith("data:video"):
+            _, b64_data = video_reference.split(",", 1)
+        else:
+            b64_data = video_reference
+
+        try:
+            video_bytes = base64.b64decode(b64_data)
+        except (binascii.Error, ValueError) as exc:  # pragma: no cover - malformed base64
+            raise InvalidInputReferenceError(f"Invalid {source}: video data is not valid base64.") from exc
+        return _decode_video_bytes(video_bytes, source=source, max_frames=max_frames, keep=keep)
+    raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
+
+
+async def decode_video_url(
+    video_url: str,
+    *,
+    max_frames: int | None = None,
+    keep: Literal["first", "last"] = "first",
+) -> list[Image.Image]:
+    if video_url.startswith("data:video"):
+        return _decode_base64_video(
+            video_url,
+            source="video_reference.video_url",
+            max_frames=max_frames,
+            keep=keep,
+        )
+
+    if video_url.startswith(("http://", "https://")):
+        async with httpx.AsyncClient(timeout=60) as client:
+            try:
+                response = await client.get(video_url)
+                response.raise_for_status()
+            except httpx.HTTPError as exc:
+                raise InvalidInputReferenceError(
+                    "Invalid video_reference.video_url: failed to download video."
+                ) from exc
+        return _decode_video_bytes(
+            response.content,
+            source="video_reference.video_url",
+            max_frames=max_frames,
+            keep=keep,
+        )
+
+    raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
+
+
 async def decode_input_reference(
     image_reference: ImageReference | None,
+    video_reference: VideoReference | None,
     input_reference_bytes: bytes | None,
-) -> Image.Image | None:
-    """Decode image input from multipart bytes, base64/data URL, or image_reference."""
+    *,
+    max_video_frames: int | None = None,
+    video_keep: Literal["first", "last"] = "first",
+) -> Image.Image | list[Image.Image] | None:
+    """Decode media input from multipart bytes, data URLs, or typed references."""
 
-    if input_reference_bytes is not None and image_reference is not None:
-        raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.")
+    provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
+    if provided > 1:
+        raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")
 
     if isinstance(input_reference_bytes, bytes):
-        return _decode_image_bytes(input_reference_bytes, source="input_reference")
+        return _decode_media_bytes(
+            input_reference_bytes,
+            source="input_reference",
+            max_video_frames=max_video_frames,
+            video_keep=video_keep,
+        )
 
     if isinstance(image_reference, UrlImageReference):
         return await decode_image_url(image_reference.image_url)
     elif isinstance(image_reference, FileImageReference):
         raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
 
+    if isinstance(video_reference, UrlVideoReference):
+        return await decode_video_url(
+            video_reference.video_url,
+            max_frames=max_video_frames,
+            keep=video_keep,
+        )
+    elif isinstance(video_reference, FileVideoReference):
+        raise InvalidInputReferenceError("Invalid video_reference: file_id is not supported yet.")
+
     return None
 
 

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -250,7 +250,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         if input_ids is not None:
             self._last_step_input_ids = input_ids
 
-        # Stash query_start_loc for _apply_audio_mode_bias
+        # Stash query_start_loc and max_query_len for prefill detection
+        _max_query_len = None
         try:
             from vllm.forward_context import get_forward_context
 
@@ -264,13 +265,19 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                 self._last_step_query_start_loc = qsl.detach().clone()
             else:
                 self._last_step_query_start_loc = None
+            _max_query_len = getattr(attn, "max_query_len", None)
         except Exception:
             self._last_step_query_start_loc = None
 
-        # Prefill-only operations: ref audio substitution and audio feedback
-        # require Python dict/list ops that break CUDA graph capture.
-        # Detect prefill (sequence length > batch size heuristic) vs decode.
-        is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1
+        # Detect prefill vs decode using attn_metadata.max_query_len.
+        # max_query_len == 1 means pure decode (even with N concurrent
+        # requests, each contributes exactly 1 token). numel > 1 is NOT
+        # reliable because decode with N>1 concurrent requests gives
+        # input_ids.numel() == N.
+        if _max_query_len is not None:
+            is_prefill = int(_max_query_len) > 1
+        else:
+            is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1
         if is_prefill:
             # Voice clone: replace -100 placeholder positions with ref audio embeddings
             info_dicts = kwargs.get("model_intermediate_buffer")
@@ -278,7 +285,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                 info_dicts = kwargs.get("runtime_additional_information")
             hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts)
 
-        # Audio feedback at decode: replace continuation token embeddings
+        # Audio feedback: replace continuation token embeddings with audio
+        # embeddings from the last decoded frame (CUDA-graph safe).
         if input_ids is not None and inputs_embeds is None:
             hidden_states = self._apply_audio_feedback(hidden_states, input_ids)
 
@@ -411,11 +419,14 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             self._decode_last_codes = self._decode_last_codes.to(hidden_states.device)
             self._decode_has_codes = self._decode_has_codes.to(hidden_states.device)
         if bs > self._decode_last_codes.shape[0]:
-            new_size = max(bs, self._decode_last_codes.shape[0] * 2)
-            self._decode_last_codes = torch.zeros(
-                new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device
-            )
-            self._decode_has_codes = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device)
+            old_size = self._decode_last_codes.shape[0]
+            new_size = max(bs, old_size * 2)
+            new_codes = torch.zeros(new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device)
+            new_has = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device)
+            new_codes[:old_size] = self._decode_last_codes
+            new_has[:old_size] = self._decode_has_codes
+            self._decode_last_codes = new_codes
+            self._decode_has_codes = new_has
 
         # Compute audio embeddings from last_codes for ALL rows (graph-safe)
         codes_slice = self._decode_last_codes[:bs]  # [bs, N]
@@ -561,7 +572,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                     num_remaining_delays = num_codebooks - last_eos_idx - 1
 
             if num_remaining_delays is not None and num_remaining_delays <= 0:
-                # Ramp-down complete — terminate
+                # Ramp-down complete — terminate without emitting this frame.
                 state["num_delay"] = 0
                 state["num_remaining_delays"] = None
                 state["should_terminate"] = True
@@ -578,6 +589,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                 state["audio_out_ids"] = this_codes.unsqueeze(-1).clone()
             else:
                 state["audio_out_ids"] = torch.cat([state["audio_out_ids"], this_codes.unsqueeze(-1)], dim=-1)
+
             new_codes_flat.append(this_codes)
 
         # Build full codes tensor [batch_size, num_codebooks]

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1457,6 +1457,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
     HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs
 
     prefer_model_sampler = True
+    logitsprocs_need_output_token_ids = True
 
     # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
     # weights are replicated and the image batch is sharded across TP ranks.

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -1,8 +1,9 @@
-"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.14) with minimal overrides."""
+"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides."""
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Iterator, Mapping
 from typing import Any
 
+import numpy as np
 import torch
 from torch import nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -49,9 +50,7 @@ from vllm.model_executor.models.utils import (
     WeightsMapper,
     init_vllm_registered_model,
     maybe_prefix,
-    split_list_into_ranges,
 )
-from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -92,20 +91,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
         self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
 
+        # Detect use_audio_in_video from mm_kwargs
         use_audio_in_video = False
         if "video" in mm_kwargs:
-            non_none_items = [item for item in mm_kwargs["video"] if item is not None]
-            if non_none_items:
-                for item in non_none_items:
-                    if item.get("use_audio_in_video"):
-                        uaiv_tensor = item["use_audio_in_video"].data
-                        if uaiv_tensor.numel() > 0:
-                            use_audio_in_video = bool(uaiv_tensor.item())
-                            break
-            elif "audio" in mm_prompt_updates:
-                tokenizer = self.info.get_tokenizer()
-                audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
-                use_audio_in_video = audio_pad_id not in prompt_ids
+            for item in mm_kwargs["video"]:
+                if item and item.get("use_audio_in_video"):
+                    use_audio_in_video_tensor = item["use_audio_in_video"].data
+                    if use_audio_in_video_tensor.numel() > 0:
+                        use_audio_in_video = bool(use_audio_in_video_tensor.item())
+                        break
             # for mutilmodality cache
             if any(item is None for item in mm_kwargs["video"]):
                 video_token_id = self.info.get_hf_config().video_token_id
@@ -421,167 +415,193 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
+    def _get_audio_for_video_mapping(self, mm_features: list[MultiModalFeatureSpec]) -> tuple[dict[int, int], set[int]]:
+        """
+        Map video offset -> paired audio_feature_length for use_audio_in_video.
+
+        When use_audio_in_video=True, audio is interleaved within video chunks.
+        The pairing is based on feature order in mm_features.
+        """
+        videos_with_audio = [
+            f
+            for f in mm_features
+            if f.modality == "video" and f.data.get("use_audio_in_video") and f.data["use_audio_in_video"].data.item()
+        ]
+        audios = [f for f in mm_features if f.modality == "audio"]
+
+        mapping: dict[int, int] = {}
+        paired_audio_offsets: set[int] = set()
+        for i, video_f in enumerate(videos_with_audio):
+            if i < len(audios):
+                audio_len = audios[i].data["audio_feature_lengths"].data.item()
+                mapping[video_f.mm_position.offset] = audio_len
+                paired_audio_offsets.add(audios[i].mm_position.offset)
+        return mapping, paired_audio_offsets
+
+    def _compute_audio_token_count(self, audio_feature_length: int) -> int:
+        return ((audio_feature_length - 1) // 2 + 1 - 2) // 2 + 1
+
+    def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator[tuple[int, str, dict[str, Any]]]:
+        thinker_config = self.config
+        spatial_merge_size = thinker_config.vision_config.spatial_merge_size
+        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
+
+        sorted_features = sorted(mm_features, key=lambda f: f.mm_position.offset)
+        audio_for_video, paired_audio_offsets = self._get_audio_for_video_mapping(sorted_features)
+
+        for mm_feature in sorted_features:
+            offset = mm_feature.mm_position.offset
+            modality = mm_feature.modality
+
+            if modality == "image":
+                t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
+                yield (
+                    offset,
+                    "image",
+                    {
+                        "grid_t": t,
+                        "grid_h": h // spatial_merge_size,
+                        "grid_w": w // spatial_merge_size,
+                        "t_factor": 1.0 * tokens_per_second,
+                    },
+                )
+            elif modality == "video":
+                t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
+                second_per_grid_ts = 1.0
+                if mm_feature.data.get("second_per_grid_ts"):
+                    second_per_grid_ts = mm_feature.data["second_per_grid_ts"].data.item()
+                use_audio_in_video = False
+                if mm_feature.data.get("use_audio_in_video"):
+                    use_audio_in_video = bool(mm_feature.data["use_audio_in_video"].data.item())
+
+                yield (
+                    offset,
+                    "video",
+                    {
+                        "grid_t": t,
+                        "grid_h": h // spatial_merge_size,
+                        "grid_w": w // spatial_merge_size,
+                        "t_factor": second_per_grid_ts * tokens_per_second,
+                        "use_audio_in_video": use_audio_in_video,
+                        "audio_feature_length": audio_for_video.get(offset),
+                    },
+                )
+            elif modality == "audio":
+                if offset not in paired_audio_offsets:
+                    audio_len = mm_feature.data["audio_feature_lengths"].data.item()
+                    yield offset, "audio", {"audio_feature_length": audio_len}
+
+    def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]:
+        grid_t = data["grid_t"]
+        grid_h = data["grid_h"]
+        grid_w = data["grid_w"]
+        t_factor = data["t_factor"]
+        audio_len = data["audio_feature_length"]
+
+        thinker_config = self.config
+        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
+        seconds_per_chunk = thinker_config.seconds_per_chunk
+        t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
+
+        t_index = (np.arange(grid_t) * t_factor).astype(np.int64)
+        t_index_split_chunk: list[list[int]] = [[] for _ in range((int(t_index.max()) // t_ntoken_per_chunk) + 1)]
+        for t_val in t_index:
+            idx = int(t_val) // t_ntoken_per_chunk
+            t_index_split_chunk[idx].append(int(t_val))
+
+        pure_audio_len = self._compute_audio_token_count(audio_len)
+        added_audio_len = 0
+        pos_ids_list: list[np.ndarray] = []
+        audio_start_idx = start_idx
+
+        for t_chunk in t_index_split_chunk:
+            if not t_chunk:
+                continue
+
+            chunk_t = len(t_chunk)
+
+            h_indices = np.tile(np.arange(grid_h).reshape(1, -1, 1), (chunk_t, 1, grid_w)).flatten()
+            w_indices = np.tile(np.arange(grid_w).reshape(1, 1, -1), (chunk_t, grid_h, 1)).flatten()
+            t_indices = np.repeat(np.array(t_chunk), grid_h * grid_w)
+
+            vision_pos = np.stack([t_indices, h_indices, w_indices]) + start_idx
+            pos_ids_list.append(vision_pos)
+
+            audio_chunk_size = min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
+            if audio_chunk_size > 0:
+                audio_pos = np.broadcast_to(np.arange(audio_chunk_size), (3, audio_chunk_size)) + audio_start_idx
+                pos_ids_list.append(audio_pos)
+                audio_start_idx = audio_start_idx + audio_chunk_size
+                added_audio_len += audio_chunk_size
+
+        if added_audio_len < pure_audio_len:
+            remaining = pure_audio_len - added_audio_len
+            remaining_audio_pos = np.broadcast_to(np.arange(remaining), (3, remaining)) + audio_start_idx
+            pos_ids_list.append(remaining_audio_pos)
+
+        vision_tokens = grid_t * grid_h * grid_w
+        total_tokens = vision_tokens + pure_audio_len
+
+        return np.concatenate(pos_ids_list, axis=1), total_tokens
+
     def get_mrope_input_positions(
         self,
         input_tokens: list[int],
         mm_features: list[MultiModalFeatureSpec],
     ) -> tuple[torch.Tensor, int]:
-        """
-        Example:
+        llm_pos_ids_list: list[np.ndarray] = []
+        st = 0
 
-            (V_i are vision position ids, A_i are audio position ids)
+        for offset, modality, data in self.iter_mm_features(mm_features):
+            text_len = offset - st
+            st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
+            if text_len > 0:
+                llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
+                st_idx += text_len
 
-            |V_1 ...    V_n|A_1 ...   A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
-            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
-        """
-        kwargs = MultiModalFeatureSpec.gather_kwargs(
-            mm_features,
-            {
-                "image_grid_thw",
-                "video_grid_thw",
-                "second_per_grid_ts",
-                "audio_feature_lengths",
-                "use_audio_in_video",
-            },
-        )
-        image_grid_thw = kwargs.get("image_grid_thw", [])
-        video_grid_thw = kwargs.get("video_grid_thw", [])
-        second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
-        audio_feature_lengths = kwargs.get("audio_feature_lengths", [])
-        use_audio_in_video = any(kwargs.get("use_audio_in_video", []))
+            if modality == "audio":
+                audio_tokens = self._compute_audio_token_count(data["audio_feature_length"])
+                llm_pos_ids_list.append(np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx)
+                st = offset + audio_tokens
 
-        image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw)
-        video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw)
+            elif modality == "image":
+                grid_t = data["grid_t"]
+                grid_h = data["grid_h"]
+                grid_w = data["grid_w"]
+                t_factor = data["t_factor"]
 
-        # TODO(fyabc): refactor and share more code with
-        #  _vl_get_input_positions_tensor.
+                grid_indices = np.indices((grid_t, grid_h, grid_w))
+                if t_factor != 1.0:
+                    grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
+                llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
+                st = offset + grid_t * grid_h * grid_w
 
-        thinker_config = self.config
-        audio_token_id = thinker_config.audio_token_index
-        image_token_id = thinker_config.image_token_index
-        video_token_id = thinker_config.video_token_index
-        audio_start_token_id = thinker_config.audio_start_token_id
-        audio_end_token_id = thinker_config.audio_end_token_id
-        vision_start_token_id = thinker_config.vision_start_token_id
-        vision_end_token_id = thinker_config.vision_end_token_id
-        seconds_per_chunk = thinker_config.seconds_per_chunk
-        spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
+            elif modality == "video":
+                grid_t = data["grid_t"]
+                grid_h = data["grid_h"]
+                grid_w = data["grid_w"]
+                t_factor = data["t_factor"]
 
-        src_item = input_tokens
-        audio_seqlens = audio_feature_lengths
-        if not second_per_grid_ts:
-            second_per_grid_ts = [1] * video_grid_thw.shape[0]
-        audio_idx = 0
-        video_idx = 0
-        image_idx = 0
-        new_src_item: list[int] = []
-        llm_pos_ids_list: list[torch.Tensor] = []
+                if not data["use_audio_in_video"]:
+                    grid_indices = np.indices((grid_t, grid_h, grid_w))
+                    if t_factor != 1.0:
+                        grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
+                    llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
+                    st = offset + grid_t * grid_h * grid_w
+                else:
+                    pos_ids, token_count = self._compute_interleaved_positions(st_idx, data)
+                    llm_pos_ids_list.append(pos_ids)
+                    st = offset + token_count
 
-        idx = 0
-        while idx < len(src_item):
-            new_src_item_len = len(new_src_item)
-            start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
-                if use_audio_in_video and idx > 0:
-                    if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
-                        # processing the <|audio_eos|> before <|vision_eos|>
-                        start_idx -= 1
-                    elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
-                        # processing the <|audio_bos|> after <|vision_eos|>
-                        start_idx -= 1
-                new_src_item.append(src_item[idx])
-                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
-                llm_pos_ids_list.append(llm_pos_ids)
-            elif src_item[idx] == audio_token_id:
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
-                new_src_item.extend([audio_token_id] * place_num)
-                llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
-                llm_pos_ids_list.append(llm_pos_ids)
-                audio_idx += 1
-            elif src_item[idx] == image_token_id:
-                grid_t = image_grid_thw[image_idx][0]
-                grid_hs = image_grid_thw[:, 1]
-                grid_ws = image_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([image_token_id] * vision_seqlen)
-                image_idx += 1
-            elif src_item[idx] == video_token_id and not use_audio_in_video:
-                grid_t = video_grid_thw[video_idx][0]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                new_src_item.extend([video_token_id] * vision_seqlen)
-                video_idx += 1
-            else:
-                # read audio from video
-                assert audio_seqlens is not None
-                audio_seqlen = audio_seqlens[audio_idx]
-                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                grid_t = video_grid_thw[video_idx][0]
-                grid_h = video_grid_thw[video_idx][1]
-                grid_w = video_grid_thw[video_idx][2]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
-                t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk)
-                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
-                pure_audio_len = place_num - 2
-                added_audio_len = 0
-                audio_llm_pos_ids_list: list[torch.Tensor] = []
-                for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
-                    vision_llm_pos_ids_list = get_llm_pos_ids_for_vision(
-                        start_idx,
-                        video_idx,
-                        spatial_merge_size,
-                        t_chunk,
-                        grid_hs,
-                        grid_ws,
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(vision_llm_pos_ids_list)
-                    new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_start_idx = (
-                        start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
-                    )
-                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
-                        audio_llm_pos_ids_list = (
-                            torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
-                            + audio_start_idx
-                        ).split(1, dim=1)
-                    else:
-                        audio_llm_pos_ids_list = []
-                    added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                if added_audio_len < pure_audio_len:
-                    new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
-                    audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
-                    ).split(1, dim=1)
-                    llm_pos_ids_list.extend(audio_llm_pos_ids_list)
-                audio_idx += 1
-                video_idx += 1
-            # move to the next token
-            idx += len(new_src_item) - new_src_item_len
+        if st < len(input_tokens):
+            st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
+            text_len = len(input_tokens) - st
+            llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
 
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
+        llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
+        mrope_position_delta = int(llm_positions.max()) + 1 - len(input_tokens)
 
-        return llm_positions, mrope_position_delta
+        return torch.from_numpy(llm_positions), mrope_position_delta
 
     def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
@@ -629,14 +649,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # Check for audio-in-video: interleaved video and audio tokens
         # in the multimodal region. Only use the interleaved path when
         # needed; otherwise fall back to the default parent implementation.
-        # vLLM's _gather_mm_embeddings builds is_multimodal on CPU for merge
-        # indexing; bitwise ops with input_ids require the same device.
-        is_mm_device = is_multimodal.to(device=input_ids.device, non_blocking=True)
         video_token_id = self.config.video_token_index
         audio_token_id = self.config.audio_token_index
 
-        is_video = is_mm_device & (input_ids == video_token_id)
-        is_audio = is_mm_device & (input_ids == audio_token_id)
+        input_ids_cpu = input_ids.cpu()
+        is_video = is_multimodal & (input_ids_cpu == video_token_id)
+        is_audio = is_multimodal & (input_ids_cpu == audio_token_id)
 
         num_video = is_video.sum().item()
         num_audio = is_audio.sum().item()
@@ -652,7 +670,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 multimodal_embeddings,
                 is_video,
                 is_audio,
-                is_mm_device,
+                is_multimodal,
                 num_video,
                 num_audio,
             )

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -758,17 +758,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
         use_audio_in_video = False
         if "video" in mm_kwargs:
-            non_none_items = [item for item in mm_kwargs["video"] if item is not None]
-            if non_none_items:
-                # Normal case: at least one non-cached item, read flag directly
-                use_audio_in_video = any(item["use_audio_in_video"].data for item in non_none_items)
-            elif "audio" in mm_prompt_updates:
-                # All video items are from cache (None); infer from prompt:
-                # use_audio_in_video=True means the prompt has no <|audio_pad|>
-                # placeholder (audio is embedded in video tokens instead)
-                tokenizer = self.info.get_tokenizer()
-                audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
-                use_audio_in_video = audio_pad_id not in prompt_ids
+            for item in mm_kwargs["video"]:
+                if item and item["use_audio_in_video"].data:
+                    use_audio_in_video = True
+                else:
+                    use_audio_in_video = False
             # for mutilmodality cache
             if any(item is None for item in mm_kwargs["video"]):
                 video_token_id = self.info.get_hf_config().video_token_id

vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml

@@ -0,0 +1,38 @@
+# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect.
+# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP.
+#
+# Use with a ModelOpt FP8 checkpoint (e.g. produced by
+# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py).
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      model_class_name: HunyuanVideo15Pipeline
+      max_num_seqs: 1
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
+      # the full video (the DiT contributes little), so tiling the decode is the
+      # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB).
+      # Default it on here; override with --vae-use-tiling/false if you have headroom.
+      vae_use_tiling: true
+      parallel_config:
+        tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: video
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml

@@ -0,0 +1,41 @@
+# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect.
+# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half).
+# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`.
+#
+# Use with a ModelOpt FP8 checkpoint (e.g. produced by
+# examples/quantization/quantize_wan2_2_modelopt_fp8.py).
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects
+      # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone
+      # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.)
+      model_class_name: Wan22Pipeline
+      max_num_seqs: 1
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so
+      # FP8 is the main lever here), but tiling adds headroom for higher
+      # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false.
+      vae_use_tiling: true
+      parallel_config:
+        tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: video
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py

@@ -135,7 +135,25 @@ def talker2code2wav(
         codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
 
         # Step 1: Revert delay pattern
-        codes_qt = _revert_delay_pattern(codes_qt)
+        try:
+            codes_qt = _revert_delay_pattern(codes_qt)
+        except ValueError:
+            logger.warning(
+                "Skipping request: insufficient frames for delay pattern "
+                "reversal (T=%d < Q=%d). This is a request-level error, "
+                "not a server-fatal condition.",
+                codes_qt.shape[1],
+                codes_qt.shape[0],
+            )
+            code2wav_inputs.append(
+                OmniTokensPrompt(
+                    prompt_token_ids=[],
+                    multi_modal_data=None,
+                    mm_processor_kwargs=None,
+                    additional_information=None,
+                )
+            )
+            continue
 
         # Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0.
         # Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023
@@ -336,7 +354,21 @@ def talker2code2wav_async_chunk(
         return None
 
     codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous()
-    de_delayed = _revert_delay_pattern(codes_qt)
+    try:
+        de_delayed = _revert_delay_pattern(codes_qt)
+    except ValueError:
+        logger.warning(
+            "async_chunk: insufficient frames for delay pattern reversal "
+            "(T=%d < Q=%d, request=%s). Returning empty chunk.",
+            codes_qt.shape[1],
+            codes_qt.shape[0],
+            request_id,
+        )
+        emitted_frames.pop(request_id, None)
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+            meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+        )
     # Replace BOC=1024/EOC=1025 (and any negative pads) with 0; matches the
     # sync-path substitution. Clamp would turn 1025 into 1023 which is a
     # VALID codec code and decodes to audible artifacts.

vllm_omni/patch.py

@@ -217,3 +217,55 @@ def _patch_scaled_mm_fp8_contiguous_activation():
 
 
 _patch_scaled_mm_fp8_contiguous_activation()
+
+
+def _patch_flashinfer_fp8_scaled_mm_output_shape():
+    """Restore the N-D output shape for the FlashInfer FP8 ScaledMM kernel.
+
+    ``FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm`` returns the raw 2-D
+    GEMM result and ignores ``output_shape``, unlike the CUTLASS / PyTorch
+    ScaledMM kernels which reshape to it. A 3-D activation ``(B, S, D)`` thus
+    collapses to ``(B*S, D)``, breaking diffusion DiTs that reshape the linear
+    output by absolute dim (e.g. Wan2.2 ``qkv.unflatten(2, ...)``) with
+    ``IndexError: Dimension out of range``. It only bites >2-D inputs, so LLM
+    (token-flattened, 2-D) paths are unaffected.
+
+    Carried here because the upstream fix may not have landed yet; this override
+    becomes a harmless no-op once vLLM honors ``output_shape`` itself.
+    """
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
+            FlashInferFP8ScaledMMLinearKernel,
+        )
+    except ImportError:
+        return
+
+    _original_apply_scaled_mm = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
+    if getattr(_original_apply_scaled_mm, "_omni_output_shape_patched", False):
+        return
+
+    def _patched_apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape):  # noqa: N803
+        out = _original_apply_scaled_mm(
+            self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape
+        )
+        if tuple(out.shape) != tuple(output_shape):
+            out = out.view(*output_shape)
+        return out
+
+    _patched_apply_scaled_mm._omni_output_shape_patched = True
+    FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = _patched_apply_scaled_mm
+
+
+_patch_flashinfer_fp8_scaled_mm_output_shape()
+
+
+def _patch_fp8_use_quack_fused_bias():
+    try:
+        from vllm_omni.quantization.quack_fp8 import install_quack_fp8_patch
+
+        install_quack_fp8_patch()
+    except Exception:  # noqa: BLE001
+        pass
+
+
+_patch_fp8_use_quack_fused_bias()

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -33,6 +33,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         # in _pad_for_sequence_parallelism during execute_model.
         # This is a workaround for vllm-ascend not passing vllm_config to enable_sp().
         enable_sp(self.vllm_config)
+        self._maybe_enable_output_token_ids_for_model_sampler()
         self._init_talker_mtp()
 
     @torch.inference_mode()

vllm_omni/quantization/quack_fp8.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FP8 scaled-MM with a fused bias epilogue, backed by quack's CuteDSL GEMM."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_TRUTHY = {"1", "true", "yes", "on"}
+_gemm_interface = None
+
+
+def _is_blackwell() -> bool:
+    try:
+        if not torch.cuda.is_available():
+            return False
+        return torch.cuda.get_device_capability()[0] >= 10
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def quack_enabled() -> bool:
+    override = os.environ.get("VLLM_OMNI_USE_QUACK_FP8")
+    if override is not None:
+        return override.lower() in _TRUTHY
+    return _is_blackwell()
+
+
+def _set_persistent_cache_dir() -> None:
+    if os.environ.get("QUACK_CACHE_DIR"):
+        return
+    root = (
+        os.environ.get("VLLM_CACHE_ROOT")
+        or os.environ.get("XDG_CACHE_HOME")
+        or os.path.join(os.path.expanduser("~"), ".cache")
+    )
+    os.environ["QUACK_CACHE_DIR"] = os.path.join(root, "vllm_omni", "quack")
+
+
+def _load_quack():
+    global _gemm_interface
+    if _gemm_interface is not None:
+        return _gemm_interface or None
+    try:
+        _set_persistent_cache_dir()
+
+        import cutlass
+        import cutlass.base_dsl
+        import cutlass.base_dsl.arch as arch
+
+        if not hasattr(cutlass.base_dsl, "Arch"):
+            cutlass.base_dsl.Arch = arch.Arch
+
+        import quack.gemm_interface as gemm_interface
+        from quack.cute_dsl_utils import torch2cute_dtype_map
+
+        torch2cute_dtype_map.setdefault(torch.float8_e4m3fn, cutlass.Float8E4M3FN)
+        torch2cute_dtype_map.setdefault(torch.float8_e5m2, cutlass.Float8E5M2)
+
+        _gemm_interface = gemm_interface
+        logger.info("Quack FP8 fused-bias GEMM enabled (CuteDSL).")
+        return gemm_interface
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("Quack FP8 unavailable, using FlashInfer: %s", exc)
+        _gemm_interface = False
+        return None
+
+
+def quack_scaled_fp8_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    out_dtype: torch.dtype,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    gemm = _load_quack()
+    if gemm is None:
+        return None
+    out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
+    alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
+    gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
+    return out
+
+
+def install_quack_fp8_patch() -> None:
+    if not quack_enabled():
+        return
+    if _load_quack() is None:
+        return
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
+            FlashInferFP8ScaledMMLinearKernel,
+        )
+    except ImportError:
+        return
+
+    original = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
+    if getattr(original, "_omni_quack_patched", False):
+        return
+
+    def apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape):  # noqa: N803
+        try:
+            out = quack_scaled_fp8_mm(A, B, As, Bs, out_dtype, bias)
+            if out is not None:
+                return out.view(*output_shape)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning_once("Quack FP8 GEMM failed (%s); using FlashInfer.", exc)
+        return original(self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape)
+
+    apply_scaled_mm._omni_quack_patched = True
+    FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = apply_scaled_mm
+    logger.info("Patched FlashInfer FP8 ScaledMM to use quack fused-bias GEMM.")
+
+
+def warmup_quack_fp8(
+    shapes: list[tuple[int, int, int]],
+    device: str = "cuda",
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    if _load_quack() is None:
+        return
+    scale = torch.ones(1, device=device, dtype=torch.float32)
+    for m, k, n in shapes:
+        a = torch.zeros(m, k, device=device, dtype=torch.float8_e4m3fn)
+        b = torch.zeros(k, n, device=device, dtype=torch.float8_e4m3fn)
+        quack_scaled_fp8_mm(a, b, scale, scale, out_dtype)
+    if torch.cuda.is_available():
+        torch.accelerator.synchronize()

vllm_omni/worker/gpu_model_runner.py

@@ -143,9 +143,14 @@ class OmniGPUModelRunner(GPUModelRunner):
     @instrument(span_name="Loading (GPU)")
     def load_model(self, *args, **kwargs) -> None:
         super().load_model(*args, **kwargs)
+        self._maybe_enable_output_token_ids_for_model_sampler()
         self._init_talker_mtp()
         self._prewarm_attention_capture_workspaces()
 
+    def _maybe_enable_output_token_ids_for_model_sampler(self) -> None:
+        if getattr(self.model, "logitsprocs_need_output_token_ids", False):
+            self.input_batch.logitsprocs_need_output_token_ids = True
+
     def _init_talker_mtp(self) -> None:
         # TODO move this model specific logic to a separate class
         # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.