docs: define module ownership contracts

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
docs: update module design navigation
2026-08-07 11:45:59 +08:00 · 2026-07-16 14:09:07 +08:00 · 2026-07-16 12:08:45 +08:00
211 changed files with 2378 additions and 13875 deletions
@@ -404,7 +404,7 @@ steps:
      - label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
        timeout_in_minutes: 90
        commands:
-          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -m "full_model and cuda" --run-level "full_model"
        mirror_hardwares: h100_1
      - label: ":full_moon: Diffusion X2V · Doc Test"
@@ -47,17 +47,15 @@ steps:
          - timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
        mirror_hardwares: l4_1
-      - label: "Diffusion · Distributed & Attention Test"
+      - label: "Diffusion · Distributed Test"
        source_file_dependencies:
          - tests/diffusion/distributed/
          - tests/diffusion/attention/test_attention_sp.py
          - vllm_omni/diffusion/distributed/
-          - vllm_omni/diffusion/attention/
+          - vllm_omni/diffusion/attention/parallel/
          - vllm_omni/diffusion/data.py
          - vllm_omni/diffusion/models/
        commands:
          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
          - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
        mirror_hardwares: l4_4
      - label: "Diffusion · Model Test"
@@ -220,70 +220,65 @@ See `plan/voxcpm2_native_ar_design.md`.
 ### Steps
-1. **Write one adapter** under `vllm_omni/entrypoints/openai/tts_adapters/`.
+1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
-   `serving_speech.py` should not need an edit — detection, stage discovery and
+   partial integration causes hard-to-debug failures. This file is modified by every
-   dispatch are all derived from what the adapter declares.
+   model PR and is the most common source of rebase conflicts — see conflict note below.
-   Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`:
+   **Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
   ```python
-   @register_tts_adapter
+   _YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
-   class YourModelAdapter(ARTTSAdapter):
+   ```
       name = "your_model"                              # registry key + log label
       stage_keys = frozenset({"your_stage_key"})       # the deploy yaml's model_stage
-       def validate(self, request) -> str | None:
+   **Point 2** — union into `_TTS_MODEL_STAGES`:
   ```python
   _TTS_MODEL_STAGES: set[str] = (
       ...
       | _YOUR_MODEL_TTS_MODEL_STAGES
   )
   ```
   **Point 3** — model type detection in `_detect_tts_model_type()`:
   ```python
   if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
       return "your_model"
   ```
   **Point 4** — validation dispatch in `_validate_tts_request()`:
   ```python
   if self._tts_model_type == "your_model":
       return self._validate_your_model_request(request)
   ```
   **Point 5** — validation + parameter-builder methods:
   ```python
   def _validate_your_model_request(self, request) -> str | None:
       if not request.input or not request.input.strip():
           return "Input text cannot be empty"
       return None
-       async def build(self, request, sampling_params_list, has_inline_ref_audio):
+   def _build_your_model_params(self, request) -> dict:
       params = {"text": [request.input]}
       if request.voice is not None:
           params["voice"] = [request.voice]
-           return PreparedRequest(
+       return params
               prompt={"prompt": request.input},
               tts_params=params,
               model_type=self.name,
           )
   ```
   Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
   model-specific param builders.
-   Then add the module to the import block at the bottom of
+   > **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
-   `tts_adapters/__init__.py` so it registers. That import line is the only shared
+   > instance attribute checked before `elif self._is_tts`, while all newer models
-   file a new model touches — which is also why the old rebase-conflict hotspot is
+   > (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
-   gone.
+   > `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
   > pattern — do not add new `_is_*` flags.
-   > **Pure-diffusion TTS does not go through adapters yet.** Under
+   > **Unused variable rule**: only extract fields in `_build_your_model_params` that
-   > `for_diffusion()`, `create_speech()` routes straight to
+   > are actually forwarded to the model. Unused extractions fail `ruff F841`.
-   > `_create_diffusion_speech()` and never calls `validate()`/`build()`.
+   > For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
-   > `DiffusionTTSAdapter` is scaffolding with no production subclass, so logic
+   > `prompt_text`), add them to the param builder and verify they reach the model call.
   > placed in one would silently never run. Diffusion-engine models follow the
   > existing diffusion path; wiring it through adapters is open work (#4855).
-   **If a stage key alone cannot identify the model**, declare
+   **Rebase conflict note**: when rebasing onto `main` after another model was merged,
-   `model_archs = frozenset({"YourModelForConditionalGeneration"})`; add
+   `serving_speech.py` will conflict. Resolution: always keep *both* the upstream
-   `arch_identifies_entry_stage = True` when the model owns no stage key at all
+   model's additions and your own — never discard either side.
   (Ming dense). For a rule that is not set membership, override `matches()` (see
   `covo_audio.py`). For a genuine overlap with another adapter, give one an
   explicit `detect_priority` — `test_tts_detection.py` fails on an unordered
   overlap. Models that only serve speech in some topologies override
   `stage_serves_speech()` (see `audex.py`).
   **Reuse shared helpers via `self.ctx.server`** rather than reimplementing them:
   `_resolve_ref_audio`, `_apply_uploaded_speaker`, `_validate_ref_audio_format`,
   `_max_instructions_length`. Read a comparable adapter first — `fish_speech.py`
   (voice cloning), `higgs_audio_v3.py` (parameter-heavy), `moss_tts.py` (family
   sharing a base class).
   > **Do not add `self._tts_model_type == ...` branches to `serving_speech.py`.**
   > Older models predate the adapter framework and still have them; they are being
   > migrated out (RFC #4327, #4855). `tools/pre_commit/check_tts_adapter.py` is a
   > ratchet on the remaining count and fails the commit if it grows. Behaviour that
   > no adapter hook can express is a missing hook — propose it on the RFC.
   > **Unused variable rule**: only extract fields in `build()` that are actually
   > forwarded to the model. Unused extractions fail `ruff F841`. For voice-cloning
   > fields (`ref_audio` -> `prompt_audio_path`, `ref_text` -> `prompt_text`), add
   > them to the params and verify they reach the model call.
 2. **Handle model-specific parameters**:
   - Voice cloning: `ref_audio` encoding and prompt injection
@@ -300,7 +295,7 @@ import base64
 from pathlib import Path
 def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
-    """Build prompt with reference audio for voice cloning, called from the adapter."""
+    """Build prompt with reference audio for voice cloning in serving_speech.py."""
    audio_bytes = Path(ref_audio_path).read_bytes()
    codes = codec.encode(audio_bytes)  # Encode on CPU using model's codec (e.g., DAC)
    token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
@@ -326,7 +321,7 @@ Classify the model's **CI priority** first (high / medium / low). High-priority
 | Level | Location | Marker | CI pipeline | Notes |
 |-------|----------|--------|-------------|-------|
-| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, adapter validation — no GPU |
+| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU |
 | **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path |
 | **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk |
 | **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` |
@@ -371,7 +366,7 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers
 ### Deliverables
- One adapter file under `tts_adapters/` plus its line in the package import block
+- Updated `serving_speech.py` with all 5 integration points (single commit)
 - Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
 - Gradio demo with streaming and voice cloning UI in the same dir
 - E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope)
@@ -519,8 +514,8 @@ Use this checklist when integrating a new TTS model:
 - [ ] README.md written
 ### Phase 3: Online Serving
- [ ] Adapter written under `tts_adapters/` and registered in the import block
+- [ ] All 5 `serving_speech.py` integration points added in one commit
- [ ] Only extract params in `build()` that are forwarded to the model call (ruff F841)
+- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
 - [ ] Prompt builder handles text input correctly
 - [ ] Voice cloning works (if supported)
 - [ ] All response formats work (wav, mp3, flac, pcm)
@@ -1,231 +0,0 @@
 ---
 name: review-pr
 description: Review pull requests and local branches for vllm-project/vllm-omni with a frozen snapshot, module-design ownership, feature-design overlays, targeted validation, and concise evidence-backed findings. Use for default, detailed, or repeat maintainer reviews; checking correctness, compatibility, tests, benchmarks, model additions, distributed changes, or breaking behavior; and identifying or explicitly requesting the most relevant code-owner reviewers. Use precheck-pr instead for an author's pre-submit self-check.
 ---
 # Review vLLM-Omni Pull Requests
 Review like a maintainer: direct, selective, and focused on issues that CI does
 not prove. Prefer a few high-confidence findings over exhaustive commentary.
 Zero findings is a valid result.
 ## Quality contract
 Make every finding:
 - **Correct:** prove a reachable failure, not a suspicion.
 - **Prioritized:** lead with merge blockers and high-impact defects.
 - **Actionable:** identify the smallest safe fix direction.
 - **Evidence-based:** cite code, tests, docs, CI, or measurements.
 - **Concise:** avoid review templates and repeated summaries.
 - **Calibrated:** match severity to user and maintainer impact.
 Do not report unrelated backlog, style already enforced by pre-commit, or a
 missing test that would not protect changed behavior.
 ## Select the input and depth
 Use `vllm-project/vllm-omni` as the base repository. Accept its forks and local
 checkouts; use another skill for unrelated repositories.
 | Input | Review surface |
 | --- | --- |
 | PR number or URL | Frozen PR metadata, full diff, and relevant threads. |
 | Local branch/worktree | Frozen target-base SHA through committed, staged, unstaged, and in-scope untracked changes. |
 | Pre-filled context | Reuse supplied metadata; fetch only missing facts and the full diff. |
 Default to maintainer brevity. A detailed or audit request expands coverage and
 lists `path:line` findings, but keeps the same confidence and severity bar.
 ## Reference guide
 Load references in review order: process, one primary module contract, matching
 feature designs, evidence checks, then delivery. Every file is linked directly
 below; do not load unrelated module references.
 Each concise reference links to the maintained
 [vLLM-Omni documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/).
 For branch-specific behavior, inspect the matching `docs/` file in the reviewed
 checkout first; use the published latest docs for current guidance and discovery.
 If docs and live code disagree, verify the code/tests and report the drift.
 ### Review process
 | Reference | Read when |
 | --- | --- |
 | [review-execution.md](references/process/review-execution.md) | Every review; freeze inputs, inspect safely, and deliver against the same snapshot. |
 | [general-checks.md](references/process/general-checks.md) | Every review; apply repository-wide correctness and evidence rules. |
 | [design-contracts.md](references/process/design-contracts.md) | Every production review; resolve branch-local module and feature design status. |
 | [review-routing.md](references/process/review-routing.md) | After the diff census; select one primary module and conditional overlays. |
 ### Primary module contract
 | Reference | Read when |
 | --- | --- |
 | [entrypoints.md](references/modules/entrypoints.md) | Offline/CLI/API ingress, validation, rendering, streaming, or sessions change. |
 | [configuration.md](references/modules/configuration.md) | Config construction, deploy/stage schema, defaults, registry, or topology changes. |
 | [input-output-modality.md](references/modules/input-output-modality.md) | Requests, messages, serialization, output types, accumulation, or completion change. |
 | [error-contracts.md](references/modules/error-contracts.md) | Error classification, fatality, propagation, sanitization, or rendering changes. |
 | [engine-orchestration.md](references/modules/engine-orchestration.md) | Cross-stage routing, request state, output ordering, RPC correlation, or terminal convergence changes. |
 | [stage-runtime.md](references/modules/stage-runtime.md) | Placement, startup, readiness, replica identity, affinity, membership, or shutdown changes. |
 | [omni-connector.md](references/modules/omni-connector.md) | Cross-stage/process/device/node transport or synchronization changes. |
 | [model-integration.md](references/modules/model-integration.md) | Registration, preprocessing, loading, runners, or model-specific execution changes. |
 | [ar-runtime.md](references/modules/ar-runtime.md) | AR scheduling, request/cache state, adapters, workers, or upstream vLLM semantics change. |
 | [diffusion.md](references/modules/diffusion.md) | Diffusion runtime, models, batching, parallelism, or offload changes. |
 | [execution-platforms.md](references/modules/execution-platforms.md) | Hardware selection, capabilities, vendor workers, kernels, or patches change. |
 | [cache-management.md](references/modules/cache-management.md) | Cache identity, reuse, validity, reset, eviction, or teardown changes. |
 | [quantization.md](references/modules/quantization.md) | Quantization selection, checkpoint metadata, layer mapping, precision, or constraints change. |
 | [observability.md](references/modules/observability.md) | Metrics, logs, units, labels, correlation, or lifecycle changes. |
 | [profiling.md](references/modules/profiling.md) | Profiling instrumentation, traces, start/stop lifecycle, or overhead changes. |
 | [benchmarking.md](references/modules/benchmarking.md) | Benchmark workload, metric calculation, CLI, or result metadata changes. |
 ### Feature-design overlays
 | Reference | Read when |
 | --- | --- |
 | [runtime-stage-execution.md](references/features/runtime-stage-execution.md) | Disaggregated inference, async chunk/output/materialization, or prefix caching changes. |
 | [communication.md](references/features/communication.md) | A concrete OmniConnector backend or its deployment contract changes. |
 | [diffusion-acceleration.md](references/features/diffusion-acceleration.md) | Diffusion parallelism, attention, quantization, cache, batching, or offload changes. |
 | [infrastructure-performance.md](references/features/infrastructure-performance.md) | Metrics infrastructure or documented speech optimization stacks change. |
 ### Evidence and quality checks
 | Reference | Read when |
 | --- | --- |
 | [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, or stage config is added. |
 | [perf-verification.md](references/checks/perf-verification.md) | The PR makes a latency, throughput, memory, or quality claim. |
 | [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
 | [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
 | [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
 ### Delivery and reviewer coordination
 | Reference | Read when |
 | --- | --- |
 | [maintainer-style-study.md](references/delivery/maintainer-style-study.md) | Findings are ready for concise maintainer-style delivery. |
 | [review-requests.md](references/delivery/review-requests.md) | The user asks to identify, suggest, request, or ping code-owner reviewers. |
 ## Workflow
 ### 1. Freeze and report the snapshot
 Pin the base and head before reading source or running validation. Within 60 seconds,
 report the pinned head, CI, mergeability, and preliminary findings in the host
 conversation. Do not wait for CI or post this update to GitHub.
 If the target changes while fetching, discard the evidence and retry once. If
 it changes again, report the churn and wait for a stable target.
 For a trusted PR head, materialize the pinned head in an isolated detached
 worktree. A worktree freezes identity but is not a security sandbox. Treat fork
 heads as untrusted unless the user and environment policy explicitly establish
 otherwise: execute them only in a disposable, secret-free sandbox with restricted
 filesystem, network, and resources; without one, use static SHA-addressed reads
 and CI evidence only. For a local review, freeze the committed, index, worktree,
 and NUL-safe in-scope untracked contents. Follow
 [review-execution.md](references/process/review-execution.md) for trust gates,
 state fingerprints, and byte-for-byte staleness checks.
 ### 2. Build the diff census
 Group files into production code, tests, docs, configuration, build/CI, and
 generated artifacts. Map each changed production file and test group to the PR
 goal. Compare the title/body claims with the actual diff; use linked issues only
 when they define the contract or reproduction.
 Mark unrelated scope and unexplained generated artifacts. Do not infer behavior
 from the PR description without tracing the live code.
 ### 3. Route from the live behavior
 Trace each claimed behavior through the changed producer to its live consumer,
 then use [design-contracts.md](references/process/design-contracts.md) and
 [review-routing.md](references/process/review-routing.md) to select one primary
 module contract, a second only for a real documented cross-boundary call path,
 and every matching feature-design and evidence overlay. Treat titles and paths
 as hints; live behavior and the frozen head's current design metadata are
 authoritative. For docs-, tests-, or CI-only changes, route to the production
 contract they protect or use only the applicable evidence checks.
 ### 4. Run the blocker scan
 Apply every category in [general-checks.md](references/process/general-checks.md) before
 lower-priority comments.
 For each changed value or behavior, trace:
 ```text
 public ingress -> validation/defaulting -> producer -> transformations
  -> stage/worker/connector boundary -> final consumer -> terminal cleanup
 ```
 Cover every applicable offline/online, streaming/non-streaming, sync/async,
 feature-on/off, topology, and compatibility path. Search bounded callers and
 sibling implementations rather than assuming the changed hunk is the only path.
 ### 5. Apply module and feature contracts
 Apply the reference set selected in step 3 and any matching repo-local skill.
 Read the exact module and feature pages in the frozen head, including status,
 ownership boundary, dependencies, candidate invariants, safe-change guide, and
 promotion gate. Candidate or draft rules are questions, not blockers, unless
 current code, tests, or policy enforce them. Inspect both sides of any config,
 registry, serialization, connector, cache, or stage boundary.
 When a diff adds or expands a helper, class, fallback, compatibility branch, or
 public behavior, run a subtraction pass: remove out-of-scope behavior and check
 whether each new abstraction can be deleted, merged, moved, or inlined.
 ### 6. Verify the changed path
 Before each validation group, verify the frozen SHA plus the tracked, index,
 untracked, and ignored-file fingerprint, or recreate a pristine snapshot. On a
 trusted head or inside the required sandbox, run an import/version preflight,
 then the narrowest relevant tests and low-cost static checks. Bind every result
 to the head SHA, snapshot fingerprint, and environment fingerprint. Never run
 imports, tests, builds, hooks, or repo-configurable tooling from an untrusted
 head on the reviewer host.
 - Treat CI as status evidence; inspect only the first overlapping failure.
 - For docs-only changes, use diff hygiene, links/build checks, and bounded live
  contract verification instead of dependency setup or pytest.
 - For hardware-dependent paths, run available static/CPU checks and name the
  exact GPU/NPU gap. Never simulate device evidence.
 - For performance or accuracy claims, require comparable base/head runs with
  the same environment, workload, warmup, repetitions, and quality criteria.
 Stop when each changed semantic path has a supported finding or an explicit
 no-issue conclusion. Do not search further only to increase confidence.
 ### 7. Consolidate and deliver
 Verify each finding against the current diff, deduplicate by root cause, and
 order by severity.
 Re-read the remote head and reverify or recreate the pristine validation
 snapshot immediately before delivery. If either changed, mark the review stale
 and restart from the new snapshot.
 Return findings first. Use
 [maintainer-style-study.md](references/delivery/maintainer-style-study.md) to keep them
 direct and brief. Each finding must include an exact `path:line`, trigger or
 call path, current behavior, impact, and smallest fix direction. If there are
 no findings, say so briefly and name material validation gaps.
 Keep the review read-only unless the user explicitly authorizes posting. Do not
 submit `APPROVE`, `COMMENT`, or `REQUEST_CHANGES`, add labels, edit code, or push
 commits as an implied part of review.
 ### 8. Optionally request focused owner reviews
 Only when the user asks to identify or request reviewers, read
 [review-requests.md](references/delivery/review-requests.md). Rank path-matched
 CODEOWNERS with the frozen module page's owners or required reviewers and
 documented governance expertise; propose one to three focused reviewers with an
 explicit contract rationale.
 Identifying or suggesting reviewers is read-only. Requesting reviewers or
 posting `@mention` comments changes external state and requires explicit user
 authorization. When authorized, recheck the head, deduplicate existing
 requests, and post at most one consolidated comment. Do not infer this
 permission from a request to review the code.
@@ -1,4 +0,0 @@
 interface:
  display_name: "vLLM-Omni PR Review"
  short_description: "Review vLLM-Omni PRs and route owner requests"
  default_prompt: "Use $review-pr to review this vLLM-Omni pull request with owner-specific rules and evidence-backed findings."
@@ -1,121 +0,0 @@
 # Model Addition Checklist
 Read this reference when a PR adds or expands a model architecture, pipeline,
 loader, processor, registry entry, or stage configuration.
 Official docs: [model contribution guides](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/),
 [adding an omni model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_omni_model/),
 [adding a diffusion model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model/),
 [adding a TTS model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_tts_model/),
 [supported models](https://docs.vllm.ai/projects/vllm-omni/en/latest/models/supported_models/),
 [diffusion features](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion_features/),
 and [feature compatibility](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/feature_compatibility/).
 ## Match claims to the diff
 - Reconcile every claimed architecture, modality, variant, feature, and example
  with changed files and registry entries.
 - Resolve config model types, architecture names, callable paths, exports, and
  dependency extras against live symbols.
 - Check sibling generations sharing an architecture name are disambiguated by
  a stable config predicate rather than import order or an accidental default.
 ## Close the integration chain
 Trace:
 ```text
 public model id -> config/pipeline selection -> registry -> loader/processor
  -> stage inputs -> model execution -> stage/public output
 ```
 Verify optional dependencies fail with an actionable message, weight names and
 dtypes map correctly, and each advertised serving mode reaches the production
 dispatcher. Confirm outputs are non-empty and valid for their modality, shape,
 sample rate, or response schema.
 ## Keep support documentation truthful
 - Build an evidence-backed inventory of the new model's exact architecture and
  checkpoint identifiers, modalities and tasks, hardware backends, serving
  modes, and supported acceleration, parallelism, cache, offload, and
  quantization features. Treat untested support as unknown, not supported.
 - Update `docs/models/supported_models.md` with the architecture, model family,
  example checkpoints, and only the backends validated by the PR.
 - For a diffusion model, add it to the applicable ImageGen, VideoGen, or
  AudioGen table in `docs/user_guide/diffusion_features.md`. Fill every feature
  column according to its legend and document partial support or constraints.
 - Update `docs/user_guide/feature_compatibility.md` when the model changes valid
  feature combinations, configuration constraints, examples, or limitations,
  and update any other branch-local support table affected by the claims.
 - Cross-check every documented support mark against the registry, configuration,
  production-path tests, and PR evidence; do not copy a sibling model's support
  claims without validation.
 ## Ship the model recipe
 - Require every new model PR to add or update one model-family recipe under
  `recipes/<vendor>/` and its row in the `recipes/README.md` Available Recipes
  table. Follow [`recipes/TEMPLATE.md`](../../../../../recipes/TEMPLATE.md) and
  the layout rules in
  [`recipes/README.md`](../../../../../recipes/README.md); the
  external recipes repository is a structural reference, not a substitute.
 - Include only tested task, serving-mode, and hardware configurations. Keep the
  exact model identifiers, flags, commands, verification, feature support,
  constraints, and limitations consistent with code, tests, examples, and the
  support documentation.
 - Cover each validated platform in its own hardware section. Use one recipe per
  model family by default; justify a split when configurations cannot be kept
  clear in one file, and do not add placeholder sections for untested platforms.
 ## Compare the reference implementation
 Require every new model PR to include a head-versus-reference accuracy or
 quality and performance comparison:
 - Name the canonical implementation and pin its repository revision, package
  version, configuration, and checkpoint. If it cannot run, explain why and use
  the closest runnable reference; do not silently omit the comparison.
 - Match inputs or dataset, seed, output settings, precision or quantization,
  hardware, software, warmup, repetitions, synchronization, and timing and
  memory scope. Include exact commands and environment details for both sides.
 - Compare accuracy or output quality with a task-appropriate metric and explicit
  tolerance. When no reliable metric exists, provide paired fixed-input outputs,
  artifacts, and evaluation criteria, and explain expected numerical or
  qualitative differences.
 - Report end-to-end latency, task-appropriate throughput, peak device memory,
  run-to-run variability, and head/reference ratios at equivalent output quality.
 - Prefer a timing split for preprocessing or encoding, each model stage,
  inter-stage transfer or communication, decoding or postprocessing, and output
  assembly. Define the boundaries, account for overlap, reconcile the split with
  end-to-end time, and explain any unavailable stage measurement.
 Use [perf-verification.md](perf-verification.md) for the comparison protocol and
 repository benchmark and profiling links.
 ## Remove accidental surface
 Search bounded call sites for:
 - inference-dead training `forward()` paths, unused factories/wrappers, never-set
  branch keys, and immediately discarded parameters;
 - duplicate payload strings, validation, or shape coercion across stages;
 - compatibility aliases in brand-new code without an existing caller;
 - private symbols re-exported as public API or module-level side effects.
 Keep code only when it has a distinct live caller, invariant, or compatibility
 contract. Prefer one typed producer-consumer schema to repeated string keys.
 ## Require proportionate evidence
 - Run a representative production-path inference and assert output content, not
  only process survival.
 - Require profiling or A/B tables when the PR makes performance, memory,
  precision, or quality claims, or when a suspected hot-path/device bug needs
  that evidence. Do not impose fixed utilization or regression thresholds that
  the repository does not define.
 - Require focused registry/config/loading tests, a representative E2E path, and
  user-facing model/usage documentation.
 Use [diffusion.md](../modules/diffusion.md) for diffusion models and
 [perf-verification.md](perf-verification.md) for quantitative claims.
@@ -1,51 +0,0 @@
 # Performance and Accuracy Verification
 Activate this reference when the PR adds a model or claims or intentionally
 changes latency, throughput, memory, scaling, precision, or output quality.
 Official docs: [profiling](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/profiling/),
 [serving benchmarks](https://docs.vllm.ai/projects/vllm-omni/en/latest/cli/bench/serve/),
 and [metrics](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/metrics/).
 ## Comparable A/B contract
 Compare frozen base and head for an existing path. For a new model absent from
 base, compare head with a pinned canonical reference implementation. Use the same:
 - hardware, driver, software versions, model/checkpoint, and dependencies;
 - input or dataset, seed, precision, batch/concurrency, topology, and feature
  flags;
 - warmup, measured repetitions, synchronization, timing scope, and memory
  collection method.
 Keep exact commands and report variability rather than a cherry-picked run.
 Use the repository's benchmark when it covers the claim; otherwise use the
 smallest reproducible workload that reaches the changed production path.
 ## Evidence table
 | Dimension | Typical evidence |
 | --- | --- |
 | Latency | End-to-end latency and, when available, TTFT plus a defined per-stage and transfer-time split. |
 | Throughput/scaling | Requests, tokens, frames, or audio duration per second across the claimed concurrency/topology. |
 | Memory | Peak allocated/reserved device memory and OOM boundary if relevant. |
 | Quality/accuracy | Repository metric or known-good output comparison with an explicit tolerance; paired samples when no metric exists. |
 Every optimization needs correctness or quality evidence. A faster result on a
 different workload, precision, seed, or topology is not a valid comparison.
 ## Reviewer verification ladder
 1. Run base/head A/B on suitable hardware when affordable.
 2. If only one side can run, verify the script and request a comparable pair.
 3. If hardware or assets are unavailable, audit methodology, code path, and
   contributor evidence; name the exact unverified claim.
 Classify discrepancies before reporting: implementation regression, benchmark
 bug, environmental drift, noise, or unsupported claim. Explain material
 regressions against the PR's stated goal or repository contract; do not invent
 universal percentage thresholds.
 Report claimed and measured values together, bind them to the frozen SHAs and
 environment, and keep unavailable hardware as a validation gap rather than a
 fabricated pass.
@@ -1,39 +0,0 @@
 # Test Quality Evaluation
 Read this reference when tests change, a risky production path has no test, or
 a test-only PR may not prove the intended behavior.
 Official docs: [test system overview](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_system_overview/)
 and [test writing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/).
 ## Static proof check
 For each changed semantic path, ask whether the test would fail if that behavior
 were reverted or broken. Check:
 - assertions pin the contract, not only non-null output, logs, or “no crash”;
 - the production dispatcher, registry, connector, or scheduler is reached
  instead of the changed behavior being mocked away;
 - fakes preserve relevant types, MRO, shapes, devices, async behavior, and
  lifecycle transitions;
 - seeds, ordering, timing, synchronization, external services, and numeric
  tolerances are deterministic or explicitly controlled;
 - normal, invalid, boundary, feature-off, failure/cancellation, and regression
  paths are covered where the diff changes them;
 - run-level and domain markers place hardware/model tests in the intended CI
  lane without silently skipping the contract.
 Map source symbols to tests with bounded `rg` searches; do not assume the test
 tree mirrors production paths.
 ## Runtime check
 Run the narrowest affected tests supported by the environment. Record skipped
 hardware/model cases as gaps. Classify failures as code, test, infrastructure,
 or flaky before turning them into findings; one passing rerun does not erase a
 flakiness signal.
 Keep grades or matrices internal. Report only concrete code bugs or the one or
 two test defects that leave changed behavior materially unprotected. Use
 [tests-docs-checklist.md](tests-docs-checklist.md) for CI and documentation
 coverage.
@@ -1,48 +0,0 @@
 # Tests and Documentation Checklist
 Use this reference to connect changed behavior to coverage, CI selection,
 examples, documentation, and reproducible PR evidence.
 Official docs: [test writing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/),
 [test execution guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_execution_guide/),
 and [documentation guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/DOCS_GUIDE/).
 ## Coverage packet
 For each high-risk change, record a compact internal packet:
 ```text
 change and failure risk -> existing unit/E2E coverage -> uncovered boundary
  -> smallest test and stable assertion that closes the gap
 ```
 Prefer a regression test that fails on the frozen base for bug fixes. Check
 numeric tolerances, deterministic inputs, realistic mocks, hardware/domain
 markers, and Buildkite/CI wiring using the documentation present on the reviewed
 branch. A skipped test is not a pass.
 ## Documentation sync
 Require user-facing updates when the diff changes a model, feature, CLI/API,
 config key, default, compatibility behavior, or supported platform. Verify:
 - the actual docs navigation and support tables used by the branch;
 - exact identifiers, defaults, constraints, known limitations, and migration;
 - a minimal runnable example for newly exposed behavior;
 - examples and docs agree on commands, parameters, output schema, and supported
  modes;
 - links and the bounded live contract build or resolve correctly.
 Do not request docs for internal behavior with no user-visible contract.
 ## PR evidence
 For contributor-run tests or benchmarks, require enough provenance to evaluate
 the claim: hardware, relevant software versions, exact commands, inputs/model,
 result, and approximate runtime. Performance or accuracy evidence must compare
 base and head under equivalent conditions; use
 [perf-verification.md](perf-verification.md).
 Turn a missing test or document into a finding only when it protects or explains
 changed behavior. Name the smallest useful addition rather than requesting a
 broad coverage expansion.
@@ -1,36 +0,0 @@
 # Active Verification
 Read this reference when a runnable affected path and suitable local or remote
 hardware are available. Never claim device evidence from static analysis or a
 simulated backend.
 Official docs: [installation](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/)
 and [test execution guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_execution_guide/).
 ## Choose the narrowest level
 | Available environment | Verification |
 | --- | --- |
 | Matching accelerator, model/assets, and server | Targeted unit/integration test plus one representative E2E request. |
 | Matching accelerator without serving setup | Targeted model/runtime test and output or metric inspection. |
 | CPU/static environment only | Import/version preflight, focused CPU tests, config/registry resolution, and static contract checks. |
 | Missing dependencies or inaccessible hardware | Record the exact gap and evaluate contributor/CI evidence. |
 ## Execute safely
 1. Bind the run to the frozen head SHA and record Python, platform, dependency,
   model, device, and relevant environment fingerprints.
 2. Run the smallest affected unit test before an E2E workload.
 3. Exercise the changed public or production-dispatch path with bounded inputs.
 4. Compare actual output, metrics, cleanup, and exit behavior with the PR claim.
 5. For base/head claims, use the equivalent A/B protocol in
   [perf-verification.md](perf-verification.md).
 Avoid disturbing shared servers, ports, models, or accelerators without user
 authorization. Do not install large dependencies, download gated assets, or
 start remote workloads merely because hardware might exist.
 Report the command, environment, result, and any skipped path. Classify failures
 as change-induced, pre-existing, infrastructure, environment, or flaky before
 making a finding. A hardware gap lowers confidence; it is not automatic proof
 of a defect.
@@ -1,105 +0,0 @@
 # Maintainer Review Style
 Use this reference only when turning proved findings into comments. The baseline
 maintainer samples favor short, direct comments and reserve long explanations
 for architecture or non-obvious failure paths.
 Official docs: [contributing and review process](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
 ## Evidence basis
 This calibration used 928 PRs merged from 2026-05-10 through 2026-08-08.
 Candidates were ranked by substantive human inline-review threads, then by
 distinct reviewers and review rounds. Author self-comments, bots and AI
 reviewers, and bodyless approvals were excluded. Counts selected the sample;
 the practices below came from reading root comments and follow-up replies.
 The main set contains the six leaders after filtering; three other high-ranked
 PRs add core-runtime and benchmark coverage. All were merged.
 | PR | Review boundary |
 | --- | --- |
 | [#2906](https://github.com/vllm-project/vllm-omni/pull/2906) | New TTS model, shared utilities, recipes, and request state. |
 | [#3907](https://github.com/vllm-project/vllm-omni/pull/3907) | Full-duplex API, session lifecycle, concurrency, and realtime evidence. |
 | [#2162](https://github.com/vllm-project/vllm-omni/pull/2162) | Diffusion world model, OpenPI serving, tests, examples, and docs. |
 | [#3454](https://github.com/vllm-project/vllm-omni/pull/3454) | Diffusion model, generic CLI surface, loaders, and configuration scope. |
 | [#4425](https://github.com/vllm-project/vllm-omni/pull/4425) | Structured configuration, typing, defaults, and compatibility. |
 | [#5397](https://github.com/vllm-project/vllm-omni/pull/5397) | Distributed offload, collectives, failure handling, and load semantics. |
 | [#4281](https://github.com/vllm-project/vllm-omni/pull/4281) | Composable parallelism, configuration contracts, and examples. |
 | [#3628](https://github.com/vllm-project/vllm-omni/pull/3628) | Stage metrics, benchmark definitions, opt-in API behavior, and cleanup. |
 | [#3855](https://github.com/vllm-project/vllm-omni/pull/3855) | Stage runtime, thread safety, startup, and compatibility. |
 Historical comments calibrate delivery; they do not create merge policy. Live
 code, tests, repository docs, and the current `CODEOWNERS` file remain
 authoritative. When refreshing this study, promote a practice only when it
 recurs across at least three PRs or is supported by current repository guidance.
 ## Recurring review behavior
 - Start from a changed contract, not a style preference. State the exact
  request, configuration, topology, or failure path that triggers the issue.
 - Protect repository boundaries. Keep model-specific behavior in its domain,
  move genuinely reusable logic to a shared module, and reuse existing loaders,
  protocols, definitions, and config sources instead of bypassing them.
 - Challenge unrelated scope and abstractions. Ask why a new flag, config, class,
  helper, fallback, or file is needed when an existing path appears sufficient.
 - Trace concurrency and terminal paths: overlapping requests, disconnect,
  cancellation, partial failure, stale replies, collective schedule mismatch,
  cleanup, and retry. Prefer an explicit startup error for unsupported
  combinations over silently running with incomplete state or wrong weights.
 - Preserve user-visible behavior across variants. An accepted CLI/API/config
  field must be applied, explicitly rejected, or visibly warned about; compare
  normal and headless, offline and online, and streaming and non-streaming paths.
 - Ask for evidence at the affected boundary: a regression test using the real
  tensor or protocol shape, correct pytest markers so CI runs it, comparable A/B
  performance results, or an SLO-oriented latency/quality measurement.
 - Keep interfaces explicit: avoid `Any` and opaque strings where a precise
  union, enum, or typed envelope exists; remove dead code and duplicated
  constants; validate requests before expensive execution.
 - Document module scope and non-obvious invariants. Update examples, recipes,
  supported-model/feature docs, and hardware assumptions when users rely on
  them.
 - When expertise crosses an ownership boundary, ask a named owner to inspect a
  specific file or contract and explain why. Do not send a generic `PTAL` list.
 - Keep non-blocking cleanup in a linked follow-up only after the current PR is
  safe. If a revision fixes only part of a failure chain, continue the existing
  thread with the remaining trigger instead of opening a duplicate finding.
 ## Write the finding
 - Lead with the issue or decisive question; omit a review preamble.
 - Name the concrete trigger and impact, then the smallest fix direction.
 - Keep obvious fixes to one sentence. Use more detail only when the call path or
  ownership boundary is otherwise unclear.
 - Prefer one root-cause comment over repeated symptom comments.
 - Hedge only when evidence is genuinely incomplete; mark that item as a
  question or validation gap rather than overstating severity.
 - Say briefly that no findings were found when the review is clean; zero
  comments is a valid outcome when posting is authorized.
 Useful forms:
 ```text
 This leaves <state> allocated when <failure/cancellation path>. Release it in
 the shared terminal cleanup path.
 ```
 ```text
 Does this producer also update <consumer>? It still reads the old field here.
 ```
 ```text
 With <configuration>, this bypasses <existing loader/validator>, so <impact>.
 Reuse that path or reject this combination before startup.
 ```
 ```text
@<owner>, could you check <specific boundary>? This changes <contract> used by
 <consumer>.
 ```
 Use a priority label only for a proved high-impact defect. Avoid generic praise,
 routine “Nit:” prefixes, decorative severity labels, audit templates, rule IDs,
 comment-count narration, and “I left a few comments.” Keep approvals, review
 events, reviewer requests, mentions, and GitHub replies subject to the user's
 explicit authorization.
@@ -1,99 +0,0 @@
 # Focused Code-Owner Review Requests
 Read only when the user asks to identify, suggest, request, add, or ping
 reviewers. Finding candidates is read-only; GitHub review requests and
 `@mention` comments are external writes.
 Owner sources, in order:
 1. the reviewed head's
   [CODEOWNERS](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS)
   for path ownership;
 2. the active module design page's `owners` and `required_reviewers` metadata
   for contract ownership and mandatory design review; and
 3. documented
   [governance expertise](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/governance/)
   for specialization and tie-breaking.
 Resolve all three from the reviewed head when available because ownership
 changes over time. Use published docs only to discover a page absent from an
 older checkout, never to override the frozen head, and do not copy a static
 owner list into review output. CODEOWNERS remains the path authority; module
 metadata identifies contract experts and required design reviewers.
 ## Rank candidates
 1. Freeze the PR head and collect changed files, author, requested reviewers,
   submitted reviews, and existing comments.
 2. Resolve the **last matching CODEOWNERS rule** for each changed file. Do not
   union an earlier broad rule with a later specific rule.
 3. Resolve the primary module and feature overlays through
   [design-contracts.md](../process/design-contracts.md). Read the active module
   page's status, `owners`, and `required_reviewers`; do not treat owners on an
   archived, superseded, or unrelated page as current authority.
 4. Group files by the selected primary owner and the exact module or feature
   contract needing review. Weight production owners and live consumers above
   tests/docs/CI owners and the repository-wide fallback. Include an applicable
   `required_reviewers` entry even when it is not the path owner, and label why.
 5. Use documented governance specialties only to choose among matching owners:
   for example serving versus configuration, diffusion model versus cache or
   parallelism, TTS versus generic model execution, or CI versus tests.
 6. Exclude the PR author, bots, users who already completed the needed review,
   and duplicate pending requests. Do not infer availability or invent
   expertise from a username.
 7. Select one reviewer for a single-owner change and at most three for a real
   cross-module change. If more are plausible, rank them but do not ping a
   broad group; ask the user before exceeding three.
 If every specific owner is excluded, consider one repository-wide fallback
 owner and label that choice as a fallback rather than a path match. Never fan
 out to the full fallback list. Prefer a qualified owner who covers multiple
 critical path groups over separate reviewers for each file. If CODEOWNERS is
 missing or invalid, state that gap and use governance only as a labeled fallback.
 Recent file history may break a tie only when CODEOWNERS and governance remain
 ambiguous. Never use commit count alone to override an explicit current owner.
 ## Preview before writing
 Show the user the exact candidate set and rationale:
 ```text
@owner — <owned paths or module role>; please verify <specific module or feature contract>.
 ```
 Draft one consolidated comment tied to the frozen head:
 ```text
 Requesting focused review for `<short-head-sha>`:
 - @owner1: could you check the <module contract> in `<path/group>`?
 - @owner2: could you check the <feature or cross-boundary contract> in `<path/group>`?
 Relevant validation: <one concise result or exact remaining gap>.
 ```
 Mention only owners being requested. Keep the reason specific enough that each
 person knows what decision or invariant needs attention; do not post a generic
 “please review” or a generated ownership audit.
 ## External-write gate
 - “Who should review?”, “suggest reviewers”, or “draft a request” authorizes
  only the preview.
 - “Request/add these reviewers” authorizes the GitHub review-request mutation.
 - “Ping/comment/tag the owners” authorizes the consolidated comment.
 - When the user explicitly asks to request reviews **with an owner comment**,
  both writes are authorized. Obey narrower wording when only one action is
  requested.
 Immediately before an authorized write, verify the PR is still open, the head
 SHA is unchanged, and the candidate is not the author or already satisfied.
 Report draft/WIP state and avoid notifications unless the user explicitly wants
 reviewers requested before ready-for-review.
 Use the GitHub connector for reviewer requests and the single top-level comment;
 use `gh` only when connector coverage is unavailable. Do not submit an approval,
 request-changes event, label, assignment, or multiple comments as part of this
 action. After writing, reread the PR and report the reviewers requested and the
 comment URL. If the head changed, discard the draft and rerank before posting.
@@ -1,24 +0,0 @@
 # Communication Features
 Use with the OmniConnector, stage-runtime, configuration, and I/O module
 contracts. Load the concrete backend design from the reviewed head:
 | Backend | Design |
 | --- | --- |
 | Mooncake store | [MooncakeStoreConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mooncake_store_connector/) |
 | Mooncake transfer engine | [MooncakeTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mooncake_transfer_engine_connector/) |
 | Mori transfer engine | [MoriTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mori_transfer_engine_connector/) |
 | Shared memory | [SharedMemoryConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/shared_memory_connector/) |
 | Yuanrong store | [YuanrongConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/yuanrong_connector/) |
 | Yuanrong transfer engine | [YuanrongTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/yuanrong_transfer_engine_connector/) |
 ## Feature checks
 - Verify dependency, service, device/network, port, memory-pool, and environment
  prerequisites with actionable startup errors.
 - Match key/metadata construction, sender and receiver roles, allocation,
  ownership, completion, and supported fast/serialized paths.
 - Exercise timeout, retry, stale sender/buffer reclamation, partial transfer,
  cancellation, and connector close semantics.
 - Validate the smallest supported same-host or cross-node topology and document
  backend limitations and operational diagnostics.
@@ -1,26 +0,0 @@
 # Diffusion Acceleration Features
 Use with the exact diffusion submodule plus platform, cache, quantization, or
 benchmarking contracts selected by the live path.
 | Feature family | Designs |
 | --- | --- |
 | Parallelism | [CFG](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel/), [Expert](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/expert_parallel/), [HSDP](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp/), [Pipeline](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/pipeline_parallel/), [Sequence](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel/), [Tensor](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel/), [VAE patch](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/vae_parallel/) |
 | Attention | [Skip-Softmax](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/skip_softmax/) |
 | Quantization | [Quantization](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/quantization/) |
 | Cache | [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/teacache/) |
 | Batching | [Diffusion Continuous Batching](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/diffusion_continuous_batching/) |
 | Offload | [Distributed Layerwise Offload](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/distributed_layerwise_offload/) |
 ## Feature checks
 - Verify feature selection, supported models/platforms/topologies, defaults,
  dependencies, and every touched feature combination.
 - Preserve feature-off and single-rank correctness; reject unsupported
  combinations before collectives, loading, or execution begin.
 - Trace shapes, shards, ranks, collectives, precision, cache validity, residency,
  and per-request state through the selected runtime hooks.
 - Compare output quality or accuracy against the disabled/reference path before
  accepting latency or memory improvements.
 - Require exact commands, topology, environment, warmup, repetitions, and
  measurements for performance claims; update feature compatibility docs.
@@ -1,22 +0,0 @@
 # Infrastructure and Performance Features
 Load the matching design from the reviewed head:
 - [Prometheus Metrics](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/metrics/)
  for logger wrapping, gating, lifecycle, stage transfer, and metric definitions.
 - [Speech Generation Performance Optimizations](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/qwen3_omni_tts_performance_optimization/)
  for stage batching, CUDA graphs, async chunk, streaming, re-prefill, compile,
  and deployment evidence.
 ## Feature checks
 - Route instrumentation changes through observability or profiling and result
  methodology through benchmarking and the performance evidence check.
 - Preserve opt-in/default behavior, units, labels, aggregation, timing scope,
  lifecycle resets, and bounded overhead.
 - For stacked optimizations, measure each increment against the same baseline
  and explain interactions, quality guards, regressions, and limitations.
 - Require exact model, checkpoint, hardware, software, workload, topology,
  concurrency, warmup, repetitions, and production commands.
 - Update the design page when metric meaning, supported optimization stack, or
  deployment guidance changes.
@@ -1,24 +0,0 @@
 # Runtime and Stage Execution Features
 Load the matching design from the reviewed head:
 | Feature | Design | Primary module intersections |
 | --- | --- | --- |
 | Disaggregated inference | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) | Configuration, stage runtime, OmniConnector, I/O contracts |
 | Async chunk | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/async_chunk/) | Model integration, orchestration, stage runtime, connector, I/O |
 | Async diffusion output | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/async_diffusion_output/) | Diffusion runtime, I/O contracts, entrypoints |
 | Async Omni output materialization | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_async_output_materialization/) | AR runtime, model integration, I/O contracts, orchestration |
 | Automatic prefix caching | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/prefix_caching/) | AR runtime, cache management, I/O contracts |
 ## Feature checks
 - Confirm the feature's supported models, entrypoints, stage types, topology,
  configuration, default, and compatibility limitations.
 - Trace enablement from public/config ingress through every participating module
  and verify the advertised production dispatch actually selects it.
 - Preserve a correct disabled or fallback path and make unsupported combinations
  fail or warn as documented.
 - Test overlap, ordering, identity, backpressure, cancellation, failure, final
  partial output, and cleanup at each asynchronous boundary.
 - Update the feature page, support matrices, examples, and recipes when behavior
  or coverage changes.
@@ -1,22 +0,0 @@
 # Autoregressive Runtime
 Primary design: [Autoregressive Runtime](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/ar_runtime/).
 Use for AR request queues, scheduling state, token and memory budgets, KV/input
 readiness, workers, runners, and upstream vLLM request/cache transitions.
 ## Contract checks
 - Preserve upstream request-state, scheduling, cache, and preemption semantics
  unless an Omni-specific difference is documented and tested.
 - Convert modality-specific data through explicit input or output adapters; do
  not inject it through unrelated scheduler state.
 - Keep workers and model runners focused on assigned execution rather than
  admission or cross-stage routing.
 - Synchronize token, sequence, batch, and memory budgets with worker-visible
  state and prevent starvation, double counting, or stale readiness.
 - Preserve legal transitions and one terminal cleanup path across completion,
  abort, timeout, preemption, worker failure, and shutdown.
 Test request lifecycle, scheduling decisions, abort, cache state, adapters, and
 every affected worker execution mode against the pinned upstream contract.
@@ -1,24 +0,0 @@
 # Benchmarking
 Primary design: [Benchmarking](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/benchmarking/).
 Use when benchmark code, workloads, metric calculations, CLI benchmark
 entrypoints, or result metadata change. Use the performance check separately
 when a PR uses benchmark results as evidence.
 ## Contract checks
 - Record model and revision, code SHA, hardware, software, configuration,
  workload, precision, topology, concurrency, warmup, repetitions, and
  measurement window.
 - Validate successful and semantically acceptable outputs before accepting
  timing, memory, or scaling results.
 - Define every metric's unit, population, aggregation, synchronization, and
  timing scope.
 - Keep workload parsing and metric computation deterministic and reusable
  across base and head.
 - Reject comparisons that change material workload or runtime settings without
  isolating and explaining that difference.
 Test workload parsing and metric calculations with deterministic fixtures and
 publish exact commands plus variability for quantitative claims.
@@ -1,22 +0,0 @@
 # Cache Management
 Primary design: [Cache Management](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/cache_management/).
 Use for AR or diffusion cache identity, reusable state, validity, lifecycle,
 reset, eviction, and teardown.
 ## Contract checks
 - Include every model, request, input, timestep, precision, device, topology,
  and feature property that can change a reused value in cache identity.
 - Preserve a correct disabled and miss path; never reuse after validity inputs
  change.
 - Isolate request-scoped state across concurrent requests and define ownership
  for model-scoped state.
 - Make reset, invalidation, eviction, capacity, and teardown explicit on
  completion, cancellation, failure, model reload, and shutdown.
 - Keep cache optimizations behind the owning runtime's hooks rather than adding
  a second request lifecycle.
 Test hit, miss, invalidation, disabled behavior, concurrency, capacity, reset,
 and teardown, with quality comparison when reuse is approximate.
@@ -1,24 +0,0 @@
 # vLLM-Omni Configuration
 Primary design: [vLLM-Omni Configuration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/vllm_omni_config/).
 Use for config construction, deploy YAML, registries, schema, defaults,
 aliases, CLI projection, stage config, and topology. Inspect the reviewed
 page's status: while its contract remains deferred or draft, current code and
 tests—not proposed precedence rules—are authoritative.
 ## Contract checks
 - Resolve each supported structured, legacy, CLI, environment, default, alias,
  and stage-override producer through its actual runtime consumer.
 - Reject unknown or owner-mismatched fields explicitly; keep process-local
  runtime objects out of transportable configuration.
 - Verify default, explicit, and feature-off values reach every live factory and
  consumer without silent reinterpretation.
 - Treat stage count, placement, connector, parallelism, device, and memory
  fields as one topology contract and reject impossible combinations early.
 - Preserve supported construction-path parity until a documented migration
  removes a path.
 Require normalization and schema tests, a live-consumer assertion, and docs for
 public keys, defaults, constraints, precedence, or migration behavior.
@@ -1,33 +0,0 @@
 # Diffusion Module Family
 Primary design: [Diffusion Modules](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/).
 Use for diffusion-stage behavior, then load the exact submodule page from the
 reviewed head:
 | Submodule | Signals | Design |
 | --- | --- | --- |
 | Runtime | Admission, scheduling, execution, progress, outputs, cancellation, hooks, cleanup | [Diffusion Runtime](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/diffusion_runtime/) |
 | Model Integration | Pipelines, registry, loader, adapters, shared layers, model-specific processing | [Diffusion Model Integration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/diffusion_model_integration/) |
 | Continuous Batching | Compatibility, request/step batches, per-request progress, capacity | [Continuous Batching](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/continuous_batching/) |
 | Parallelism | Rank topology, process groups, sharding, collectives, distributed execution | [Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/parallelism/) |
 | Offloader | Residency, transfers, memory accounting, prefetch, teardown | [Offloader](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/offloader/) |
 ## Family-wide checks
 - Keep one scheduler-owned lifecycle from admission through completion,
  cancellation, failure, and cleanup; optional features must use defined hooks.
 - Select pipelines through registry and loader contracts instead of scattered
  model-name conditionals; keep genuine model differences in model directories.
 - Trace latent, conditioning, timestep, generator, shape, layout, dtype, device,
  batch expansion, and output conversion through the actual consumer.
 - Make batch compatibility and per-request isolation explicit; never associate
  output by unstable batch position.
 - Derive topology from validated configuration, define every sharding boundary,
  keep collectives symmetric, and preserve a supported single-rank path.
 - Track one residency owner per offloaded component, wait for transfer readiness,
  preserve model state, bound retained copies, and clean up deterministically.
 Load matching feature designs for cache, quantization, parallelism, continuous
 batching, or offload. Require feature-off correctness, representative inference,
 terminal-path tests, and quality evidence for optimizations.
@@ -1,23 +0,0 @@
 # Engine Orchestration
 Primary design: [Engine Orchestration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/engine_orchestration/).
 Use for request-state creation, cross-stage routing, output ordering, companion
 tracking, RPC correlation, cancellation propagation, and terminal convergence.
 Orchestration does not own replica placement, stage-process startup, public
 rendering, connector transport, or payload and error schemas.
 ## Contract checks
 - Keep one owner for cross-stage routing; entrypoints, models, and stage clients
  must not independently select the next logical stage.
 - Preserve request identity and output order across concurrent stages and
  correlate every control or RPC result to the originating request and wave.
 - Make terminal state monotonic: no new forwarding after success, abort,
  request failure, or fatal engine failure.
 - Propagate cancellation and failure to every active stage and release
  orchestrator-owned tasks, queues, companions, and request state.
 - Preserve shutdown ordering across the orchestrator and stage runtime.
 Test representative multi-stage routing, ordering, correlation, cancellation,
 failure, and repeated shutdown.
@@ -1,25 +0,0 @@
 # Entrypoints and Serving Boundaries
 Primary design: [Entrypoints and Serving Boundaries](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/entrypoints/).
 Use for offline APIs, CLI and server composition, OpenAI-compatible routes,
 request validation, response conversion, streaming, sessions, and engine
 handoff. Entrypoints adapt public contracts; they do not own cross-stage
 routing, stage lifecycle, configuration precedence, payload schemas, or error
 classification.
 ## Contract checks
 - Normalize and validate public values once before engine submission.
 - Preserve request identity, modality, ordering, and exactly one terminal event
  across streaming, disconnect, cancellation, and failure.
 - Translate internal outputs and errors explicitly for each advertised offline,
  HTTP, SSE, WebSocket, or asynchronous-job path.
 - Keep model-specific behavior behind a common adapter or processor instead of
  adding model-name policy to generic entrypoints.
 - Require API or CLI changes to explain necessity, alternatives, compatibility,
  migration, and the supported route or transport matrix.
 Test the affected public protocol, normalization, engine handoff, rendering,
 and terminal outcomes. Route orchestration, I/O schema, and semantic error
 changes to their module contracts.
@@ -1,24 +0,0 @@
 # Error Classification, Propagation, and Rendering
 Primary design: [Error Classification, Propagation, and Rendering](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/error_contracts/).
 Use when a change classifies failures, crosses process boundaries, decides
 fatal versus request-scoped impact, sanitizes context, or renders offline,
 HTTP, SSE, WebSocket, or asynchronous-job errors. The I/O module owns the wire
 schema; this module owns the meaning of its error fields.
 ## Contract checks
 - Keep request-scoped failure isolated unless the underlying engine is no
  longer usable.
 - Preserve stable classification and request context across process and
  protocol boundaries without exposing prompts, outputs, secrets, or internal
  stack details.
 - Keep cancellation distinguishable from internal failure.
 - Terminate each public transport exactly once with consistent status, body,
  event, finish reason, and metrics classification.
 - Maintain compatibility across mixed or older producers only when the
  reviewed contract defines that behavior.
 Test table-driven classification, wire round trips, fatal/request-scoped
 propagation, cancellation, sanitization, and every affected public renderer.
@@ -1,23 +0,0 @@
 # Execution Platforms
 Primary design: [Execution Platforms](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/execution_platforms/).
 Use for hardware capability detection, worker selection, patches, kernels,
 attention backends, and vendor-specific configuration or connectors.
 ## Contract checks
 - Guard hardware behavior through platform selection or capability detection
  and fail clearly when unsupported.
 - Keep portable modules free of direct vendor imports when a common selector
  exists; verify portable imports without the target accelerator.
 - Override only behavior that differs from common or upstream implementations,
  and compare sibling worker hooks when inheritance does not provide parity.
 - Preserve dtype, layout, rank, stream, synchronization, workspace,
  graph-capture, and resource-lifetime contracts.
 - Do not let simplified CPU doubles hide platform MRO, initialization, or
  dispatch failures.
 Run portable checks first, then the smallest matching-device test in a fresh
 environment. Report unavailable hardware as a validation gap, never simulated
 device evidence.
@@ -1,23 +0,0 @@
 # Input, Output, and Modality Contracts
 Primary design: [Input, Output, and Modality Contracts](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/input_output_modality_contracts/).
 Use for requests, prompt and modality types, message schemas, serialization,
 output types, accumulation, completion, and compatibility shims crossing
 entrypoint, orchestration, stage, connector, or model boundaries.
 ## Contract checks
 - Give every boundary value an explicit modality and validated typed schema.
 - Preserve request and stage identity across conversion, serialization,
  streaming updates, batching, cancellation, and errors.
 - Distinguish partial, final, empty, and failed output states for every modality.
 - Verify producer and consumer agreement on shape, layout, dtype, device,
  ownership, metadata, ordering, and optional or unknown fields.
 - Keep internal objects out of public protocols; translate them at the owning
  entrypoint.
 - Preserve wire compatibility deliberately and give deprecated shims a bounded
  migration path.
 Require construction, round-trip, accumulation, completion, rejection, and
 compatibility tests at each changed producer-consumer boundary.
@@ -1,25 +0,0 @@
 # Model Integration
 Primary design: [Model Integration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/model_integration/).
 Use for model registration, configuration adapters, preprocessing, checkpoint
 loading, stage inputs, model-specific hooks, shared layers, runners, and model
 execution. Use the diffusion family reference for diffusion models.
 ## Contract checks
 - Declare explicitly how the model class, loader, input processor, stage
  configuration, and runtime capabilities are selected.
 - Preserve weight-name mapping, dtype/device placement, optional-dependency
  errors, and checkpoint-specific loading semantics.
 - Trace each typed input and output through preprocessing, runner state, model
  call, and its actual consumer.
 - Reuse upstream vLLM behavior when sufficient; document and test the precise
  difference for every override or patch.
 - Keep model code out of admission, cross-stage routing, public policy, and
  terminal lifecycle ownership.
 - Place reusable behavior in shared modules and test it with more than one
  integration.
 Require focused registry, loader, processor, runner, and representative model
 execution tests. Apply the model-addition check for new or expanded support.
@@ -1,21 +0,0 @@
 # Observability
 Primary design: [Observability](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/observability/).
 Use for metrics, logs, correlation fields, and cross-stage request visibility.
 ## Contract checks
 - Preserve stable request and stage identity across entrypoints, orchestration,
  stages, workers, errors, and terminal cleanup.
 - Treat metric names, units, labels, aggregation, and lifecycle as public
  contracts; document any semantic change.
 - Keep labels bounded: never use request IDs, prompts, outputs, paths, or other
  unbounded user values.
 - Register and reset metrics exactly once across repeated startup and shutdown,
  and keep disabled instrumentation inexpensive.
 - Align emitted values with their actual timing and lifecycle owner rather than
  reconstructing them from ambiguous client strings.
 Test registration, label sets, units, resets, error classification, and
 representative multi-stage emission.
@@ -1,24 +0,0 @@
 # OmniConnector
 Primary design: [OmniConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/omni_connector/).
 Use for connector contracts and model-agnostic data transfer across stages,
 processes, devices, or nodes. Connector implementations transport and
 synchronize data; they do not choose the next stage or implement model policy.
 ## Contract checks
 - Prove producer and consumer agreement on request/stage identity, shape,
  layout, dtype, placement, ownership, ordering, and completion.
 - Keep backend selection and metadata/key construction explicit; do not silently
  reinterpret or drop payloads across SHM, store, network, or vendor paths.
 - Bound connect, put/get, retry, synchronization, and readiness waits with
  backpressure, cancellation, and actionable errors.
 - Prevent route, endpoint, port, key, and buffer collisions across repeated
  startup, concurrent requests, and replica replacement.
 - Release connections, handles, buffers, listener work, and stale sender state
  on success, cancellation, failure, and shutdown.
 Load the matching communication feature design for a concrete backend. Require
 focused producer-consumer tests and the smallest runnable multi-process or
 multi-node path.
@@ -1,21 +0,0 @@
 # Profiling
 Primary design: [Profiling](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/profiling/).
 Use for opt-in instrumentation and trace collection across supported profiler
 backends.
 ## Contract checks
 - Keep expensive collection disabled by default and bounded by explicit start,
  stop, timeout, and cleanup behavior.
 - Preserve scheduling, generated outputs, resource ownership, and concurrency
  semantics apart from documented synchronization overhead.
 - Include stage, rank, worker, operation, and request correlation where useful
  without embedding user payloads or secrets.
 - Make repeated start/stop and partial failure safe, including multi-rank trace
  naming and artifact ownership.
 - Measure disabled and enabled overhead before making performance claims.
 Test the disabled path, repeated lifecycle, failure cleanup, trace context, and
 each affected backend.
@@ -1,22 +0,0 @@
 # Quantization
 Primary design: [Quantization](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/quantization/).
 Use for method registration, checkpoint interpretation, layer replacement,
 reduced precision, platform capability, and runtime constraints.
 ## Contract checks
 - Validate checkpoint metadata, selected method, supported layers, device
  capability, required dependencies, and parallel layout before loading.
 - Preserve loader and weight-name mapping semantics, dtype, scale/granularity,
  zero point, padding, and buffer handling through the real kernel consumer.
 - Fail clearly for unsupported combinations; do not silently choose a different
  method or precision.
 - Keep online and pre-quantized checkpoint paths explicit and ensure patches or
  fallbacks do not bypass model-owned loaders.
 - Require numerical comparison against a defined reference with explicit
  tolerance before treating performance results as valid.
 Test registration, loading, representative layers, metadata, platform guards,
 parallel combinations, and numerical correctness separately from performance.
@@ -1,23 +0,0 @@
 # Stage Runtime and Replica Lifecycle
 Primary design: [Stage Runtime and Replica Lifecycle](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/stage_runtime/).
 Use for local or distributed placement, startup, readiness, stage pools,
 replica identity and membership, selection, affinity, liveness, draining, and
 stage-client or process lifecycle. Cross-stage request policy belongs to engine
 orchestration; model scheduling belongs to the owning runtime.
 ## Contract checks
 - Acquire, replace, and retire replicas through one stage-runtime lifecycle.
 - Preserve stable replica identity and request affinity until termination or a
  declared replica loss.
 - Make readiness explicit before admission; surface partial-startup and
  membership failures without leaving half-initialized capacity.
 - Keep local and distributed semantics aligned for supported modes.
 - Make drain, shutdown, and cleanup idempotent across repeated signals and
  process failure.
 Test startup, readiness, selection, affinity, membership changes, abort,
 process loss, draining, and repeated shutdown on affected local/distributed
 paths.
@@ -1,34 +0,0 @@
 # Resolve Design Contracts
 Use after freezing the review head and before selecting module or feature
 references. The maintained [design index](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/)
 defines two routing axes:
 - **Module designs** own code boundaries and invariants. Select one primary
  module and a second only when the live producer-consumer path crosses it.
 - **Feature designs** describe behavior spanning modules. Load every feature
  page whose implementation, default, support, or contract changes.
 ## Resolution order
 1. Open `docs/design/index.md` in the reviewed head, not the reviewer's checkout
   or published latest tree.
 2. Match changed behavior against each module page's `primary_code_paths`,
   `primary_path_exceptions`, ownership boundary, and `depends_on` metadata.
 3. Read `status`, `architecture_state`, `last_verified_commit`, and any linked
   in-flight work before treating prose as current behavior.
 4. Use current code and tests when a page is draft, deferred, stale, absent, or
   conflicts with the frozen implementation. Report material documentation
   drift separately.
 5. Treat candidate invariants and promotion gates as review questions until
   repository policy, current code/tests, or owner-approved normative text
   makes them enforceable. Do not turn a draft identifier into an automatic
   blocker.
 6. Use published latest docs only for discovery when the reviewed head lacks a
   page; never import a newer contract into an older branch silently.
 When behavior changes an active module or feature contract, require the same PR
 to update that page. Require a new page only when repository policy, an accepted
 RFC/issue, or explicit maintainer acceptance criteria do so; otherwise treat it
 as a non-blocking design suggestion. Do not demand a second page when an existing
 design already owns the exact contract.
@@ -1,129 +0,0 @@
 # General Review Checks
 Apply these checks to every PR after freezing the review surface. A pattern is
 not a finding until the diff introduces or exposes a reachable trigger, impact,
 and smallest safe fix. Do not report unrelated backlog.
 Official docs: [contributing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/)
 [feature compatibility](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/feature_compatibility/),
 and [design documents](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/).
 ## Contract and scope
 - Match the title/body claims to every changed production path, test, config,
  dependency, generated artifact, and user-visible document.
 - Trace each changed value through public ingress, validation/defaulting,
  producer, transformations, process or device boundaries, final consumer, and
  terminal cleanup.
 - Check every applicable sync/async, offline/online, streaming/non-streaming,
  feature-off, topology, and compatibility path without demanding unsupported
  modes merely for symmetry.
 ## Module, API, and class design
 Apply these rules to new or materially changed production code; do not require
 unrelated cleanup from untouched files.
 1. Start each source module with a concise module docstring that defines its
   responsibility and boundaries and identifies its primary classes and key
   functions. Place it after any required license header and before imports.
 2. Give every added or materially changed class and function a clear docstring
   covering its responsibility, parameters, return value, raised exceptions,
   and non-obvious invariants. Use explicit typed parameters at owned API
   boundaries; avoid opaque `**kwargs` and `Any`. Allow them only when an
   upstream or compatibility interface requires them, and document and validate
   the accepted keys and concrete types.
 3. Move stable helpers shared by multiple features into an appropriately scoped
   utility module. Do not let feature modules accumulate large collections of
   unrelated private utility functions; keep a private helper local only when
   it is single-purpose and tightly coupled to that feature.
 4. Keep every method on a key class within that class's stated responsibility,
   and document non-obvious design decisions and invariants. Before adding a
   method, prefer reuse, consolidation, or extension of an existing abstraction;
   split the class when new behavior creates a separate responsibility or makes
   the class unnecessarily long.
 5. Require changes to user-facing behavior to update `docs/`. When an existing
   active module or feature design page owns the changed contract, update that
   page. Require a new design page only when the reviewed branch's repository
   policy, an accepted RFC/issue, or explicit maintainer acceptance criteria do
   so; otherwise suggest it as non-blocking. For a major architectural change,
   apply the branch's contributing policy for a design RFC or issue; the current
   policy uses more than 500 changed production-code lines, excluding kernel,
   data, config, and test changes. Resolve page status and authority through
   [design-contracts.md](design-contracts.md), and do not demand a duplicate page
   when an existing design owns the exact contract. When required, cover scope,
   data flow, public contracts, failure behavior, alternatives, validation,
   limitations, and rollout.
 6. Require every performance PR to provide reproducible base-versus-head A/B
   results under the same hardware, software, model, workload, precision,
   topology, warmup, repetitions, and measurement scope. Include exact commands,
   variability, and a correctness or quality guard; one-sided or non-comparable
   numbers are incomplete evidence. For a new model absent from the frozen base,
   compare head with its pinned canonical reference implementation under rule 12
   instead. Follow
   [perf-verification.md](../checks/perf-verification.md) for the full protocol.
 7. Require every addition, removal, default change, or semantic change to a
   user-facing inference API or CLI to explain why it is necessary, which
   alternatives were considered, and why the chosen interface is preferable.
   Record compatibility, deprecation or migration behavior, and the rationale
   in the PR. Update the corresponding module and feature design documents when
   they define the changed interface or workflow, or when policy requires one.
 8. Require every bug-fix PR to add an automated regression test that reproduces
   the defect on the frozen base and passes on the fixed head. Exercise the
   production path or nearest stable boundary and assert the corrected behavior,
   not only process survival. Record the exact base/head command and result, and
   place hardware- or model-specific coverage in the correct CI lane. A manual
   reproduction supplements but does not replace the regression test.
 9. Require every new or materially changed recipe to follow
   [`recipes/TEMPLATE.md`](../../../../../recipes/TEMPLATE.md) and the naming and
   layout conventions in
   [`recipes/README.md`](../../../../../recipes/README.md).
   Complete every applicable summary, use-case, reference, hardware,
   environment, exact-command, verification, and notes section with tested,
   reproducible details; explain any intentional deviation in the PR.
 10. Require every new model to verify its supported modalities, platforms,
    execution modes, and features with production-path evidence. Update
    `docs/models/supported_models.md` and every applicable feature support or
    compatibility table; do not claim support that the PR does not validate.
    Follow
    [model-addition-checklist.md](../checks/model-addition-checklist.md) for the
    model-specific documentation checks.
 11. Require every new model PR to add or update its in-repository model-family
    recipe under `recipes/` and its `recipes/README.md` index entry. Apply rule
    9, and keep the recipe's model identifiers, tasks, serving modes, hardware,
    commands, supported features, and limitations aligned with validated PR
    evidence. An external recipe link does not replace the in-repository recipe.
 12. Require every new model to compare accuracy or output quality and performance
    against a pinned baseline or canonical reference implementation. Match the
    checkpoint, inputs, precision, hardware, workload, and measurement method;
    provide exact commands, results, and explanations for material differences.
    Prefer a detailed timing split for each stage and explain when it is
    unavailable. Follow
    [model-addition-checklist.md](../checks/model-addition-checklist.md) and
    [perf-verification.md](../checks/perf-verification.md) for the evidence
    contract.
 ## Blocking risk scan
 | Risk | Prove before reporting |
 | --- | --- |
 | Correctness | A live input reaches a wrong output, exception, silent drop, shape/dtype/device mismatch, or partially initialized state. |
 | Compatibility | A supported API, config key, default, serialization schema, model identifier, or caller breaks without validation, migration, or deprecation. |
 | Lifecycle | Allocation or startup lacks cleanup on partial failure, timeout, cancellation, shutdown, or repeated requests. |
 | Concurrency/distribution | Ordering, identity, rank, world size, completion, timeout, backpressure, or collective assumptions fail on a supported path. |
 | Cache/shared state | Identity omits a correctness-affecting input, or isolation, invalidation, disabled behavior, eviction, or cleanup is incomplete. |
 | Async behavior | Blocking I/O, sleep, a lock across `await`, or unnecessary serialization harms a reachable async path. |
 | Security/data handling | Secrets, user payloads, unsafe deserialization/eval/shell, unvalidated paths, or unbounded metric labels are reachable. |
 | Validation evidence | Tests bypass the production dispatcher, mock away the change, use unrealistic types/MRO, or cannot fail when the defect returns. |
 | User contract | User-visible model, feature, API, CLI, config, default, or compatibility behavior changes without accurate docs or examples. |
 Search bounded callers and sibling implementations before reporting dead code,
 duplication, or a missed consumer. Keep new abstractions only when they have a
 distinct live caller, invariant, or compatibility need.
 ## Finding bar
 Anchor each finding to a changed `path:line`; name the trigger or call path,
 current behavior, impact, and smallest fix direction. Treat pending CI, missing
 hardware, or unsupported measurements as validation gaps unless the repository
 contract makes that evidence a merge requirement.
@@ -1,191 +0,0 @@
 # Review Execution and Delivery
 Use this reference for every review. It defines snapshot collection, gates,
 validation records, delivery permissions, maintainer tone, and re-review.
 Official docs: [contributing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/)
 and [CI failure triage](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/failures/).
 ## Contents
 - [Freeze the review surface](#freeze-the-review-surface)
 - [Report status before analysis](#report-status-before-analysis)
 - [Apply review gates](#apply-review-gates)
 - [Run bounded validation](#run-bounded-validation)
 - [Deliver maintainer-style findings](#deliver-maintainer-style-findings)
 - [Re-review safely](#re-review-safely)
 ## Freeze the review surface
 For a GitHub PR, read the base/head SHA before and after fetching metadata and
 the diff:
 ```bash
 gh api "repos/vllm-project/vllm-omni/pulls/<PR>" \
  --jq '{base_sha: .base.sha, head_sha: .head.sha}'
 REVIEW_FIELDS="number,url,title,body,isDraft,baseRefName,headRefName,mergeable,mergeStateStatus,statusCheckRollup,files"
 gh pr view <PR> --repo vllm-project/vllm-omni \
  --json "${REVIEW_FIELDS}"
 gh pr diff <PR> --repo vllm-project/vllm-omni
 gh api "repos/vllm-project/vllm-omni/pulls/<PR>" \
  --jq '{base_sha: .base.sha, head_sha: .head.sha}'
 ```
 Discard the snapshot if either SHA changed. Do not mix comments or validation
 from different heads.
 After the SHAs stabilize, fetch and materialize a trusted `head_sha` in an
 isolated detached worktree. Run every source read, `rg` search, import, and test
 from that snapshot, not the caller's checkout. A detached worktree provides
 snapshot isolation, not a security boundary.
 Treat a fork head as untrusted unless the user and environment policy explicitly
 establish trust. Never run its imports, tests, builds, hooks, package setup, or
 repo-configurable linters/plugins on the reviewer host. Execute them only in a
 disposable sandbox or VM with no credentials or inherited secrets, no host agent
 or service sockets, minimal read-only host mounts, disabled or explicitly
 allowlisted network, resource/time limits, and destruction after the run. If
 that boundary is unavailable, limit the review to the remote diff, SHA-addressed
 reads such as `git show <head_sha>:<path>`, and existing CI evidence; report all
 executable validation as a gap.
 Before the first source read, record a pristine snapshot fingerprint outside
 the reviewed worktree. Include `HEAD`, the NUL-delimited index entries and blob
 IDs, and a NUL-safe manifest of every worktree entry except Git metadata,
 including tracked, untracked, and ignored paths with file type, mode, symlink
 target, and content hash. Before each validation group and delivery, recompute
 and byte-compare the fingerprint as well as asserting `HEAD == head_sha`. If a
 tool changes or creates any entry, discard affected evidence and recreate the
 snapshot before another group; do not rely on `HEAD` alone or clean an unknown
 worktree in place. If an exact snapshot cannot be created, use SHA-addressed
 reads and report filesystem-dependent validation as a gap.
 For a local branch/worktree, determine the target ref from the user, current PR,
 or configured upstream; never infer it from the branch name. Resolve the target
 once and include all task-owned worktree state:
 ```bash
 git status --porcelain=v2 -z
 git rev-parse HEAD
 git rev-parse <target-ref>
 git merge-base <target-base-sha> HEAD
 git diff --stat <comparison-commit>
 git diff --name-status <comparison-commit>
 git diff --binary <comparison-commit>
 git diff --cached --binary <comparison-commit>
 git diff --binary
 git ls-files --others --exclude-standard -z
 ```
 Freeze the exact bytes of every in-scope untracked file from the NUL-delimited
 list, plus a NUL-safe path, file-type, mode, and content-hash manifest. Also
 fingerprint `HEAD`, the target and merge base, porcelain status, and both binary
 index and worktree patches. Recording only untracked names is insufficient.
 Keep these snapshots in the evidence packet and do not mutate the reviewed
 checkout during a read-only review.
 For local changes without PR metadata, mark the title/body as synthetic. Base
 them on the user request and commit messages without inventing validation claims.
 If no PR, branch, or usable worktree is identifiable, ask for a PR URL/number or
 explicit base and head rather than guessing.
 ## Report status before analysis
 Within 60 seconds of starting, and before source searches or tests, send this
 host update:
 ```text
 Pinned head: <SHA>
 Base/comparison: <ref and SHA>
 CI: <pass/fail/pending/not applicable>
 Mergeability: <state/not applicable>
 Preliminary findings: <brief finding or none yet>
 ```
 Mark early findings as preliminary and continue. This is not a GitHub comment.
 ## Apply review gates
 - Report draft/WIP state. Continue a local review when explicitly requested,
  but do not publish review events without separate authorization.
 - Record DCO, pre-commit, required CI, and mergeability. Pending or unknown gates
  do not block source review.
 - Treat a failed gate as its own evidence. Do not restate its formatting/lint
  output as a new code-review finding.
 - Open CI logs only when the first failing step overlaps the frozen diff or
  blocks the verdict. Start with the first error, not the last cascade.
 - Treat the PR body as navigation. Commands, benchmark tables, and claimed tests
  become evidence only after their provenance and relevance are checked.
 ## Run bounded validation
 Keep one evidence packet for files read, bounded searches, callers, tests, CI,
 hardware, routes, and findings. Reuse it rather than refetching facts.
 Before pytest, run a short import/version compatibility check. Record each
 validation result with:
 ```text
 repo, head SHA, command, result, Python/platform, dependency or lock fingerprint
 ```
 After preflight, run targeted tests and low-cost static checks alongside source
 inspection when possible. Imports, tests, builds, and linters may create caches
 or rewrite files, so isolate each potentially mutating group in a disposable
 snapshot or recreate the pinned snapshot before continuing. Map source symbols
 to tests with bounded `rg` searches rather than assuming the test directory
 mirrors production paths.
 Classify failures as code, test, infrastructure, or flaky before reporting.
 Skipped hardware tests are gaps, not passes. For available hardware, verify the
 smallest representative unit/E2E path and compare actual output with PR claims.
 ## Deliver maintainer-style findings
 Default to roughly 1-5 short comments, but treat that as calibration rather than
 a quota: report every real blocker and report no finding when there is no issue.
 Write each finding as:
 ```text
 [P1] Short imperative title — path/to/file.py:<line>
 <Trigger or call path>. <Current behavior and impact>. <Smallest fix direction>.
 ```
 Use [maintainer-style-study.md](../delivery/maintainer-style-study.md) for
 comment tone.
 Keep rule IDs, grades, and audit matrices internal unless the user asks for the
 complete audit.
 Verify every inline `path:line` against the frozen diff before delivery. For a
 PR, re-read the remote head and byte-compare the detached snapshot fingerprint,
 or recreate it from `head_sha`, before asserting current lines. For a local
 review, recompute and byte-compare the frozen `HEAD`, target, merge base, status,
 index patch, worktree patch, and untracked-content manifest. Any mismatch makes
 the review stale; discard the affected evidence and restart. Prefer one
 root-cause comment to several symptom comments.
 Local presentation is the default. Only explicit authorization permits GitHub
 posting. When authorized, publish one consolidated final review after
 validation; do not post preliminary or incremental comments. Submit a review
 event (`APPROVE`, `COMMENT`, or `REQUEST_CHANGES`) only when the user explicitly
 chooses that event.
 Reviewer requests and owner `@mention` comments are separate external writes;
 follow [review-requests.md](../delivery/review-requests.md) when the user asks
 for them.
 ## Re-review safely
 Freeze the new head and compare it with the previously reviewed SHA. Then:
 1. inspect the delta and every file affected by conflict resolution or rebase;
 2. revalidate earlier findings against current lines and behavior;
 3. read unresolved, non-outdated threads and the author's evidence;
 4. rerun only checks invalidated by the delta;
 5. look for new regressions introduced by the fix;
 6. do not repeat resolved or outdated comments.
 If the target head changes during re-review, discard stale validation and start
 again from the new snapshot.
@@ -1,76 +0,0 @@
 # Module and Feature Review Routing
 Read [design-contracts.md](design-contracts.md) first. Route from the frozen
 head's live producer-consumer behavior, then confirm the result against its
 `docs/design/` metadata. Titles and paths are evidence, not the deciding rule.
 ## Primary module contract
 Choose one primary module and a second only when the live call path crosses a
 documented dependency or primary-path exception.
 | Module | Live behavior signals | Read |
 | --- | --- | --- |
 | Entrypoints | Offline/CLI/API ingress, validation, rendering, streaming, sessions | [entrypoints.md](../modules/entrypoints.md) |
 | Configuration | Schema, defaults, deploy/stage construction, registry, topology | [configuration.md](../modules/configuration.md) |
 | I/O and Modality | Requests, messages, serialization, output types, accumulation, completion | [input-output-modality.md](../modules/input-output-modality.md) |
 | Error Contracts | Classification, fatality, propagation, sanitization, public rendering | [error-contracts.md](../modules/error-contracts.md) |
 | Engine Orchestration | Cross-stage routing, request state, ordering, RPC correlation, terminal convergence | [engine-orchestration.md](../modules/engine-orchestration.md) |
 | Stage Runtime | Placement, startup, readiness, replicas, affinity, membership, shutdown | [stage-runtime.md](../modules/stage-runtime.md) |
 | OmniConnector | Cross-stage/process/device/node transport and synchronization | [omni-connector.md](../modules/omni-connector.md) |
 | Model Integration | Registration, preprocessing, loading, runners, model-specific execution | [model-integration.md](../modules/model-integration.md) |
 | AR Runtime | Scheduling, request/cache state, adapters, workers, upstream vLLM semantics | [ar-runtime.md](../modules/ar-runtime.md) |
 | Diffusion Family | Diffusion runtime, models, batching, parallelism, or offload | [diffusion.md](../modules/diffusion.md) |
 | Execution Platforms | Hardware selection, capabilities, vendor workers, kernels, patches | [execution-platforms.md](../modules/execution-platforms.md) |
 | Cache Management | Cache identity, validity, reuse, eviction, reset, teardown | [cache-management.md](../modules/cache-management.md) |
 | Quantization | Method/metadata selection, layer mapping, reduced precision, constraints | [quantization.md](../modules/quantization.md) |
 | Observability | Metrics, logs, labels, units, correlation, lifecycle | [observability.md](../modules/observability.md) |
 | Profiling | Opt-in instrumentation, traces, start/stop lifecycle, overhead | [profiling.md](../modules/profiling.md) |
 | Benchmarking | Workloads, metric computation, benchmark CLI, result metadata | [benchmarking.md](../modules/benchmarking.md) |
 For tests-, docs-, CI-, or recipe-only changes, route to the production module
 whose contract they protect. When none exists, use the applicable evidence
 check without inventing a production owner.
 ## Feature-design overlays
 Load every category whose documented behavior, enablement, default, support, or
 compatibility changes.
 | Feature-design section | Read |
 | --- | --- |
 | Runtime and stage execution | [runtime-stage-execution.md](../features/runtime-stage-execution.md) |
 | Communication and concrete connectors | [communication.md](../features/communication.md) |
 | Diffusion acceleration | [diffusion-acceleration.md](../features/diffusion-acceleration.md) |
 | Infrastructure and performance | [infrastructure-performance.md](../features/infrastructure-performance.md) |
 ## Evidence and change overlays
 | Signal | Read | Optional repo-local skill |
 | --- | --- | --- |
 | New or expanded model, loader, processor, registry, or stage config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) |
 | Latency, throughput, memory, scaling, precision, or quality claim | [perf-verification.md](../checks/perf-verification.md) | [`diffusion-perf-opt`](../../../diffusion-perf-opt/SKILL.md) for diffusion |
 | Tests changed, absent for risky behavior, or test-only | [test-quality-evaluation.md](../checks/test-quality-evaluation.md) | [`vllm-omni-test`](../../../vllm-omni-test/SKILL.md) |
 | CI, examples, docs, public behavior, or contributor evidence | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | None |
 | Suitable hardware/server and runnable affected path | [verification.md](../checks/verification.md) | None |
 | User asks who should review or requests owner notification | [review-requests.md](../delivery/review-requests.md) | None |
 Use [`quantization`](../../../quantization/SKILL.md) for a quantization module
 change and [`$vllm-omni-npu-model-runner-upgrade`](../../../vllm-omni-npu-upgrade/SKILL.md)
 for an NPU execution-platform change when those skills are available.
 For a bug fix, require a reachable reproduction and regression test. For a
 refactor, prove parity and remove obsolete paths. For a feature, verify the
 public/config contract, module integration, compatibility/default behavior,
 production dispatch, feature design, and docs.
 ## Calibrate findings
 - **P0:** security exposure, data corruption, or broad project unusability.
 - **P1:** reachable runtime failure, wrong output, compatibility break, or
  unsafe lifecycle in the changed behavior.
 - **P2:** real non-blocking defect with a concrete future failure mode.
 Treat draft candidate invariants, missing hardware, and unsupported claims as
 questions or validation gaps unless current code, tests, or policy makes them a
 merge requirement.
@@ -61,13 +61,6 @@ repos:
      types: [python]
      files: ^tests/
    - id: check-tts-adapter-migration
      name: Ratchet TTS per-model branches out of serving_speech.py
      entry: python tools/pre_commit/check_tts_adapter.py
      language: python
      pass_filenames: false
      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
    # Keep `suggestion` last
    - id: suggestion
      name: Suggestion
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
@@ -16,21 +16,31 @@ nav:
  - Examples:
    - examples/README.md
    - Offline Inference:
      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
      - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
    - Online Serving:
      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
-      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
  - General:
    - usage/*
@@ -63,5 +63,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
@@ -672,147 +672,100 @@ down and spawn a new one mid-module. A few rules that save real CI debugging tim
 ## Online Serving Integration
-To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, write
+To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
-one adapter under `vllm_omni/entrypoints/openai/tts_adapters/`. The adapter declares
+**all five** of the following integration points to
-how your model is recognized and owns its request handling; `serving_speech.py` should
+`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
-not need an edit.
+piecemeal causes partial-integration failures that are hard to debug.
-### 1. Write the adapter
+### 1. Stage constant
-Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`:
+Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
 ```python
-# SPDX-License-Identifier: Apache-2.0
+_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
-"""YourModel serving adapter."""
+```
-from typing import TYPE_CHECKING
+### 2. Union into `_TTS_MODEL_STAGES`
-from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+Add to the `_TTS_MODEL_STAGES` set union:
 from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
-if TYPE_CHECKING:
+```python
-    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+_TTS_MODEL_STAGES: set[str] = (
    ...
    | _YOUR_MODEL_TTS_MODEL_STAGES
 )
 ```
 ### 3. Model type detection
-@register_tts_adapter
+In `_detect_tts_model_type()`, add before the final `return None`:
 class YourModelAdapter(ARTTSAdapter):
    # The registry key, and the model-type string used in logs.
    name = "your_model"
    # The engine ``model_stage`` value(s) your deployment yaml sets. Detection,
    # stage discovery and `/v1/audio/speech` routing all derive from this.
    stage_keys = frozenset({"your_model_stage_key"})
-    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+```python
-        """Return an error string, or None if the request is valid."""
+if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
    return "your_model"
 ```
 ### 4. Request validation dispatch
 In `_validate_tts_request()`, add before the fallback `return`:
 ```python
 if self._tts_model_type == "your_model":
    return self._validate_your_model_request(request)
 ```
 ### 5. Validation and parameter-builder methods
 Add two new methods:
 ```python
 def _validate_your_model_request(
    self, request: OpenAICreateSpeechRequest
 ) -> str | None:
    """Validate YourModel request. Returns an error string or None."""
    if not request.input or not request.input.strip():
        return "Input text cannot be empty"
    return None
-    async def build(
+def _build_your_model_params(
-        self,
+    self, request: OpenAICreateSpeechRequest
-        request: "OpenAICreateSpeechRequest",
+) -> dict[str, Any]:
-        sampling_params_list: list,
+    """Build additional_information dict for YourModel."""
-        has_inline_ref_audio: bool,
+    params: dict[str, Any] = {"text": [request.input]}
    ) -> PreparedRequest:
        """Build the engine prompt and per-model parameters."""
        params: dict[str, list] = {"text": [request.input]}
    if request.voice is not None:
        params["voice"] = [request.voice]
-        return PreparedRequest(
+    # Add any other model-specific fields here
-            prompt={"prompt": request.input},
+    return params
            tts_params=params,
            model_type=self.name,
        )
 ```
-Then add it to the import block at the bottom of
+Then wire `_build_your_model_params` into the request-dispatch block in
-`vllm_omni/entrypoints/openai/tts_adapters/__init__.py` so it registers on import.
+`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
-That import line is the only shared file a new model has to touch.
+existing model to find the right location). If the model supports voice cloning
 (`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
 here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
 `_build_moss_tts_params` for the voice-cloning variant) for the pattern.
-> **Pure-diffusion TTS does not go through adapters yet.** When the server is built
+> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
-> by `OmniOpenAIServingSpeech.for_diffusion()`, `create_speech()` routes straight to
+> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
-> `_create_diffusion_speech()` and never calls `validate()` or `build()`.
+> string pattern shown above. For new models, always use the string pattern — do not
-> `DiffusionTTSAdapter` is scaffolding for that migration and has no production
+> add new `_is_*` boolean flags.
 > subclass today, so validation or prompt construction placed in one would silently
 > never run. If your model is served by the diffusion engine rather than
 > `engine_client`, follow the existing diffusion path and say so on #4855 — wiring
 > that path through the adapter layer is open work, not something to do inside a
 > new-model PR.
-### 2. If a stage key alone cannot identify your model
+> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
-
+> that you actually pass to the model's generate / `inference_stream` call. Extracting
-Two escape hatches, in order of preference:
+> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
 ```python
 # The architecture is authoritative (e.g. your model has no dedicated
 # model_stage value, or shares one with another model).
 model_archs = frozenset({"YourModelForConditionalGeneration"})
 # Set this too if your model owns NO stage key, so stage discovery has to find
 # your AR entry stage by architecture.
 arch_identifies_entry_stage = True
 ```
 If your rule is not a set-membership test, override `matches()` — see
 `covo_audio.py`, which shares the generic `fused_thinker_talker` stage key with
 non-CoVo deployments and confirms with the architecture:
 ```python
@classmethod
 def matches(cls, model_stage: str | None, model_arch: str | None) -> bool:
    return super().matches(model_stage, model_arch) and bool(model_arch) and "CovoAudio" in model_arch
 ```
 If your model can match the same deployment as another adapter, give one of them an
 explicit `detect_priority` (lower runs first). `tests/entrypoints/openai_api/
 test_tts_detection.py` fails if two same-priority detectors can both match one input,
 so an unordered overlap cannot be merged.
 Models that are speech-capable only in some deployment topologies override
 `stage_serves_speech()` — see `audex.py`, whose omni thinker is text-final unless the
 speech decoder is deployed alongside it.
 ### 3. Reuse the shared helpers
 Adapters reach shared serving helpers through `self.ctx.server`: reference-audio
 resolution (`_resolve_ref_audio`), uploaded-speaker handling
 (`_apply_uploaded_speaker`), format validation (`_validate_ref_audio_format`), and
 prompt-length limits (`_max_instructions_length`). Read a comparable adapter before
 writing your own — `fish_speech.py` for voice cloning, `higgs_audio_v3.py` for
 parameter-heavy models, `moss_tts.py` for a family sharing one base class.
 ### 4. Test it
 `tests/entrypoints/openai_api/test_tts_adapter.py` covers registry invariants and
 `test_tts_detection.py` covers detection. Add your model to
 `EXPECTED_MODEL_TYPES` in the former, and a case to the latter if you added a
 `model_archs` or `matches()` rule.
 ### Do not add branches to `serving_speech.py`
 Older models predate the adapter framework and still have per-model branches in
 `serving_speech.py`; they are being migrated out (RFC #4327, #4855). Do not copy
 that pattern. `tools/pre_commit/check_tts_adapter.py` is a ratchet on the remaining
 `self._tts_model_type == ...` comparisons and fails the commit if the count grows.
 If your model genuinely needs behaviour that no adapter hook can express, that is a
 missing hook — propose it on the RFC rather than adding a branch.
 ### Merge conflicts
-Your adapter is a new file, so it cannot conflict. The one shared line is your entry
+`serving_speech.py` is modified by every new model PR and is the most common source of
-in the import block at the bottom of `tts_adapters/__init__.py`, which conflicts only
+rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
-when another model lands at the same time. The resolution is always to **keep both**
+is always to **keep both** the upstream model's additions and your own — never discard
-imports — never discard either side. After resolving:
+either side. After resolving:
 ```bash
-git add vllm_omni/entrypoints/openai/tts_adapters/__init__.py
+git add vllm_omni/entrypoints/openai/serving_speech.py
 git rebase --continue
 ```
 If you find yourself resolving a conflict inside `serving_speech.py`, something in
 your model is still going through the legacy path; see the note above.
 ## Single-Stage Models
 Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
@@ -1,6 +1,6 @@
 # Architecture Overview
-This document outlines the architecture of vLLM-Omni.
+This document outlines the architectural design for vLLM-Omni.
 <p align="center">
  <picture>
@@ -9,456 +9,194 @@ This document outlines the architecture of vLLM-Omni.
  </picture>
 </p>
-## Design goals
+# Goals
-The primary goal of vLLM-Omni is to provide a fast and easy-to-use inference
+The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
 and serving engine for omni-modality models. vLLM-Omni extends vLLM's
 text-oriented autoregressive (AR) runtime with stage-based execution for
 non-textual outputs and non-autoregressive model components.
-The architecture is designed to:
+* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
 * **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
 * **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
 * **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
 * support text, image, audio, video, and action inputs and outputs;
 * compose autoregressive, generation, and diffusion stages in one pipeline;
 * reuse vLLM scheduling, cache, distributed-execution, and serving
  primitives where they fit; and
 * keep model topology, deployment placement, runtime lifecycle, transport,
  and public API concerns in separate layers.
-## Model execution
+# Representative omni-modality models
-These four examples show how vLLM-Omni maps model components to stage-local
+According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
 execution policies. Each summary covers the model-owned path, batching,
 attention, parallelism, quantization, and primary serving targets. The profiles
 are representative; a model may support multiple `PipelineConfig` and
 `DeployConfig` combinations.
-| Model | Model-owned path | vLLM-Omni stage view | Primary serving target |
+**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
-| --- | --- | --- | --- |
+  A powerful image generation foundation model capable of complex text rendering and precise image editing.
 | Qwen3-Omni | Thinker → Talker → Code2Wav | Three stages with asynchronous chunk hand-off | TTFT/TTFP, TPOT, E2EL, audio RTF, and concurrent throughput |
 | HunyuanImage-3.0 | Multimodal AR understanding/reasoning plus image-generation DiT | AR-only, DiT-only, or split AR → DiT deployment | DiT E2EL/denoising latency and image throughput; TTFT/TPOT for AR tasks |
 | MiniMax-H3 | Text encoder → task-specific FL2VA or Ref2VA DiT → video/audio VAE decoder | Three logical component/phase boundaries within one registered diffusion stage | Video/audio E2EL and media throughput; RTF when audio streaming is measured |
 | Cosmos3 | Unified MoT reasoner tower plus diffusion generator tower | One diffusion pipeline materializes both towers inside the stage | Image/video/action E2EL and throughput; RTF for synchronized audio |
 ### Serving metric vocabulary
 The target metric follows the output modality, not just the model name.
 | Metric | Meaning in this document | Most useful for |
 | --- | --- | --- |
 | **TTFT** | Time to the first text token | AR reasoning or text output |
 | **TPOT** | Time per output token after the first token | AR decode stages |
 | **TTFP** | Time to the first streamed media packet, such as audio | Interactive speech and multimodal chat |
 | **E2EL** | End-to-end latency from request admission to the final output | Image, video, audio, and action requests |
 | **RTF** | Wall-clock processing time divided by generated audio duration; lower is faster than real time | Audio generation |
 | **Throughput** | Requests, tokens, or media seconds produced per wall-clock second | Offline and concurrent serving |
 ### Qwen3-Omni: streaming AR pipeline
 Qwen3-Omni illustrates a multi-stage streaming AR pipeline: Thinker → Talker →
 Code2Wav. The [technical report](https://arxiv.org/abs/2509.17765) describes
 the model components, and the figure shows the stage-level streaming schedule
 used by vLLM-Omni.
 <p align="center">
-  <a href="https://github.com/vllm-project/vllm-omni/blob/main/docs/source/architecture/qwen3-omni-async-chunk.png">
+  <picture>
-    <img alt="Qwen3-Omni Thinker, Talker, and Code2Wav streaming stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width="80%">
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png">
-  </a>
+    <img alt="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
  </picture>
 </p>
-#### Model architecture by stage
+**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
-
+  A unified multimodal comprehension and generation model, with cot text output and visual generation.
 | Stage | Model component | Stage output |
 | --- | --- | --- |
 | **Thinker** | AuT audio encoder, SigLIP2 vision encoder, and the 30B-total/3B-active Thinker MoE Transformer | Text tokens and conditioning for the Talker |
 | **Talker** | 3B-total/0.3B-active Talker MoE Transformer with the MTP code-prediction path | Text tokens and audio codec tokens |
 | **Code2Wav** | Approximately 200M-parameter ConvNet audio decoder | Waveform/audio chunks |
 #### Representative stage execution
 | Stage | Batching | Attention and execution | Parallelism | Quantization |
 | --- | --- | --- | --- | --- |
 | **Thinker** | Continuous batching with stage-local `max_num_seqs` and token budgets | vLLM KV-cached AR attention; CUDA Graph execution when the stage is not eager | Independent stage placement; use the AR runtime's configured tensor parallelism | Supported ModelOpt FP8/NVFP4 or AutoRound checkpoint paths target the Thinker; encoders remain BF16 in the documented paths |
 | **Talker** | Continuous batching over autoregressive decode requests | KV-cached token decode with async-chunk output to Code2Wav | Independent stage placement; typically a separate GPU or replica from Thinker | BF16 baseline; no generic Talker quantization path is assumed |
 | **Code2Wav** | Static/chunk batching for codec-to-waveform decode | Non-AR chunked audio decode; eager execution is commonly retained | Separate stage or colocated placement, depending on the deployment | BF16 baseline; keep codec/vocoder precision model-specific |
 For this pipeline, the main serving target is interactive responsiveness:
 TTFT for text, TTFP for the first audio packet, TPOT for autoregressive decode,
 E2EL for the complete answer, and RTF/throughput for sustained audio serving.
 Async chunking and streaming primarily reduce TTFP and improve stage overlap;
 batching and CUDA Graphs primarily improve E2EL and throughput. See the
 [Qwen3-Omni recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md)
 and [async chunk design](feature/async_chunk.md) for the deployment details.
 ### HunyuanImage-3.0: shared multimodal model, split deployment choices
 HunyuanImage-3.0 illustrates a shared multimodal model with flexible AR, DiT,
 and split AR → DiT deployment. The official [technical
 report](https://arxiv.org/abs/2509.23951) provides the model view; the tables
 below map it to vLLM-Omni stages.
 <p align="center">
-  <a href="https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/assets/framework.png">
+  <picture>
-    <img alt="HunyuanImage-3.0 framework for understanding, language modeling, and image generation" src="https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/framework.png" width="95%">
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png">
-  </a>
+    <img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
  </picture>
 </p>
-#### Model architecture by stage
+**AR+DiT (e.g. Qwen-Omni)**
-
+  A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
 | Stage or component | Model component | Stage output |
 | --- | --- | --- |
 | **AR understanding/reasoning** | Understanding and generation encoders feeding the Hunyuan-A13B decoder-only Transformer and text detokenizer | Text responses or image-generation conditioning |
 | **Image generation** | Generation encoder, diffusion prediction path/Gen. Decoder, and the image VAE decode path | Image latents and final images |
 The vLLM-Omni deployment may expose the AR path only, the DiT path only, or
 both as an AR → DiT pipeline with a connector between them. That is a runtime
 decomposition of the model's capabilities; it should not be read as a claim
 that the published framework figure contains two unrelated backbones.
 #### Representative stage execution
 | Stage | Batching | Attention and execution | Parallelism | Quantization |
 | --- | --- | --- | --- | --- |
 | **AR** | Standard vLLM continuous batching when serving understanding/text tasks | KV-cached causal attention | Tensor/expert parallelism follows the selected AR deployment | Checkpoint-specific; do not infer DiT FP8 settings for the AR checkpoint |
 | **DiT + VAE** | Request batching or step-wise batching; Hunyuan step batching with more than one request requires `TORCH_SDPA` | `TORCH_SDPA` is used because the model mixes causal and full attention | Validated profiles include TP4, TP2 + sequence parallelism, TP2 + CFG parallelism, and expert parallelism for MoE weights | Online FP8 and ModelOpt mixed FP8/NVFP4 are documented for the DiT; VAE precision stays separate |
 Image generation is therefore measured mainly with per-image E2EL/denoising
 latency, peak memory, and images-per-second throughput. TTFT and TPOT remain
 meaningful for the optional AR understanding path, but they are not the right
 headline metrics for the DiT path. See the
 [HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
 for the validated batching, attention, and quantization combinations.
 ### MiniMax-H3: conditioning, denoising, and decode
 MiniMax-H3 illustrates logical component and phase boundaries between text
 encoding, task-specific denoising, and media decoding:
 `Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
 these components in one diffusion stage. These boundaries are useful for
 profiling and offload; independent placement would require a future
 multi-stage `PipelineConfig`. The stage mapping follows the implementation and
 serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
 notes that a detailed technical report is forthcoming.
 #### Model architecture by component
 | Logical component | Model component | Component output |
 | --- | --- | --- |
 | **Text Encoder** | Tokenizer/processor and Qwen3-VL text encoder; vision/audio references are prepared alongside the text condition | Packed multimodal conditioning for denoising |
 | **DiT** | Task-specific `FL2VA` DiT for text/first-frame generation or `Ref2VA` DiT for mixed references | Video and audio latent tokens |
 | **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
 Both DiT partitions remain behind one task-selection point; `task` chooses
 which partition runs. In the current implementation, the three logical
 components execute inside one diffusion stage. This is different from
 Qwen3-Omni's three independently scheduled stages, even though both models
 produce synchronized audio and video-like outputs.
 #### Representative component execution
 | Logical component | Batching | Attention and execution | Parallelism | Quantization |
 | --- | --- | --- | --- | --- |
 | **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
 | **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
 | **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
 The primary target is video/audio E2EL and media throughput, with the logical
 component boundaries making prompt encoding, denoising, and VAE decode costs
 separately measurable. RTF is useful when the audio stream is evaluated
 independently. TTFT and TPOT do not describe the main H3 generation path
 because the user-visible output is diffusion media rather than a token stream.
 The
 [MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
 contains the hardware-specific profiles and warmup requirements.
 ### Cosmos3: unified MoT reasoner and generator
 Cosmos3 illustrates a unified Mixture-of-Transformers (MoT) reasoner and
 diffusion generator. The official [technical
 report](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)
 and NVIDIA's [technical explanation](https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/)
 provide the two-tower model view used below.
 <p align="center">
-  <a href="https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/">
+  <picture>
-    <img alt="Cosmos3 autoregressive reasoner and diffusion generator towers" src="https://developer-blogs.nvidia.com/wp-content/uploads/2026/05/image-11-e1780000686151.webp" width="95%">
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png">
-  </a>
+    <img alt="Qwen-Omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png" width=30%>
  </picture>
 </p>
-#### Model architecture by stage
+# vLLM-Omni main architecture
-| Stage or component | Model component | Stage output |
+<p align="center">
-| --- | --- | --- |
+  <picture>
-| **Reasoner tower** | Autoregressive VLM path with causal attention over language and optional vision/action context; runs once per generation request | Discrete reasoning/context tokens |
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png">
-| **Generator tower** | Diffusion path with full attention over the generator tokens and the reasoner K/V context; runs for each denoising step | Image, video, audio, or action tokens |
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png" width=55%>
-| **Shared pipeline** | `Cosmos3OmniDiffusersPipeline` plus modality encoders/decoders and VAE paths | Final image/video/audio or action response |
+  </picture>
 </p>
-The reasoner and generator are model-owned components, not two independent
+## Key Components
 vLLM-Omni stages by default. The current Cosmos3 pipeline materializes both
 inside one diffusion stage; model-level CPU offload can swap the reasoner and
 generator components between their phases.
-#### Representative stage execution
+| Component         | Description                                                                                                                              |
 | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
 | **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
 | **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
 | **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
 | **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
-| Stage | Batching | Attention and execution | Parallelism | Quantization |
+Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
 | --- | --- | --- | --- | --- |
 | **Reasoner + generator pipeline** | Use request batching or step execution only when the selected pipeline advertises that capability; otherwise use `max_num_seqs=1` | Causal attention for reasoner tokens and full attention for diffusion tokens; the validated recipe uses a platform-selected aiter FlashAttention path with `TORCH_SDPA` fallback | Ulysses/context parallelism or tensor parallelism for the transformer; VAE tiling and layerwise/model-level offload address activation and weight memory | Online FP8 is supported for the diffusion pipeline; AutoRound/checkpoint-specific paths and quality validation remain separate from the BF16 baseline |
 For Cosmos3, generation E2EL and output throughput are the primary targets.
 Action-policy workloads add a control-loop latency target, while synchronized
 audio workloads add RTF. TTFT/TPOT are appropriate for a reasoner-only service,
 but not for the full diffusion generation path. See the
 [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
 and [diffusion execution modes](../user_guide/diffusion/execution_modes.md)
 for the current runtime choices.
 ## Why the stage-based architecture follows from these examples
 The four representative pipelines show that omni-modality serving is not one
 execution problem. Autoregressive decoding, diffusion denoising, multimodal
 encoding, and media decoding have different scheduling, attention, parallelism,
 memory, and quantization requirements. The architecture therefore needs a
 unit that can be optimized and resourced independently while remaining part of
 one request pipeline.
 | Observation | Design implication | vLLM-Omni response |
 | --- | --- | --- |
 | AR, diffusion, encoder, and decoder components have different execution loops | One scheduler and one execution policy cannot fit every component | Specialized AR, generation, and diffusion stage runtimes |
 | A model can be deployed as AR-only, DiT-only, or AR → DiT | Model structure must be separate from deployment topology | `PipelineConfig` describes logical stages and relationships; `DeployConfig` describes placement and resources |
 | Text encoders, DiTs, and VAE decoders have different batch and memory behavior | Batching, attention, parallelism, and quantization must be stage-local | Each stage carries its own runtime and engine configuration |
 | MiniMax-H3 can co-locate or logically split its text encoder and VAE decoder | A stage boundary should be configurable rather than forcing a process boundary | Logical stages can be co-located or assigned independent runtime placement |
 | Qwen3-Omni transfers hidden states and codec chunks across stages | Large intermediate data needs typed transport and synchronization | `OmniConnector` transports payloads and KV-cache data without owning model routing |
 | Requests span multiple stages and may be cancelled or produce ordered streaming output | Cross-stage request state needs a separate control plane | `Orchestrator` owns correlation, cancellation, routing state, and output ordering |
 | Qwen3-Omni optimizes first-packet latency while image/video models optimize E2EL and throughput | Performance objectives must be model- and stage-aware | TTFT/TPOT/TTFP, E2EL, RTF, and throughput are measured according to the output modality |
 This leads to a precise definition of a stage: it is a logical execution unit
 with its own model component, execution policy, resource ownership, and
 performance objectives. A stage is not necessarily a separate process. Cosmos3
 keeps its reasoner and generator inside one diffusion pipeline, while
 MiniMax-H3 exposes useful logical boundaries between text encoding, denoising,
 and VAE decoding. Both are valid mappings of the same abstraction.
 The resulting separation can be summarized as:
 ```text
 Model structure       What is computed and how components are related
        ↓
 PipelineConfig        Which logical stages exist and how they are connected
        ↓
 Stage runtime policy  How each stage batches, attends, parallelizes, and quantizes
        ↓
 DeployConfig           Where stages run, with which devices, replicas, and connectors
 ```
 `AsyncOmniEngine` and `Orchestrator` coordinate the pipeline, `OmniConnector`
 transports stage outputs, and `StageRuntime` materializes the selected
 placement. This separation lets one logical model support multiple serving
 topologies without coupling model code to process layout or deployment
 resources.
 ## System architecture
 The figure shows the main components of the stage-based runtime: entrypoints,
 the shared `AsyncOmniEngine`, request orchestration, stage lifecycle
 management, AR and diffusion execution modules, model/layer operations,
 connector transport, and multimodal outputs. `AsyncOmniEngine` is the
 composition root between the entrypoints and the orchestrator.
 ```mermaid
 flowchart TB
    subgraph omni["vLLM-Omni"]
        subgraph entry["Entry points"]
            api["API server"]
            client["Omni / AsyncOmni"]
        end
        engine["AsyncOmniEngine<br/>composition root and background loop"]
        orchestrator["Orchestrator<br/>request state and cross-stage routing"]
        runtime["StageRuntime<br/>placement, replicas, readiness, and lifecycle"]
        ar["AR module<br/>vLLM engine, scheduler, cache, executor, worker, and model runner"]
        diffusion["Diffusion module<br/>scheduler, worker, model runner, and pipeline"]
        ops["Model / layer / ops<br/>attention, parallelism, and quantization"]
        connector["OmniConnector<br/>payload and KV transport plus synchronization"]
        outputs["Multimodal outputs<br/>streaming chunks and final artifacts"]
        api --> engine
        client --> engine
        engine --> orchestrator
        engine --> runtime
        orchestrator --> runtime
        runtime --> ar
        runtime --> diffusion
        ar -. uses .-> ops
        diffusion -. uses .-> ops
        ar -. transports through .-> connector
        diffusion -. transports through .-> connector
        ar --> outputs
        diffusion --> outputs
    end
 ```
 ### Key components
 | Component | Responsibility |
 | --- | --- |
 | **Entrypoints** | Translate offline, CLI, OpenAI-compatible, and duplex requests into engine operations and render outputs back to public protocols. |
 | **Configuration resolution** | Combines pipeline topology, deployment settings, model metadata, and user overrides into a validated control-plane configuration. |
 | **AsyncOmniEngine** | Owns engine composition, the background event loop, stage initialization, request submission, and output collection. |
 | **Orchestrator** | Owns cross-stage request state, stage-to-stage routing, request correlation, cancellation, and output ordering. It does not own model selection or deployment placement. |
 | **StageRuntime** | Expands logical stages into local or distributed replicas, starts stage clients and processes, and manages readiness, affinity, failure, and shutdown. |
 | **AR runtime** | Extends vLLM's scheduler, KV-cache, worker, and model-runner path for omni-modality inputs and inter-stage outputs. |
 | **Diffusion runtime** | Schedules and executes denoising workloads through diffusion executors, workers, pipelines, acceleration backends, and output materialization. |
 | **OmniConnector** | Transports stage payloads and KV-cache data and provides synchronization. Connectors transport data; they do not choose the next logical stage. |
 | **Multimodal outputs** | `MultimodalPayload` separates tensor content from metadata, while `OmniRequestOutput` carries pipeline and diffusion results through the common output path. |
 ## Configuration and runtime resolution
 The control plane has five conceptual layers. Authoring inputs are resolved
 once into a complete, transport-safe configuration before runtime processes are
 started. This keeps stage topology and model capabilities distinct from
 deployment placement and from process-local engine objects.
 The figure shows only the primary hand-off object for each layer; the detailed
 inputs, fields, and ownership rules are described below.
 ```mermaid
 flowchart TB
    layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
    layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
    layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
    layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
    layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
    layer1 --> layer2 --> layer3 --> layer4 --> layer5
 ```
 The single resolution boundary is implemented by
 `StageConfigFactory.create_from_model()` and
 `VllmOmniConfig.from_pipeline_config()` in
 [`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
 The legacy `stage_args` YAML path remains only for models that have not yet
 migrated to `PipelineConfig` and `DeployConfig`.
 In the typed path, each stage configuration derives from
 `BaseVllmOmniStageConfig` and is specialized as
 `VllmOmniARStageConfig`, `VllmOmniGenerationStageConfig`, or
 `VllmOmniDiffusionStageConfig`. The request and engine-spec fields belong to the
 control-plane boundary; the current implementation stores their equivalent
 projections in the structured stage configuration and materializes
 backend-specific engine objects during stage initialization.
 The important ownership rules are:
 1. `PipelineConfig` is the source of truth for stage topology, execution type,
   model capabilities, and stage relationships.
 2. `DeployConfig` describes placement, replicas, devices, connectors, and
   deploy-time defaults; it does not redefine the model graph.
 3. CLI and Python overrides are applied at the resolution boundary, with
   per-stage overrides taking precedence over global values where supported.
 4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
   is runtime-private state, not a user configuration object.
 5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
   process that owns the corresponding engine.
 ## Main features
-The feature surface is grouped to match the
+vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
 [feature design documents](index.md#feature-design-documents). This page
 summarizes each feature's architectural role; the linked design document is
 the source for configuration, compatibility, and implementation details.
-### Runtime and stage execution
+### Performance and Acceleration
-* **Disaggregated inference:** Logical stages can run in separate processes,
+The framework achieves high performance through several optimization techniques:
  devices, or nodes while the orchestrator preserves their declared
  relationships. `OmniConnector` implementations transfer stage data and
  control-plane metadata. See [Disaggregated Inference](feature/disaggregated_inference.md).
 * **Asynchronous stage and output execution:** [Async Chunk](feature/async_chunk.md)
  forwards partial stage outputs as they become available. [Async Diffusion
  Output](feature/async_diffusion_output.md) overlaps device-to-host output
  packing with the next diffusion request, while [Async Omni Output
  Materialization](feature/omni_async_output_materialization.md) moves
  CPU-side payload construction off the AR decode critical path.
 * **Automatic prefix caching:** [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
  reuses KV-cache-aligned stage outputs and multimodal tensors for requests
  with common prefixes.
-### Communication
+* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
 * **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
 * **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
 * **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
    * **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
    * **Parallelism:** Supports TP, CP, USP, and CFG.
    * **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
    * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
    * **FusedOps:** Allows for custom and third-party integration.
-* **OmniConnector transport:** The connector contract carries tensors, KV-cache
+### Classifier-Free Guidance (CFG) Companion Flow
  data, and transport metadata across stage boundaries. The available
  implementations cover shared memory and multi-node Mooncake, Mori, and
  Yuanrong transports; see [Disaggregated Inference](feature/disaggregated_inference.md)
  for the connector choices and configuration model.
-### Diffusion acceleration
+vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
 1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
 2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
 3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
-* **Request and step batching:** [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
+### Flexibility and Usability
  defines request-batch and step-batch execution, scheduler admission, and the
  common streaming output path.
 * **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
  [Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
  [Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
  [Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
  according to the pipeline and hardware topology.
 * **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
  [Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
  provide backend and denoising-step optimizations without changing the
  stage contract.
 * **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
  resolves per-pipeline or per-component quantization configurations, while
  [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
  streams diffusion blocks from host memory within the existing parallel
  topology.
-## Interfaces
+vLLM-Omni is designed to be flexible and straightforward for users:
-The public interfaces map onto the same engine and stage boundaries:
+* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
 * **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
 * **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
 * **Streaming Outputs:** Supports streaming outputs.
 * **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
 * **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
 # Interface design
 If you use vLLM, then you know how to use vLLM-Omni from Day 0:
 <p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png">
    <img alt="vLLM-Omni interface design" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png" width=55%>
  </picture>
 </p>
 Taking **Qwen3-Omni** as an example:
 ## Offline Inference
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 ```mermaid
 flowchart LR
    offline["Offline Python<br/>Omni.generate()"] --> engine["AsyncOmniEngine"]
    online["OpenAI-compatible serving<br/>vllm serve ... --omni"] --> engine
    duplex["Experimental duplex WebSocket<br/>/v1/duplex or realtime duplex"] --> engine
    engine --> stages["Configured AR and diffusion stages"]
    stages --> result["Streaming or final multimodal output"]
 ```
-
+# Create an omni runtime with HF model name.
 ### Offline inference
 The **Omni** class provides a Python interface for offline batched inference:
 ```python
 from vllm_omni.entrypoints.omni import Omni
 omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
-om_inputs = {
+# Example prompts.
-    "prompt": prompt,
+om_inputs = {"prompt": prompt,
             "multi_modal_data": {
                 "video": video_frames,
                 "audio": audio_signal,
-    },
+             }}
 }
 # Generate texts and audio from the multi-modality inputs.
 outputs = omni.generate(om_inputs, sampling_params_list)
 ```
-### Online serving
+## Online Serving
 Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
-The OpenAI-compatible server uses the same stage configuration and engine
+```
 boundaries:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
-For example, a Qwen3-Omni chat request can contain text, image, audio, or
+Users can send requests to the server using curl:
 video content and a `sampling_params_list` for its configured stages. See the
 [Qwen3-Omni serving example](../user_guide/examples/online_serving/qwen3_omni.md)
 and the [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples)
 for complete requests.
-Some pipelines expose additional OpenAI-compatible endpoints, such as joint
+```
-video/audio generation. Endpoint support remains model-specific; consult the
+# prepare user content
-relevant model guide before assuming that every OpenAI route applies to every
+user_content='[
-pipeline.
+        {
          "type": "video_url",
          "video_url": {
            "url": "'"$SAMPLE_VIDEO_URL"'"
          }
        },
        {
          "type": "text",
          "text": "Why is this video funny?"
        }
      ]'
    sampling_params_list='[
      '"$thinker_sampling_params"',
      '"$talker_sampling_params"',
      '"$code2wav_sampling_params"'
    ]'
    mm_processor_kwargs="{}"
 # send the request
 curl -sS -X POST http://localhost:8091/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d @- <<EOF
 {
  "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
  "sampling_params_list": $sampling_params_list,
  "mm_processor_kwargs": $mm_processor_kwargs,
  "messages": [
    {
      "role": "system",
      "content": [
        {
          "type": "text",
          "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
        }
      ]
    },
    {
      "role": "user",
      "content": $user_content
    }
  ]
 }
 ```
 For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
@@ -20,12 +20,11 @@ extracts it into a shared, typed contract:
  (matching the bespoke caches it replaces), and reports per-device byte usage
  via `stats()` for observability.
-Autoregressive generation KV is *not* managed here: for DreamZero it is owned
+Attention KV is *not* managed here: for DreamZero it is owned by the
-by the AR-Diffusion engine's paged KV pool (PR
+AR-Diffusion engine's paged KV pool (PR
-[#4534](https://github.com/vllm-project/vllm-omni/pull/4534)). Cosmos3's fixed,
+[#4534](https://github.com/vllm-project/vllm-omni/pull/4534)). The manager
-encode-once UND conditioning K/V is different and is covered by the dense
+covers the model's non-KV session state; integrating the engine's per-session
-Cosmos3 adapter described below. Integrating the engine's per-session paged KV
+KV handle into the same byte accounting is RFC #4480 Phase 1.
 handle into the same byte accounting is RFC #4480 Phase 1.
 ## Enabling
@@ -159,62 +158,3 @@ of `num_frame_per_block` frames instead: identical output, bounded memory, and
 the per-step quantisation no longer scales with session length. Asking for more
 frames than the ring retains raises rather than silently returning a shorter
 window. The bespoke path is unchanged.
 ### Cosmos3
 When `enable_session_state_manager` is enabled, the Cosmos3 adapter stores the
 dense, encode-once UND text K/V for every transformer layer and active CFG
 branch, together with that branch's GEN M-RoPE cosine/sine tensors. These are
 GPU-resident tensors and can contribute to CUDA OOM. A useful description of
 the retained storage on one rank is:
 ```text
 sum_over_branches(
  sum_over_layers(storage_bytes(K) + storage_bytes(V))
  + storage_bytes(freqs_gen_cos) + storage_bytes(freqs_gen_sin)
 )
 ```
 This is not a peak-memory upper bound: construction of the next branch can
 temporarily overlap the already retained branch with UND activations, the new
 full K/V result, RoPE tensors, and generator activations. K/V slices also keep
 their underlying allocation alive, so physical storage can follow the padded
 text length rather than only the exposed trimmed length. The manager's byte
 accounting deliberately measures storage for this reason.
 Branch placement changes the per-rank cost:
 - no CFG retains one conditional branch;
 - sequential CFG can retain both conditional and unconditional branches on the
  same GPU;
 - CFG-parallel retains one branch on each CFG rank;
 - UND K/V is replicated across sequence-parallel ranks rather than divided by
  the SP degree.
 Consequently, risk increases with text batch/sequence length, media token count
 (which sizes `freqs_gen`), sequential CFG, dtype width, and concurrent live
 requests. Phase 0 records `nbytes:cuda:*` but does not enforce `byte_budget`, so
 the manager cannot prevent an allocation from exhausting the device. The
 count-based session limit is not an OOM guarantee.
 The adapter transfers ownership by tensor reference; it does not clone K/V or
 RoPE tensors. For a normal serial request this should not materially increase
 tensor storage relative to Cosmos3's existing bespoke cond/uncond cache, which
 retains the same branch values. It does add small CPU-side session/object
 metadata. This claim is based on the reference flow in the implementation;
 there is not yet a current-head GPU peak-memory A/B result for this PR.
 Request-finally cleanup drops the manager-owned session, so branches do not
 accumulate with completed request IDs. The transformer handoff field still
 references the last loaded branch until `reset_cache()` runs at the start of
 the next generation (or the pipeline is destroyed). Thus this is bounded
 post-request residency, not growth per completed request, but manager
 `stats()` after `drop_session()` no longer includes that last transformer-held
 alias. A profiler that only reads manager statistics can therefore under-report
 live process memory.
 For OOM validation, measure complete cache-off and true flag-on generations in
 otherwise identical fresh processes. Synchronize CUDA and report process-level
 peak memory plus PyTorch peak allocated and peak reserved memory. Include the
 period between requests when checking cleanup; do not infer safety from the
 manager byte total, a single denoising step, or an operator-only profile.
@@ -123,7 +123,6 @@ def extract_omni_serve_subparser_init():
                        "TrackingNamespace": None,  # type annotation only
                        "make_arg_parser": lambda parser: parser,  # no-op for doc
                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                        "_nonneg_finite_float": float,
                        "VLLM_SUBCMD_PARSER_EPILOG": "",
                        "logger": logger,
                        "DummySubparsers": DummySubparsers,
@@ -24,23 +24,6 @@ SERVING_MODE_TITLES = {
    "online_serving": "Online serving",
 }
 # Keep the Examples navigation focused on reusable modality/task entry points.
 # Model-specific scripts remain available under ``examples/`` and are surfaced
 # from these shared pages instead of getting one page per model.
 GENERAL_EXAMPLE_SLUGS = frozenset(
    {
        "image_to_image",
        "image_to_video",
        "speech_to_video",
        "text_to_audio",
        "text_to_image",
        "text_to_speech",
        "text_to_video",
        "x_to_text",
        "x_to_video_audio",
    }
 )
 def load_model_display_names() -> dict[str, str]:
    try:
@@ -63,13 +46,6 @@ def load_model_display_names() -> dict[str, str]:
 MODEL_DISPLAY_NAMES = load_model_display_names()
 def is_general_example(example: "Example") -> bool:
    """Return whether an example should get a generated serving page."""
    if example.category not in SERVING_MODE_TITLES:
        return True
    return example.path.stem in GENERAL_EXAMPLE_SLUGS
 def fix_case(text: str) -> str:
    subs = {
        "api": "API",
@@ -364,21 +340,22 @@ def update_nav_file(examples: list[Example]):
            examples_by_category[category] = []
        examples_by_category[category].append(example)
-    def category_items(category: str) -> list[dict[str, str]]:
+    # Build the new Examples section - start with preserved items
-        items = []
+    examples_section = preserved_items.copy()
-        for example in sorted(examples_by_category[category], key=lambda e: e.path.stem):
+
    # Add examples grouped by category, sorted by category name
    for category in sorted(examples_by_category.keys()):
        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
        category_items = []
        for example in category_examples:
            doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
            rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            items.append({example.nav_title: rel_path.as_posix()})
+            category_items.append({example.nav_title: rel_path.as_posix()})
        return items
-    # Build the new Examples section - start with preserved items.
+        if category_items:
-    examples_section = preserved_items.copy()
+            # Format category name (e.g., "offline_inference" -> "Offline Inference")
    for category in sorted(examples_by_category.keys()):
        items = category_items(category)
        if items:
            category_title = fix_case(category.replace("_", " ").title())
-            examples_section.append({category_title: items})
+            examples_section.append({category_title: category_items})
    # Update the nav structure
    nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
@@ -413,8 +390,6 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
        for path in category.glob("*/*.md"):
            examples.append(Example(path.parent, category.stem))
    examples = [example for example in examples if is_general_example(example)]
    # Generate the example documentation
    for example in sorted(examples, key=lambda e: e.path.stem):
        example_name = f"{example.path.stem}.md"
@@ -17,99 +17,83 @@ th {
 }
 </style>
-| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
+| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
-|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
+|--------------|--------|-------------------|------------|---------|-----|-----------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   | — |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
-| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
+| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
-| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
+| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
-| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
-| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
+| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   |
-| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
+| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
-| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
+| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
-| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
+| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
-| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
-| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
+| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | |
-| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
-| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
-| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | |
-| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
+| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
-| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
-| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
-| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
-| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
-| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
+| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
-| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
+| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
-| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
+| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
+| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
-| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
+| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
-| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
-| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
+| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
-| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
+| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
-| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
+| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
-| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
+| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
-| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
+| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
-| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
-| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
+| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
-| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
+| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
-| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
+| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
-| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
+| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
-| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
+| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
-| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
-| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
-| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
+| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
-| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
+| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
-|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
-| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
+| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
-| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
 link points to the repository recipe. `—` means that no recipe link has been
 added for that model row yet.
 For rows with an audited recipe link, the hardware checkmarks reflect hardware
 explicitly documented in that recipe. Other rows retain the implementation
 support metadata until their recipe is audited.
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
 <sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
 gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
 Other AMD SKUs (e.g. MI325X) are not listed until their
 own validation evidence is added. See the
 [MiniMax H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for
 per-architecture serving commands.
@@ -17,7 +17,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
 | `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
-| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
+| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
@@ -84,7 +84,7 @@ A backend that needs configuration exposes it as a typed field on the spec:
 When constructing `OmniDiffusionConfig` directly:
 ```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
 config = OmniDiffusionConfig(
    diffusion_attention_config=AttentionConfig(
@@ -99,27 +99,6 @@ config = OmniDiffusionConfig(
 A plain dict is also accepted and normalized to `AttentionConfig`.
 #### FlashInfer QK16/V8 quantized attention
 `FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
 ```python
 config = OmniDiffusionConfig(
    diffusion_attention_config=AttentionConfig(
        default=AttentionSpec(
            backend="FLASHINFER_ATTN",
            quant=AttnQuantSpec(
                dtype_qk="bfloat16",
                dtype_vo="fp8_e4m3",
            ),
        ),
    ),
    ...,
 )
 ```
 `dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
 ## Platform Defaults
 ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
@@ -6,7 +6,6 @@
 - [Overview](#overview)
 - [Quick Start](#quick-start)
 - [Example Script](#example-script)
 - [Request-Scoped Quality](#request-scoped-quality-minimax-h3)
 - [Acceleration Methods](#acceleration-methods)
 - [Configuration Parameters](#configuration-parameters)
 - [Best Practices](#best-practices)
@@ -141,43 +140,6 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
  --cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
 ```
 ## Request-Scoped Quality (MiniMax H3)
 MiniMax H3 can switch Cache-DiT at the request boundary without enabling a
 cache backend at server startup. Add a quality field to an online request:
 ```bash
 # Native reference path for this request.
 -F 'quality=lossless'
 # H3's conservative Cache-DiT profile for this request.
 -F 'quality=high'
 ```
 For offline inference, set the same field on the sampling parameters:
 ```python
 OmniDiffusionSamplingParams(
    quality="high",
    num_inference_steps=50,
    extra_args={"task": "t2va", "duration": 5.0, "audio_flow_shift": 3.0},
 )
 ```
 | Server startup | Request quality | Behavior |
 |---|---|---|
 | `--cache-backend cache_dit` | omitted | Use or restore the startup Cache-DiT profile |
 | `--cache-backend cache_dit` | `lossless` | Remove Cache-DiT hooks and run the reference path |
 | `--cache-backend cache_dit` | `high` | Use or install H3's conservative Cache-DiT profile |
 | no Cache-DiT | omitted or `lossless` | Run the reference path |
 | no Cache-DiT | `high` | Install H3's conservative Cache-DiT profile |
 The startup option is only needed when omitted-quality requests should use the
 server-configured Cache-DiT profile by default.
 See the [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#request-scoped-quality)
 for a complete request and measured trade-off.
 ---
 ## Acceleration Methods
@@ -159,7 +159,7 @@ The following tables show which models support each feature:
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
-| **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
+| **MiniMax-H3**               |     ❌     |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
 **Frame Interpolation Support**
@@ -33,25 +33,6 @@ This folder provides a unified CLI script for image-to-video generation using vL
 Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
 ### LingBot-Video TI2V
 LingBot-Video accepts exactly one first-frame image and uses the same dense or
 MoE checkpoint as its T2I and T2V modes:
 ```bash
 python examples/offline_inference/image_to_video/image_to_video.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --model-class-name LingBotVideoPipeline \
  --image /path/to/input.png \
  --prompt "the fox looks toward the camera" \
  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
  --output lingbot_ti2v.mp4
 ```
 The input image keeps its original geometry until the LingBot pipeline applies
 its model-specific resize and center crop.
 ## Prerequisites
 Download the example image used in the snippets below:
@@ -1,240 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
 Examples:
    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
    CUDA_VISIBLE_DEVICES=0,1 \
    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    CUDA_VISIBLE_DEVICES=0,1 \
    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
        --model /path/to/MiniMax-H3/FL2VA --mode request
 """
 from __future__ import annotations
 import argparse
 import asyncio
 import copy
 import json
 import multiprocessing
 import time
 from pathlib import Path
 from typing import Any
 import numpy as np
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 DEFAULT_PROMPTS = (
    "At night, three cats march into a bedroom playing tiny brass instruments, "
    "then abruptly file out, with synchronized room ambience.",
    "A paper boat crosses a rain-filled street while distant traffic and water sounds remain synchronized.",
 )
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
    parser.add_argument(
        "--mode",
        choices=("dlo-dp2", "request"),
        required=True,
        help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
    )
    parser.add_argument("--steps", type=int, default=2)
    parser.add_argument("--duration", type=float, default=5.0)
    parser.add_argument("--width", type=int, default=1344)
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--batch-wait-ms", type=float, default=500.0)
    parser.add_argument("--init-timeout", type=float, default=1800.0)
    parser.add_argument("--output", type=Path)
    return parser.parse_args()
 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
    common: dict[str, Any] = {
        "model": args.model,
        "trust_remote_code": True,
        "num_gpus": 2,
        "usp": 1,
        "ring": 1,
        "vae_parallel_mode": "tile",
        "vae_use_tiling": True,
        "diffusion_attention_backend": "CUDNN_ATTN",
        "request_batch_max_wait_ms": args.batch_wait_ms,
        "enforce_eager": True,
        "stage_init_timeout": args.init_timeout,
        "init_timeout": args.init_timeout,
    }
    if args.mode == "dlo-dp2":
        common.update(
            tensor_parallel_size=1,
            data_parallel_size=2,
            text_encoder_tp_size=1,
            vae_patch_parallel_size=1,
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
            dlo_resident_layers=0,
        )
    else:
        common.update(
            tensor_parallel_size=2,
            data_parallel_size=1,
            text_encoder_tp_size=2,
            vae_patch_parallel_size=2,
            enable_distributed_layerwise_offload=False,
        )
    return common
 def sampling_params(
    engine: AsyncOmni,
    args: argparse.Namespace,
    seed: int,
 ) -> list[Any]:
    params = copy.deepcopy(engine.default_sampling_params_list)
    diffusion = params[0]
    diffusion.width = args.width
    diffusion.height = args.height
    diffusion.fps = 24
    diffusion.num_inference_steps = args.steps
    diffusion.seed = seed
    diffusion.extra_args = {
        "task": "t2va",
        "duration": args.duration,
        "aspect_ratio": "16:9",
        "flow_shift": 12.0,
        "audio_flow_shift": 3.0,
    }
    return params
 async def generate_one(
    engine: AsyncOmni,
    args: argparse.Namespace,
    *,
    request_id: str,
    prompt: str,
    seed: int,
 ) -> Any:
    final_output = None
    async for output in engine.generate(
        prompt=prompt,
        request_id=request_id,
        sampling_params_list=sampling_params(engine, args, seed),
    ):
        if output.finished:
            final_output = output
    if final_output is None:
        raise RuntimeError(f"{request_id} finished without an output")
    return final_output
 def output_summary(output: Any, args: argparse.Namespace) -> dict[str, Any]:
    if not output.images:
        raise RuntimeError(f"{output.request_id} returned no video")
    frames = np.asarray(output.images[0])
    audio = np.asarray(output.multimodal_output.get("audio"))
    if frames.ndim != 4 or tuple(frames.shape[1:]) != (
        args.height,
        args.width,
        3,
    ):
        raise RuntimeError(f"{output.request_id} returned invalid video shape {tuple(frames.shape)}")
    if audio.ndim != 3 or tuple(audio.shape[:2]) != (1, 2):
        raise RuntimeError(f"{output.request_id} returned invalid audio shape {tuple(audio.shape)}")
    if args.duration == 5.0 and args.height == 768 and args.width == 1344:
        if tuple(frames.shape) != (124, 768, 1344, 3) or tuple(audio.shape) != (1, 2, 165600):
            raise RuntimeError(
                f"{output.request_id} default shape mismatch: video={tuple(frames.shape)}, audio={tuple(audio.shape)}"
            )
    return {
        "request_id": output.request_id,
        "frames_shape": list(frames.shape),
        "audio_shape": list(audio.shape),
        "peak_memory_mb": output.peak_memory_mb,
        "stage_durations": output.stage_durations,
    }
 async def run(args: argparse.Namespace) -> dict[str, Any]:
    engine = AsyncOmni(**engine_kwargs(args))
    summary: dict[str, Any] = {
        "mode": args.mode,
        "engine_kwargs": engine_kwargs(args),
    }
    try:
        if args.mode == "dlo-dp2":
            started = time.perf_counter()
            asymmetric = await asyncio.gather(
                generate_one(
                    engine,
                    args,
                    request_id="invalid-empty",
                    prompt="",
                    seed=1001,
                ),
                generate_one(
                    engine,
                    args,
                    request_id="invalid-peer",
                    prompt=DEFAULT_PROMPTS[0],
                    seed=1002,
                ),
                return_exceptions=True,
            )
            summary["asymmetric_wave_s"] = time.perf_counter() - started
            summary["asymmetric_errors"] = [
                f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
            ]
            if len(summary["asymmetric_errors"]) != 2:
                raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
        started = time.perf_counter()
        request_count = 2 if args.mode == "dlo-dp2" else 1
        outputs = await asyncio.gather(
            *(
                generate_one(
                    engine,
                    args,
                    request_id=f"recovery-{index}",
                    prompt=DEFAULT_PROMPTS[index],
                    seed=2000 + index,
                )
                for index in range(request_count)
            )
        )
        summary["valid_wave_s"] = time.perf_counter() - started
        summary["outputs"] = [output_summary(output, args) for output in outputs]
    finally:
        started = time.perf_counter()
        engine.close()
        summary["shutdown_s"] = time.perf_counter() - started
    children = multiprocessing.active_children()
    summary["active_children"] = [{"name": child.name, "pid": child.pid} for child in children]
    if children:
        raise RuntimeError(f"Worker processes remain after shutdown: {children}")
    return summary
 def main() -> None:
    args = parse_args()
    summary = asyncio.run(run(args))
    rendered = json.dumps(summary, indent=2, sort_keys=True)
    print(f"E2E_RESULT {rendered}", flush=True)
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered + "\n", encoding="utf-8")
 if __name__ == "__main__":
    main()
@@ -43,23 +43,6 @@ This folder provides several entrypoints for experimenting with text-to-image di
 Default model: `Qwen/Qwen-Image`
 ### LingBot-Video
 LingBot-Video uses the same checkpoint for image and video generation. Run the
 shared T2I entry point with:
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --prompt "a red fox standing in fresh snow" \
  --height 192 --width 320 --num-inference-steps 2 \
  --guidance-scale 3.0 --extra-body '{"flow_shift":3.0}' \
  --output lingbot_t2i.png
 ```
 The LingBot prompt builder selects the `image` output modality and requests one
 frame. Passing an input image is not supported in LingBot T2I mode.
 ## Quick Start
 ### Python API
@@ -55,21 +55,6 @@ python text_to_video.py \
  --output vace_t2v_output.mp4
 ```
 ### LingBot-Video
 The dedicated LingBot example explicitly selects the video output modality:
 ```bash
 python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
  --output lingbot_t2v.mp4
 ```
 Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
 checkpoint. Requested frame counts are rounded upward to the causal
 VAE's `4n+1` grid.
 ### LTX-2
 ```bash
@@ -20,9 +20,9 @@ from vllm_omni.platforms import current_omni_platform
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
-    parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
-    parser.add_argument("--negative-prompt", default=None)
+    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--output", default="lingbot_video_output.mp4")
    parser.add_argument("--height", type=int, default=192)
    parser.add_argument("--width", type=int, default=320)
@@ -58,12 +58,9 @@ def _to_frame_array(video: Any) -> np.ndarray:
    if isinstance(video, torch.Tensor):
        tensor = video.detach().cpu()
        if tensor.dim() == 5:
-            tensor = tensor[0]
+            tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
-        if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
+        elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
            if tensor.shape[0] in (3, 4):
            tensor = tensor.permute(1, 2, 3, 0)
            elif tensor.shape[1] in (3, 4):
                tensor = tensor.permute(0, 2, 3, 1)
        if tensor.is_floating_point():
            if float(tensor.min()) < 0.0:
                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
@@ -90,12 +87,8 @@ def main() -> None:
        parallel_config=DiffusionParallelConfig(),
    )
-    prompt: dict[str, Any] = {
+    prompt: dict[str, str] = {"prompt": args.prompt}
-        "prompt": args.prompt,
+    if args.negative_prompt:
        "modalities": ["video"],
        "num_frames": args.num_frames,
    }
    if args.negative_prompt is not None:
        prompt["negative_prompt"] = args.negative_prompt
    sampling_params = OmniDiffusionSamplingParams(
@@ -1,74 +0,0 @@
 # LingBot-Video
 LingBot-Video uses one `LingBotVideoPipeline` for text-to-image (T2I),
 text-to-video (T2V), and text-image-to-video (TI2V) generation. Both the dense
 and MoE checkpoints use the same request format.
 ## Start the server
 ```bash
 MODEL=robbyant/lingbot-video-dense-1.3b bash run_server.sh
 ```
 The MoE checkpoint uses the same server and request scripts, but requires
 substantially more GPU memory:
 ```bash
 MODEL=robbyant/lingbot-video-moe-30b-a3b bash run_server.sh
 ```
 ## Text to image
 The image endpoint selects T2I mode and always generates one frame:
 ```bash
 bash run_curl_text_to_image.sh
 ```
 The script sends a `320x192`, two-step smoke request and writes
 `lingbot_t2i.png`.
 ## Text or text-image to video
 Run the video script without an image to select T2V mode:
 ```bash
 bash run_curl_text_image_to_video.sh
 ```
 Pass a first-frame image to the same script to select TI2V mode:
 ```bash
 INPUT_IMAGE=/path/to/input.png bash run_curl_text_image_to_video.sh
 ```
 The client scripts omit the optional `model` request field, so they target
 whichever dense or MoE checkpoint the server loaded. The video example uses the
 lightweight `320x192`, 9-frame, two-step configuration.
 Until the shared `/v1/videos` reference-image resizing is removed, TI2V target
 dimensions must be sent through `extra_params`, for example
 `{"size":"320x192"}`. Do not use the top-level `size`, `width`, or `height`
 fields for TI2V because the serving layer currently applies those dimensions
 to the reference image before the model receives it. T2V requests can continue
 to use the top-level dimension fields.
 LingBot video frame counts use the causal VAE `4n+1` grid. The pipeline rounds
 any requested frame count upward to the next valid value. An explicit
 `num_frames` takes precedence over `seconds`; otherwise, the server first
 resolves `seconds * fps` and the pipeline applies the same alignment.
 Official `resolution`/`ratio` presets can be sent through `extra_params`, for
 example `{"resolution":"720p","ratio":"16:9"}`. The `2k` and `4k` entries
 only define output dimensions; whether they run successfully depends on the
 checkpoint, GPU memory, and memory optimizations available in the deployment.
 For `/v1/images/generations`, the server resolves these aliases to their final
 pixel dimensions before applying `--max-generated-image-size`. Requests above
 the configured limit return HTTP 400 before engine dispatch. LingBot produces
 one output per prompt; image requests with `n>1` are also rejected with HTTP
 400.
 LingBot TI2V accepts exactly one image reference. Image editing, video
 references, audio references, batching, and Refiner execution are not supported
 by this pipeline mode.
@@ -1,67 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 if [ -n "${INPUT_IMAGE:-}" ] && [ ! -f "${INPUT_IMAGE}" ]; then
    echo "INPUT_IMAGE does not exist: ${INPUT_IMAGE}"
    exit 1
 fi
 BASE_URL="${BASE_URL:-http://localhost:8099}"
 POLL_INTERVAL="${POLL_INTERVAL:-2}"
 if [ -n "${INPUT_IMAGE:-}" ]; then
    default_output="lingbot_ti2v.mp4"
 else
    default_output="lingbot_t2v.mp4"
 fi
 OUTPUT_PATH="${OUTPUT_PATH:-${default_output}}"
 curl_args=(
    -sS -X POST "${BASE_URL}/v1/videos"
    -H "Accept: application/json"
    -F "prompt=the subject turns toward the camera with smooth natural motion"
    -F "num_frames=9"
    -F "fps=24"
    -F "num_inference_steps=2"
    -F "guidance_scale=3.0"
    -F "flow_shift=3.0"
    -F "seed=42"
 )
 if [ -n "${INPUT_IMAGE:-}" ]; then
    curl_args+=(
        -F "input_reference=@${INPUT_IMAGE}"
        -F 'extra_params={"size":"320x192"}'
    )
 else
    curl_args+=(-F "width=320" -F "height=192")
 fi
 create_response="$(curl "${curl_args[@]}")"
 video_id="$(echo "${create_response}" | jq -er '.id')"
 echo "Created video job ${video_id}"
 while true; do
    status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
    status="$(echo "${status_response}" | jq -er '.status')"
    case "${status}" in
        queued|in_progress)
            sleep "${POLL_INTERVAL}"
            ;;
        completed)
            break
            ;;
        failed)
            echo "${status_response}" | jq .
            exit 1
            ;;
        *)
            echo "Unexpected video job status: ${status}"
            exit 1
            ;;
    esac
 done
 curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
 echo "Saved video to ${OUTPUT_PATH}"
@@ -1,20 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 BASE_URL="${BASE_URL:-http://localhost:8099}"
 OUTPUT_PATH="${OUTPUT_PATH:-lingbot_t2i.png}"
 curl -sS -X POST "${BASE_URL}/v1/images/generations" \
    -H "Content-Type: application/json" \
    -d "{
        \"prompt\": \"a red fox standing in fresh snow\",
        \"size\": \"320x192\",
        \"num_inference_steps\": 2,
        \"guidance_scale\": 3.0,
        \"flow_shift\": 3.0,
        \"seed\": 42,
        \"response_format\": \"b64_json\"
    }" | jq -er '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
 echo "Saved image to ${OUTPUT_PATH}"
@@ -1,10 +0,0 @@
 #!/bin/bash
 set -euo pipefail
 MODEL="${MODEL:-robbyant/lingbot-video-dense-1.3b}"
 PORT="${PORT:-8099}"
 vllm serve "${MODEL}" --omni \
    --model-class-name LingBotVideoPipeline \
    --port "${PORT}"
@@ -0,0 +1,247 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Verify a ModelOpt FP8 diffusers checkpoint exported by
 quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
 Three checks:
  A. transformer/config.json has a sane quantization_config block.
  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
  C. transformer disk size is materially smaller than a BF16 baseline.
 Example:
    python examples/quantization/check_modelopt_fp8_export.py \\
        --output ./hv15-480p-modelopt-fp8
    # Optional: compare disk size against a local or HF BF16 baseline.
    python examples/quantization/check_modelopt_fp8_export.py \\
        --output ./hv15-480p-modelopt-fp8 \\
        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
 """
 from __future__ import annotations
 import argparse
 import json
 import sys
 from collections import Counter
 from pathlib import Path
 SUPPORTED_ALGOS = ("FP8",)
 def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
    cfg_path = transformer_dir / "config.json"
    if not cfg_path.exists():
        print(f"[FAIL] {cfg_path} missing.")
        return 1, None
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
        return 1, None
    print(f"[A] quantization_config from {cfg_path}:")
    print(json.dumps(qc, indent=2))
    quant_algo = qc.get("quant_algo")
    issues = []
    if qc.get("quant_method") != "modelopt":
        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
    if quant_algo not in SUPPORTED_ALGOS:
        issues.append(
            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
        )
    if issues:
        print("[A] WARN — config looks incomplete:")
        for issue in issues:
            print(f"    - {issue}")
        return 2, quant_algo
    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
    return 0, quant_algo
 def _read_safetensors_header(path: Path) -> dict:
    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
    """
    import struct
    with open(path, "rb") as f:
        header_len = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(header_len))
    header.pop("__metadata__", None)
    return header
 def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
    """
    if not weight_scale_shapes:
        return "no weight_scale tensors found"
    def meaningful_dims(shape: list[int]) -> int:
        return sum(1 for d in shape if d > 1)
    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
    total = len(weight_scale_shapes)
    if per_tensor == total:
        return "per-tensor (all scalar scales)"
    if per_channel == total:
        return "per-channel (1 meaningful dim)"
    if per_block == total:
        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
 def _check_safetensors(transformer_dir: Path) -> int:
    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
    files = sorted(transformer_dir.glob("*.safetensors"))
    if not files:
        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
        return 1
    header_dtype_counts: Counter[str] = Counter()
    sample_quant_weight_keys: list[str] = []
    sample_scale_keys: list[str] = []
    weight_scale_shapes: list[list[int]] = []
    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
    for f in files:
        try:
            header = _read_safetensors_header(f)
        except Exception as exc:
            print(f"[B] WARN — could not parse header of {f}: {exc}")
            continue
        for k, info in header.items():
            dtype = info.get("dtype", "?")
            header_dtype_counts[dtype] += 1
            # FP8 stores weights as F8_E4M3 directly.
            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
                sample_quant_weight_keys.append(k)
            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
                sample_scale_keys.append(k)
            if k.endswith(".weight_scale"):
                weight_scale_shapes.append(info.get("shape", []))
                if len(sample_weight_scale_entries) < 5:
                    sample_weight_scale_entries.append((k, info.get("shape", [])))
    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
        marker = "  <-- FP8" if dtype.startswith("F8") else ""
        print(f"    {dtype:10s} {count:>6d}{marker}")
    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
    if quant_count == 0:
        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
        return 1
    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
    if sample_quant_weight_keys:
        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
    if sample_scale_keys:
        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
    for key, shape in sample_weight_scale_entries[:3]:
        print(f"      {key}: shape {shape}")
    return 0
 def _disk_size_gib(p: Path) -> float:
    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
 def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
    """Returns 0 always (informational only)."""
    quant_size = _disk_size_gib(transformer_dir)
    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
    if baseline is None:
        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
        return 0
    baseline_path = Path(baseline)
    if not baseline_path.exists():
        # Try HF download.
        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
            return 0
        print(f"    Downloading baseline transformer from HF: {baseline}")
        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
    bf16_size = _disk_size_gib(bf16_dir)
    if bf16_size == 0:
        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
        return 0
    # Expected reduction for FP8: ~50%.
    min_reduction = 30
    reduction = (1 - quant_size / bf16_size) * 100
    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
    if reduction < min_reduction:
        print(
            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
        )
    return 0
 def main() -> None:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
    p.add_argument(
        "--baseline",
        default=None,
        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
    )
    args = p.parse_args()
    out_root = Path(args.output).expanduser().resolve()
    transformer_dir = out_root / "transformer"
    if not transformer_dir.exists():
        print(f"[FAIL] {transformer_dir} does not exist.")
        sys.exit(1)
    print(f"Checking: {out_root}\n")
    fail = 0
    config_status, _quant_algo = _check_config(transformer_dir)
    fail |= config_status
    fail |= _check_safetensors(transformer_dir)
    _check_size_vs_baseline(transformer_dir, args.baseline)
    print()
    if fail == 0:
        print("=" * 60)
        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
    elif fail == 1:
        print("=" * 60)
        print("FAILURES detected — calibration may need to be re-run.")
        sys.exit(1)
    else:
        print("=" * 60)
        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
 if __name__ == "__main__":
    main()
@@ -0,0 +1,448 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
 Calibrates the DiT transformer using a small video prompt set and exports a
 diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
 The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
 (see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
 Layers kept full precision match the #2728 / #2795 pattern: modulation,
 AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
 and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
 empirically degrades under FP8 (see #2920 ablation).
 Example:
    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
        --output ./hv15-480p-modelopt-fp8 \\
        --overwrite
 """
 from __future__ import annotations
 import argparse
 import copy
 import json
 import re
 import shutil
 import sys
 from pathlib import Path
 from typing import Any
 import torch
 from diffusers import DiffusionPipeline
 DEFAULT_PROMPTS = [
    "A dog running across a field of golden wheat.",
    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
    "A close-up of a blooming rose covered in morning dew, soft natural light.",
    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
 ]
 def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
    p.add_argument("--height", type=int, default=480)
    p.add_argument("--width", type=int, default=832)
    p.add_argument(
        "--num-frames",
        type=int,
        default=33,
        help="Frames per calibration sample. 33 matches the typical short benchmark.",
    )
    p.add_argument("--guidance-scale", type=float, default=6.0)
    p.add_argument(
        "--calib-steps",
        type=int,
        default=10,
        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
    )
    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--prompt",
        action="append",
        default=[],
        help="Custom calibration prompt. Repeat to provide multiple.",
    )
    p.add_argument(
        "--quantize-mha",
        action="store_true",
        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
    )
    p.add_argument(
        "--weight-block-size",
        type=str,
        default=None,
        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
    )
    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
    return p
 def _parse_block_size(spec: str | None) -> list[int] | None:
    if spec is None:
        return None
    parts = [int(x) for x in spec.split(",") if x.strip()]
    if len(parts) != 2:
        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
    return parts
 def _require_modelopt() -> Any:
    try:
        import modelopt.torch.quantization as mtq
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "NVIDIA ModelOpt is not installed. Install with:\n"
            "  pip install 'nvidia-modelopt[all]'\n"
            f"Original error: {exc}"
        ) from exc
    return mtq
 def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
    model_path = args.model
    output_dir = Path(args.output).expanduser().resolve()
    if output_dir.exists():
        if not args.overwrite:
            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
        shutil.rmtree(output_dir)
    return model_path, output_dir
 def _select_dtype(name: str) -> torch.dtype:
    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
 def _build_prompts(args: argparse.Namespace) -> list[str]:
    prompts = args.prompt or DEFAULT_PROMPTS
    if args.calib_size <= 0:
        raise SystemExit("--calib-size must be positive.")
    if len(prompts) < args.calib_size:
        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
        prompts = (prompts * repeats)[: args.calib_size]
    return prompts[: args.calib_size]
 # Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
 # - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
 # - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
 # - proj_out: final output projection
 # - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
 def _filter_func_hv15(name: str) -> bool:
    pattern = re.compile(
        r"(proj_out.*|"
        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
        r"time_embed|cond_type_embed|"
        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
        r"token_refiner).*)"
    )
    return pattern.match(name) is not None
 def _mha_filter_func(name: str) -> bool:
    pattern = re.compile(
        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
    )
    return pattern.match(name) is not None
 def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
    if not hasattr(mtq, "disable_quantizer"):
        return
    mtq.disable_quantizer(backbone, _filter_func_hv15)
    if not quantize_mha:
        mtq.disable_quantizer(backbone, _mha_filter_func)
 def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
    if hasattr(pipe, "set_progress_bar_config"):
        pipe.set_progress_bar_config(disable=True)
    pipe.to("cuda")
    return pipe
 def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
    generator = torch.Generator(device="cuda")
    # Try to set guidance on the pipeline's guider object up front (modern
    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
    # back silently — calibration uses whatever default the pipeline ships with.
    guider = getattr(pipe, "guider", None)
    if guider is not None and hasattr(guider, "guidance_scale"):
        try:
            guider.guidance_scale = args.guidance_scale
        except Exception:
            pass
    base_kwargs = dict(
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.calib_steps,
        output_type="latent",
    )
    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
        with torch.inference_mode():
            for idx, prompt in enumerate(prompts):
                generator.manual_seed(args.seed + idx)
                # Try with guidance_scale first; fall back without on TypeError
                # for pipelines (like HV-1.5) that take CFG via guider config.
                try:
                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
                except TypeError as exc:
                    if "guidance_scale" not in str(exc):
                        raise
                    pipe(prompt=prompt, generator=generator, **base_kwargs)
    return forward_loop
 def _summarize_export(output_dir: Path) -> None:
    cfg_path = output_dir / "transformer" / "config.json"
    if not cfg_path.exists():
        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
        return
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
        return
    print("Export summary:")
    print(f"  quant_method: {qc.get('quant_method')}")
    print(f"  quant_algo:   {qc.get('quant_algo')}")
    producer = qc.get("producer")
    if isinstance(producer, dict):
        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
    print(f"  config path:  {cfg_path}")
 def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
    """Convert in-memory weights of quantized modules to actual FP8 storage.
    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
    in ModelOpt's recognized-model registry), so we must call the per-weight
    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
    helper.
    """
    from modelopt.torch.export.quant_utils import (
        QUANTIZATION_NONE,
        get_quantization_format,
        quantizer_attr_names,
        weight_attr_names,
    )
    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
    exported = 0
    for name, module in backbone.named_modules():
        try:
            quantization_format = get_quantization_format(module)
        except Exception as exc:
            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
            continue
        if quantization_format == QUANTIZATION_NONE:
            continue
        for weight_name in weight_attr_names(module):
            quantizer_attrs = quantizer_attr_names(weight_name)
            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
                continue
            _export_quantized_weight(module, dtype, weight_name)
            exported += 1
    return exported
 def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
    `weight_block_size` is given, advertise block-wise weight quantization in
    the saved metadata (so consumers know to expect multi-element scale tensors).
    """
    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
    if weight_block_size is not None:
        weights_cfg["strategy"] = "block"
        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
    return {
        "config_groups": {
            "group_0": {
                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
                "weights": weights_cfg,
                "targets": ["Linear"],
            }
        },
        "ignore": [
            "context_embedder*",
            "context_embedder_2*",
            "cond_type_embed*",
            "image_embedder*",
            "norm1.linear*",
            "norm1_context.linear*",
            "norm2*",
            "norm2_context*",
            "norm_out*",
            "proj_out*",
            "time_embed*",
            "token_refiner*",
            "x_embedder*",
        ],
        "producer": {"name": "modelopt"},
        "quant_algo": "FP8",
        "quant_method": "modelopt",
    }
 def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
    cfg_path = output_dir / "transformer" / "config.json"
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
    existing = cfg.get("quantization_config")
    if isinstance(existing, dict):
        producer = existing.get("producer")
        if isinstance(producer, dict):
            new_qc["producer"] = producer
    cfg["quantization_config"] = new_qc
    with cfg_path.open("w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)
 def _save_pipeline_with_fp8_transformer(
    pipe: DiffusionPipeline,
    model_path: str,
    output_dir: Path,
    max_shard_size: str = "5GB",
 ) -> None:
    """Save the pipeline with the (now FP8) transformer.
    Copies the source directory verbatim except for `transformer/`, then
    saves the transformer with quantizers hidden so the state dict contains
    only the FP8 weights + scale tensors.
    """
    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
    src = Path(model_path)
    if not src.exists():
        from huggingface_hub import snapshot_download
        src = Path(snapshot_download(model_path))
    if output_dir.exists():
        shutil.rmtree(output_dir)
    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
    transformer_out = output_dir / "transformer"
    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
    # nn.Module (transformer), not the diffusers Pipeline wrapper.
    with hide_quantizers_from_state_dict(pipe.transformer):
        pipe.transformer.save_pretrained(
            str(transformer_out),
            safe_serialization=True,
            max_shard_size=max_shard_size,
        )
 def main() -> None:
    args = _build_parser().parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
    mtq = _require_modelopt()
    model_path, output_dir = _ensure_paths(args)
    dtype = _select_dtype(args.dtype)
    prompts = _build_prompts(args)
    print("Quantization plan:")
    weight_block_size = _parse_block_size(args.weight_block_size)
    print(f"  input:           {args.model}")
    print(f"  output:          {output_dir}")
    print(f"  dtype:           {dtype}")
    print(f"  height/width:    {args.height}x{args.width}")
    print(f"  num_frames:      {args.num_frames}")
    print(f"  calib_size:      {len(prompts)}")
    print(f"  calib_steps:     {args.calib_steps}")
    print(f"  quantize_mha:    {args.quantize_mha}")
    print(
        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
    )
    pipe = _load_pipeline(model_path, dtype)
    backbone = pipe.transformer
    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
    if weight_block_size is not None:
        # Switch from per-tensor (default) to block-wise weight quantization.
        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
        quant_config["quant_cfg"]["*weight_quantizer"] = {
            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
        }
        print(
            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
        )
    forward_loop = _build_forward_loop(pipe, args, prompts)
    quantized = mtq.quantize(backbone, quant_config, forward_loop)
    if quantized is not None:
        pipe.transformer = quantized
        backbone = quantized
    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
    print("so we have to call the per-weight export helper ourselves)...")
    exported = _force_export_quantized_weights(backbone, dtype)
    print(f"  -> {exported} weights converted to FP8 in memory")
    if exported == 0:
        raise SystemExit(
            "No quantized weights were exported. Calibration may have skipped every layer "
            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
            "weight quantizers."
        )
    print("\nSaving pipeline with FP8 transformer...")
    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
    print(f"Saved to: {output_dir}")
    _summarize_export(output_dir)
    print("\nNext: validate the checkpoint with vllm-omni:")
    print(
        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
        f"    --model {output_dir} \\\n"
        "    --quantization fp8 \\\n"
        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
        "    --enforce-eager"
    )
    print(
        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
        "checkpoint's config.json has modelopt metadata.)"
    )
 if __name__ == "__main__":
    main()
@@ -0,0 +1,471 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
 Calibrates the DiT transformer using a small video prompt set and exports a
 diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
 The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
 (see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
 Layers kept full precision match the #2728 / #2795 pattern: condition embedder
 (time/text/image), patch embedding, modulation (scale_shift_table), final
 norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
 are quantized — static calibration handles the numerics that online FP8
 couldn't (see #2920 ablation).
 Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
 fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
 Example:
    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
        --output ./wan22-ti2v-modelopt-fp8 \\
        --overwrite
 """
 from __future__ import annotations
 import argparse
 import copy
 import json
 import re
 import shutil
 import sys
 from pathlib import Path
 from typing import Any
 import torch
 from diffusers import DiffusionPipeline
 DEFAULT_PROMPTS = [
    "A dog running across a field of golden wheat.",
    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
    "A close-up of a blooming rose covered in morning dew, soft natural light.",
    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
 ]
 def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
    p.add_argument(
        "--num-frames",
        type=int,
        default=49,
        help="Frames per calibration sample. 49 matches the typical short benchmark; "
        "use 17 to reduce memory pressure during calibration.",
    )
    p.add_argument("--guidance-scale", type=float, default=5.0)
    p.add_argument(
        "--calib-steps",
        type=int,
        default=10,
        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
    )
    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--prompt",
        action="append",
        default=[],
        help="Custom calibration prompt. Repeat to provide multiple.",
    )
    p.add_argument(
        "--quantize-mha",
        action="store_true",
        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
        "sequences amplified FP8 drift in the online ablation (see #2920).",
    )
    p.add_argument(
        "--weight-block-size",
        type=str,
        default=None,
        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
        "for the HV-1.5 investigation status before relying on this.",
    )
    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
    return p
 def _parse_block_size(spec: str | None) -> list[int] | None:
    if spec is None:
        return None
    parts = [int(x) for x in spec.split(",") if x.strip()]
    if len(parts) != 2:
        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
    return parts
 def _require_modelopt() -> Any:
    try:
        import modelopt.torch.quantization as mtq
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "NVIDIA ModelOpt is not installed. Install with:\n"
            "  pip install 'nvidia-modelopt[all]'\n"
            f"Original error: {exc}"
        ) from exc
    return mtq
 def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
    model_path = args.model
    output_dir = Path(args.output).expanduser().resolve()
    if output_dir.exists():
        if not args.overwrite:
            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
        shutil.rmtree(output_dir)
    return model_path, output_dir
 def _select_dtype(name: str) -> torch.dtype:
    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
 def _build_prompts(args: argparse.Namespace) -> list[str]:
    prompts = args.prompt or DEFAULT_PROMPTS
    if args.calib_size <= 0:
        raise SystemExit("--calib-size must be positive.")
    if len(prompts) < args.calib_size:
        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
        prompts = (prompts * repeats)[: args.calib_size]
    return prompts[: args.calib_size]
 # Layers to KEEP at full precision. Wan2.2's module naming:
 # - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
 # - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
 # - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
 # - norm_out: AdaLayerNorm final
 # - proj_out: final nn.Linear
 # - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
 def _filter_func_wan22(name: str) -> bool:
    pattern = re.compile(
        r"(proj_out.*|"
        r".*(condition_embedder|patch_embedding|"
        r"norm_out|scale_shift_table|"
        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
    )
    return pattern.match(name) is not None
 def _mha_filter_func(name: str) -> bool:
    pattern = re.compile(
        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
    )
    return pattern.match(name) is not None
 def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
    if not hasattr(mtq, "disable_quantizer"):
        return
    mtq.disable_quantizer(backbone, _filter_func_wan22)
    if not quantize_mha:
        mtq.disable_quantizer(backbone, _mha_filter_func)
 def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
    if hasattr(pipe, "set_progress_bar_config"):
        pipe.set_progress_bar_config(disable=True)
    pipe.to("cuda")
    return pipe
 def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
    generator = torch.Generator(device="cuda")
    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
    guider = getattr(pipe, "guider", None)
    if guider is not None and hasattr(guider, "guidance_scale"):
        try:
            guider.guidance_scale = args.guidance_scale
        except Exception:
            pass
    base_kwargs = dict(
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.calib_steps,
        output_type="latent",
    )
    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
        with torch.inference_mode():
            for idx, prompt in enumerate(prompts):
                generator.manual_seed(args.seed + idx)
                # Try with guidance_scale first; fall back without on TypeError
                # for pipelines that take CFG via guider config only.
                try:
                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
                except TypeError as exc:
                    if "guidance_scale" not in str(exc):
                        raise
                    pipe(prompt=prompt, generator=generator, **base_kwargs)
    return forward_loop
 def _summarize_export(output_dir: Path) -> None:
    cfg_path = output_dir / "transformer" / "config.json"
    if not cfg_path.exists():
        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
        return
    with cfg_path.open(encoding="utf-8") as f:
        cfg = json.load(f)
    qc = cfg.get("quantization_config")
    if not isinstance(qc, dict):
        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
        return
    print("Export summary:")
    print(f"  quant_method: {qc.get('quant_method')}")
    print(f"  quant_algo:   {qc.get('quant_algo')}")
    producer = qc.get("producer")
    if isinstance(producer, dict):
        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
    print(f"  config path:  {cfg_path}")
 def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
    """Convert in-memory weights of quantized modules to actual FP8 storage.
    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
    in ModelOpt's recognized-model registry), so we must call the per-weight
    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
    calibration helpers.
    """
    from modelopt.torch.export.quant_utils import (
        QUANTIZATION_NONE,
        get_quantization_format,
        quantizer_attr_names,
        weight_attr_names,
    )
    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
    exported = 0
    for name, module in backbone.named_modules():
        try:
            quantization_format = get_quantization_format(module)
        except Exception as exc:
            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
            continue
        if quantization_format == QUANTIZATION_NONE:
            continue
        for weight_name in weight_attr_names(module):
            quantizer_attrs = quantizer_attr_names(weight_name)
            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
                continue
            _export_quantized_weight(module, dtype, weight_name)
            exported += 1
    return exported
 def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
    if weight_block_size is not None:
        weights_cfg["strategy"] = "block"
        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
    return {
        "config_groups": {
            "group_0": {
                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
                "weights": weights_cfg,
                "targets": ["Linear"],
            }
        },
        "ignore": [
            "condition_embedder*",
            "norm_out*",
            "output_scale_shift_prepare*",
            "patch_embedding*",
            "proj_out*",
            "scale_shift_table*",
            "timestep_proj_prepare*",
        ],
        "producer": {"name": "modelopt"},
        "quant_algo": "FP8",
        "quant_method": "modelopt",
    }
 def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
    for sub in ("transformer", "transformer_2"):
        cfg_path = output_dir / sub / "config.json"
        if not cfg_path.exists():
            continue
        with cfg_path.open(encoding="utf-8") as f:
            cfg = json.load(f)
        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
        existing = cfg.get("quantization_config")
        if isinstance(existing, dict):
            producer = existing.get("producer")
            if isinstance(producer, dict):
                new_qc["producer"] = producer
        cfg["quantization_config"] = new_qc
        with cfg_path.open("w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2)
 def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
    """Return [(attr_name, module)] for every DiT on the pipeline.
    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
    (high-noise). TI2V-5B has only ``transformer``.
    """
    out = []
    for attr in ("transformer", "transformer_2"):
        mod = getattr(pipe, attr, None)
        if mod is not None:
            out.append((attr, mod))
    return out
 def _save_pipeline_with_fp8_transformer(
    pipe: DiffusionPipeline,
    model_path: str,
    output_dir: Path,
    max_shard_size: str = "5GB",
 ) -> None:
    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
    as the single-DiT TI2V-5B.
    """
    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
    src = Path(model_path)
    if not src.exists():
        from huggingface_hub import snapshot_download
        src = Path(snapshot_download(model_path))
    if output_dir.exists():
        shutil.rmtree(output_dir)
    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
    for attr, backbone in _list_transformers(pipe):
        transformer_out = output_dir / attr
        # Pass the nn.Module (transformer), not the Pipeline wrapper.
        with hide_quantizers_from_state_dict(backbone):
            backbone.save_pretrained(
                str(transformer_out),
                safe_serialization=True,
                max_shard_size=max_shard_size,
            )
 def main() -> None:
    args = _build_parser().parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
    mtq = _require_modelopt()
    model_path, output_dir = _ensure_paths(args)
    dtype = _select_dtype(args.dtype)
    prompts = _build_prompts(args)
    weight_block_size = _parse_block_size(args.weight_block_size)
    print("Quantization plan:")
    print(f"  input:           {args.model}")
    print(f"  output:          {output_dir}")
    print(f"  dtype:           {dtype}")
    print(f"  height/width:    {args.height}x{args.width}")
    print(f"  num_frames:      {args.num_frames}")
    print(f"  calib_size:      {len(prompts)}")
    print(f"  calib_steps:     {args.calib_steps}")
    print(f"  quantize_mha:    {args.quantize_mha}")
    print(
        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
    )
    pipe = _load_pipeline(model_path, dtype)
    transformers = _list_transformers(pipe)
    if not transformers:
        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
    if weight_block_size is not None:
        quant_config["quant_cfg"]["*weight_quantizer"] = {
            "num_bits": (4, 3),
            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
        }
        print(
            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
        )
    forward_loop = _build_forward_loop(pipe, args, prompts)
    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
    # on every attached transformer (A14B: transformer + transformer_2) get
    # calibrated during the same passes.
    for attr, backbone in transformers:
        print(f"\nQuantizing {attr}...")
        quantized = mtq.quantize(backbone, quant_config, forward_loop)
        if quantized is not None:
            setattr(pipe, attr, quantized)
        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
    # inside the loop would convert transformer's weights to real FP8, and the
    # next transformer's forward_loop (full pipeline) would then fake-quantize an
    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
    print("so we have to call the per-weight export helper ourselves)...")
    total_exported = 0
    for attr, backbone in _list_transformers(pipe):
        exported = _force_export_quantized_weights(backbone, dtype)
        print(f"  -> {exported} weights converted to FP8 in {attr}")
        total_exported += exported
    if total_exported == 0:
        raise SystemExit(
            "No quantized weights were exported. Calibration may have skipped every layer "
            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
            "weight quantizers."
        )
    print("\nSaving pipeline with FP8 transformer(s)...")
    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
    print(f"Saved to: {output_dir}")
    _summarize_export(output_dir)
    print("\nNext: validate the checkpoint with vllm-omni:")
    print(
        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
        f"    --model {output_dir} \\\n"
        "    --quantization fp8 \\\n"
        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
        "    --output outputs/wan22_modelopt_fp8.mp4"
    )
    print(
        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
        "checkpoint's config.json has modelopt metadata.)"
    )
 if __name__ == "__main__":
    main()
@@ -1,346 +0,0 @@
 # MiniMax-H3 on RTX PRO 6000 Blackwell GPUs
 This recipe runs MiniMax-H3 in BF16 on 96 GiB RTX PRO 6000 Blackwell GPUs.
 The additional HBM over the RTX PRO 5000 profile removes the four-GPU
 minimum: two GPUs are enough to keep the model resident with TP2. Four and
 eight GPUs raise the Ulysses degree to shard the attention sequence further.
 CPU offload and distributed layerwise offload are not required in any of
 these configurations.
 Validated on:
 - Host: YLX Y762
 - GPUs: 8 × RTX PRO 6000 Blackwell (96 GiB)
 - Device order: default (`CUDA_VISIBLE_DEVICES` not set)
 - Driver / CUDA: 580.105.08 / 13.0
 - Container image: `vllm/vllm-omni:minimax-h3`
 - Workload: T2VA, 1344×768, `duration=5.0`, 50 steps, `flow_shift=12`, `seed=1101`
 ## Capacity requirements
 | Resource | Two GPUs | Four GPUs | Eight GPUs |
 | --- | ---: | ---: | ---: |
 | GPU HBM | 96 GiB per GPU | 96 GiB per GPU | 96 GiB per GPU |
 | Checkpoint storage | 135 GiB per partition | 135 GiB per partition | 135 GiB per partition |
 | Available system RAM | 200 GiB minimum | 200 GiB minimum | 200 GiB minimum |
 | Recommended system RAM | 384 GiB | 384 GiB | 384 GiB |
 | Measured peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
 `FL2VA` and `Ref2VA` are separate checkpoint partitions. Start one server at
 a time on a host sized for the minimum system-memory requirement, or pass
 `--task-type fl2va` or `--task-type ref2va` to download and load a single
 partition. The two-server layout described in the eight-GPU section doubles
 both the storage and the host-memory requirement.
 ## Parallelism summary
 | Configuration | TP | Ulysses | Text-encoder TP | VAE patch parallel |
 | --- | ---: | ---: | ---: | ---: |
 | Two GPUs | 2 | 1 | 2 | 2 |
 | Four GPUs | 2 | 2 | 4 | 4 |
 | Eight GPUs | 2 | 4 | 8 | 8 |
 | Eight GPUs, headroom variant | 4 | 2 | 8 | 8 |
 Per-GPU weight residency is set by the tensor-parallel degree alone. Raising
 the Ulysses degree lowers activation memory and per-step latency but leaves
 weight residency unchanged. Ring parallelism stays at 1, and H3 is
 CFG-distilled, so `--cfg-parallel-size` must also remain 1.
 ## PCIe topology and GPU order
 RTX PRO 6000 Blackwell does not provide NVLink. Before starting the server,
 inspect the host topology:
 ```bash
 nvidia-smi topo -m
 nvidia-smi nvlink -s
 ```
 Tensor-parallel groups are consecutive ranks. Ulysses groups are strided by
 the tensor-parallel degree. Confirm this grouping against the server log
 before tuning the device order on a new host.
 For the two-GPU configuration, pick a single `PXB` pair on one NUMA node.
 For the four-GPU configuration, pick two `PXB` pairs on the same node and
 order them so that each Ulysses group lands inside one pair. If physical
 GPUs `(0,1)` and `(2,3)` are the two local pairs, the order is
 `CUDA_VISIBLE_DEVICES=0,2,1,3`. Do not copy these IDs blindly: reproduce the
 same relationship on the target host.
 Device-order pinning and NUMA binding are tuning knobs, not prerequisites.
 The measurements in this recipe were taken without either, so any gain from
 them is additional to the numbers reported below and is host-specific.
 ## Two-GPU serving configuration
 Two 96 GiB GPUs hold the BF16 model with TP2 alone. There is no Ulysses
 group, so the DiT attention sequence is not sharded and activation memory is
 higher than in the four-GPU profile at the same output shape. Re-measure
 peak memory before increasing the output shape, the reference-input length,
 or concurrency.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 2 \
  --tensor-parallel-size 2 \
  --usp 1 \
  --ring 1 \
  --text-encoder-tp-size 2 \
  --vae-patch-parallel-size 2 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ## Four-GPU serving configuration
 Four GPUs keep TP2 for the weight shard and add Ulysses2, which shards the
 attention sequence and lowers per-GPU activation memory. Text-encoder TP4
 and VAE patch parallelism 4 spread the Qwen3-VL encoder and tiled decode
 across all four GPUs. TP1 is not an option on this card: an unsharded BF16
 partition does not fit in 96 GiB.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 4 \
  --tensor-parallel-size 2 \
  --usp 2 \
  --ring 1 \
  --text-encoder-tp-size 4 \
  --vae-patch-parallel-size 4 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ## Eight-GPU serving configuration
 Eight GPUs extend the Ulysses degree to 4 while keeping TP2. This does not
 lower per-GPU weight residency relative to the four-GPU profile; it lowers
 activation memory and per-step latency by sharding the attention sequence
 four ways.
 ```bash
 export MODEL_ROOT=/path/to/MiniMax-H3
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8000
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 8 \
  --tensor-parallel-size 2 \
  --usp 4 \
  --ring 1 \
  --text-encoder-tp-size 8 \
  --vae-patch-parallel-size 8 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 ### Rank ordering across two sockets
 With TP2 and Ulysses4 the two Ulysses groups are ranks `(0,2,4,6)` and
 `(1,3,5,7)`, and the four tensor-parallel pairs are `(0,1)`, `(2,3)`,
 `(4,5)`, and `(6,7)`. On a dual-socket host these two collectives cannot
 both be socket-local. If you choose to pin the device order, the trade-off
 is:
 | Device order | Socket-local collective | Crosses sockets |
 | --- | --- | --- |
 | `0,4,1,5,2,6,3,7` | Ulysses all-to-all | TP all-reduce |
 | `0,1,2,3,4,5,6,7` | TP all-reduce | Ulysses all-to-all |
 The first order assumes physical GPUs `0-3` are on NUMA node 0 and `4-7` on
 node 1. Reproduce that relationship on the target host rather than copying
 the IDs. The reported measurements use the default order and no NUMA
 binding; the two orders were not compared on this host.
 ### Headroom variant: TP4 with Ulysses2
 Raising the tensor-parallel degree to 4 halves per-GPU weight residency and
 leaves room for long Ref2VA references, larger output shapes, or
 concurrency greater than one. It also adds a four-way all-reduce per DiT
 block over a link without NVLink. Use it when the TP2 profile runs out of
 headroom, not as the default. Replace the parallel flags above with:
 ```bash
  --tensor-parallel-size 4 \
  --usp 2 \
 ```
 ### Two independent four-GPU servers
 An eight-GPU host can instead run one FL2VA server and one Ref2VA server
 side by side, each pinned to its own NUMA node with the four-GPU
 configuration. This removes every cross-socket collective and lifts the
 one-server-at-a-time restriction, at the cost of doubling the requirements
 to 400 GiB available system RAM and 270 GiB of model storage.
 Use `CUDA_VISIBLE_DEVICES=0,2,1,3`, `--cpunodebind=0 --membind=0`, and
 `--port 8091` for the FL2VA server. Use `CUDA_VISIBLE_DEVICES=4,6,5,7`,
 `--cpunodebind=1 --membind=1`, and `--port 8092` for the Ref2VA server.
 This layout is described for completeness and was not part of the measured
 runs.
 ## Shared serving notes
 Do not add `--enforce-eager` for a performance run. Warm the server before
 measuring so regional compilation is outside the measured request. On this
 host the first request after startup ran 19% slower than the steady state
 (107.58 s versus 90.48 s on eight GPUs); two warmup requests were enough to
 converge.
 For Ref2VA on a single-server layout, stop the FL2VA server and restart the
 same command with `MODEL="${MODEL_ROOT}/Ref2VA"`.
 ## Attention backend
 RTX PRO 6000 Blackwell is an SM120 part, so the datacenter-Blackwell
 `TRTLLM_ATTN` default does not apply here. Every configuration above selects
 cuDNN BF16 attention explicitly. This keeps the recipe independent of
 platform-default backend changes and does not depend on the experimental
 SM120 kernel.
 ## Optional: online FP8
 `--quantization fp8` quantizes eligible DiT linears at load time and is
 compatible with tensor parallelism and VAE tiling. Use it when the resident
 BF16 peak leaves too little headroom for long Ref2VA references, a larger
 output shape, or concurrency greater than one. It cannot be combined with
 layerwise offload.
 ## Target-hardware validation
 T2VA, 1344×768, `duration=5.0`, `fps=24`, 50 steps, `flow_shift=12`,
 `seed=1101`, BF16, `CUDNN_ATTN`, tiled VAE, one request at a time. Servers
 were started with the commands above — default device order, no NUMA
 binding — plus `--enable-diffusion-pipeline-profiler`. Two warmup requests
 preceded each measured request.
 | Measurement | Two GPUs | Four GPUs | Eight GPUs |
 | --- | ---: | ---: | ---: |
 | Topology | TP2 × USP1 | TP2 × USP2 | TP2 × USP4 |
 | Text encode | 0.040 s | 0.044 s | 0.035 s |
 | Denoise, 50 steps | 278.55 s | 168.73 s | 87.90 s |
 | Per step | 5.571 s | 3.375 s | 1.758 s |
 | VAE decode | 5.396 s | 2.791 s | 1.798 s |
 | Client E2E | 284.76 s | 172.32 s | 90.48 s |
 | Peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
 | Headroom below 96 GiB | 18.5 GiB | 29.6 GiB | 34.9 GiB |
 Stage times are read from the `X-Stage-Durations` response header of
 `/v1/videos/sync`. Peak memory is the maximum of `nvidia-smi
 --query-gpu=memory.used` sampled at 1 Hz across every device for the
 duration of the measured request. Per step is denoise wall time divided by
 the 50 requested steps. Stage times sum to roughly 0.75 s less than
 end-to-end in all three configurations; queueing, result transfer, and MP4
 muxing sit outside the profiled stages.
 ### Scaling behaviour
 Denoise speedup is not uniform across the sweep: 1.65× from two to four
 GPUs (83% efficiency) but 1.92× from four to eight (96%). Two GPUs run
 Ulysses1, so there is no sequence parallelism at all; going to four GPUs
 pays the one-time cost of introducing the all-to-all, while going from four
 to eight only widens an all-to-all that already exists. Expect the second
 doubling to pay off better than the first on this topology.
 VAE decode behaves in the opposite direction — 1.93× from two to four GPUs
 but only 1.55× from four to eight — as patch parallelism produces smaller
 tiles and tile-boundary overhead grows.
 Text encoding is 0.04 s in every configuration and is not worth optimising;
 denoise is 96-98% of end-to-end throughout.
 ### Per-GPU memory model
 Peak memory falls by 11.05 GiB from two to four GPUs and by a further
 5.37 GiB from four to eight — halving each time the Ulysses degree doubles,
 because all three configurations shard the DiT with TP2 and hold identical
 weights per GPU. Fitting the first two points gives:
 ```
 peak HBM per GPU ≈ 55.4 GiB + 22.1 GiB / ulysses_degree
 ```
 That fit predicts 60.9 GiB at Ulysses4 against 61.07 GiB measured, a 0.25%
 error. The practical consequence is a floor near 55 GiB per GPU: adding
 GPUs past eight will not reduce the per-GPU requirement further. To serve
 on smaller cards, raise the tensor-parallel degree instead.
 Re-measure memory for longer reference inputs, concurrency greater than
 one, or a different output shape.
 ### Not yet measured
 Ref2VA latency and memory have not been measured on this host. Ref2VA
 drives the Qwen3-VL presentation from tens of tokens for text-only prompts
 to several thousand with image or video references, so T2VA numbers do not
 transfer to it. The eight-GPU device-order comparison in the section above
 was also not run.
 ## T2VA request example
 `t2va` requires an explicit `aspect_ratio` even when `width` and `height`
 are supplied.
 ```bash
 export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
 curl -sS --max-time 1800 -X POST "${API_URL}" \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=1344' \
  -F 'height=768' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":5.0,"audio_flow_shift":3.0}' \
  -o t2va.mp4
 ```
 To collect the stage breakdown and peak memory for a request, add
 `--enable-diffusion-pipeline-profiler` to the server command and `-D
 headers.txt` to the curl invocation, then read `X-Stage-Durations` and
 `X-Peak-Memory-MB` from the saved headers.
 ## Ref2VA request example
 `ref2va` defaults to a 16:9 aspect ratio, so `aspect_ratio` is optional
 here.
 ```bash
 curl -X POST "http://127.0.0.1:${PORT}/v1/videos/sync" \
  --fail-with-body -w '\nHTTP %{http_code}\n' \
  --max-time 1200 \
  -F "input_reference=@/root/hand.jpg;type=image/jpeg" \
  -F "audio_reference=</root/audio_ref.json" \
  -F "prompt=2D动画融合在一起的影像。夕阳余晖残留在窗边，生活感十足的小厨房里有旧木桌、洗到一半的马克杯、起雾的玻璃瓶、悬挂的抹布。画面带有智能手机单手拍摄的手抖、近距离对焦的犹豫、逆光曝光波动。要像在家中慌忙拍下某个不可思议事件的自然质感，不要广告影像的精心整理。声音只用厨房环境声与手绘生物柔和的电子音、小小的叫声。" \
  -F 'width=1344' -F 'height=768' -F 'fps=24' \
  -F 'num_inference_steps=60' -F 'flow_shift=12' -F 'seed=1101' \
  -F 'extra_params={"task":"ref2va","duration":8,"audio_flow_shift":3.0}' \
  -o /root/out_ref2va.mp4
 ```
@@ -1,262 +0,0 @@
 # MiniMax-H3 on DGX Spark (GB10)
 This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
 resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
 recipes it uses **no offload of any kind** — see the capacity note below.
 Validated on:
 - Host: DGX Spark (GB10), aarch64
 - GPUs: 1 (unified memory)
 - Driver: <FILL_DRIVER_VERSION>
 - vLLM: 0.26.0
 - vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
 - Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
  `seed=1101`
 ## Capacity requirements
 | Resource | DGX Spark (GB10) |
 | --- | ---: |
 | Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
 | GPUs | 1 |
 | Checkpoint storage | 135 GiB per partition |
 | Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
 | Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
 On GB10 the CPU and GPU share one physical memory pool, which changes the capacity
 math relative to the RTX and datacenter recipes:
 - **Do not use `--enable-distributed-layerwise-offload`.** DLO stages rank-local
  weights into pinned host memory, which on GB10 is the same pool the weights are
  already in. Peak usage roughly doubles and the process is killed by the OOM
  killer (`Exit code: -9`) shortly after `Enabling offloader backend`.
 - **Do not use `--enable-cpu-offload`.** Moving weights from "VRAM" to "host RAM"
  frees nothing here.
 - **`--quantization fp8` is mandatory.** A BF16 partition is 135 GiB and does not
  fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
  Qwen3-VL text encoder and both VAEs stay BF16.
 With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
 23 GiB of headroom for the OS, page cache, and any other process on the box.
 Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
 and the allocator figure understates true pool usage.
 Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
 `MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
 a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
 roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
 which is the CUDA context, non-PyTorch allocations, and page cache the header
 does not count.
 Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
 second process, and adding reference images, reference videos, or a larger output
 shape will run into the OOM killer.
 ## One GB10: 960x576, 8 seconds
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
 FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=7200 \
 vllm serve /path/to/MiniMax-H3/FL2VA \
  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
  --init-timeout 3600 \
  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
  --vae-parallel-mode tile --vae-use-tiling \
  --quantization fp8 --enforce-eager \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 Wait for `Application startup complete`, then run a short smoke test before
 committing to a full 50-step request.
 `t2va` requires an explicit `aspect_ratio` even when `width` and `height` are
 supplied; omitting it fails the request with
 `OmniClientError: t2va requires an explicit aspect_ratio`.
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=10' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":4.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-smoke.mp4
 ```
 The full-quality request uses 50 steps:
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'aspect_ratio=16:9' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"t2va","duration":8.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-t2va.mp4
 ffprobe -v error -show_entries \
  stream=index,codec_name,width,height,r_frame_rate,sample_rate,channels \
  -of json h3-gb10-t2va.mp4
 ```
 Both requests return `200 OK` with an MP4 body: H.264 video at the requested
 geometry plus a 32 kHz stereo AAC track.
 ## One GB10: Ref2VA, 960x576, 8 seconds
 `FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
 Stop the FL2VA server before starting this one.
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
 FLASHINFER_DISABLE_VERSION_CHECK=1 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
 vllm serve /path/to/MiniMax-H3/Ref2VA \
  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
  --init-timeout 3600 \
  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
  --vae-parallel-mode tile --vae-use-tiling \
  --quantization fp8 --enforce-eager \
  --diffusion-attention-backend CUDNN_ATTN
 ```
 The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
 takes over an hour on this platform. Use 14400 s.
 Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
 omitted when `width` and `height` are supplied.
 ```bash
 curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
  -D ref2va-headers.txt \
  -w 'client_e2e_s=%{time_total}\n' \
  --max-time 14400 \
  -F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
  -F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
  -F 'width=960' \
  -F 'height=576' \
  -F 'fps=24' \
  -F 'num_inference_steps=50' \
  -F 'flow_shift=12' \
  -F 'seed=1101' \
  -F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
  -o h3-gb10-ref2va.mp4
 ```
 The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
 32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
 ## Measured latency
 Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
 | Workload | Steps | Text encode | Denoise | Per step | VAE decode | MP4 mux | Client E2E | Peak memory |
 | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
 | T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
 | T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
 | Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
 | Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
 How these were measured:
 - Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
  `X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
  `MiniMaxH3Pipeline.encode_prompt` (text encode),
  `MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
  (VAE decode).
 - MP4 mux is not a header field. It is the
  `Video response encoding (MP4 bytes)` line in the server log.
 - Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
  high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
  the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
  higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
  memory has no discrete VRAM to report and the field returns `[N/A]`. See the
  capacity section for a `/proc/meminfo` cross-check.
 - For T2VA, text encode and VAE decode were measured on a separate 10-step run at
  the same 960x576 / 8 s geometry. Both are independent of step count, so they
  transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
  same way, from the 10-step run at identical output geometry and duration.
 - Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
  (`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
  geometry and reference input rather than step count.
 - Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
  engine bookkeeping sit outside the profiled stages.
 Denoise dominates at 96% of end-to-end, so step count and geometry are the only
 knobs that materially change the wall clock on this platform.
 Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
 4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
 versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
 The cost is in denoise. A single reference image expands the Qwen3-VL
 presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
 6883 tokens, 1 reference images`), and those tokens sit in the attention context
 of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
 1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
 expected for an identical output shape.
 Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
 (397 s to 490 s) on a machine that had already been under sustained load. Treat
 these as single-run measurements on one machine, not a throughput guarantee, and
 expect the slower end of that range for back-to-back jobs.
 Ref2VA per-step time is load-dependent in the same way. On an idle machine the
 first 10-step request ran at 66.55 s per step. The 50-step request that followed
 occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
 further 10-step requests taken immediately afterwards settled at 74.88 s and
 74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
 is a stable thermal state rather than measurement noise.
 Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
 only reachable on a cold machine and does not hold for the length of one request.
 Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
 8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
 VAE decode was stable at 67.9-69.8 s throughout.
 Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
 default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
 shorter than a 50-step Ref2VA run.
 ## Known limitations
 - One GPU means no Ulysses, ring, TP, or VAE patch parallelism. Pass
  `--usp 1 --ring 1 --vae-patch-parallel-size 1` explicitly if a profile or
  script would otherwise supply higher degrees.
 - `--enforce-eager` is used because regional compile adds startup time without a
  parallelism win on a single rank.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
  `height` are both supplied they take precedence and `aspect_ratio` only has to
  name one of the supported ratios.
 - Sustained load degrades throughput. See the variance notes above.
 - Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
  with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
  above; reference videos and additional images push the token count higher and
  were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
  single-image request, there is very little room to grow — re-measure before
  adding references or raising the output shape.
 - Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
  under **Validated on**: FP8 weight loading and image-only Ref2VA are both
  broken on `release/v0.26.0`.suggestion is  `v0.26.1`.
 - Online FP8 is incompatible with layerwise offload — the offload path produces a
  weight stride the Cutlass FP8 kernel rejects. This is not a practical
  restriction here since offload is unusable on GB10 anyway.
 ```
@@ -58,10 +58,6 @@ FlashAttention-4 extra:
 uv pip install -e '.[fa4]'
 ```
 On AMD ROCm, install without the `[fa4]` extra (FA4 is CUDA-only) and use
 `--diffusion-attention-backend FLASH_ATTN`; see
 [AMD ROCm (gfx942 / gfx950)](#amd-rocm-gfx942--gfx950).
 `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
 reference-video preparation and MP4 output.
@@ -391,101 +387,6 @@ incompatible with H3 layerwise offload because the offload path produces a
 weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
 parallelism and VAE tiling instead.
 ## AMD ROCm (gfx942 / gfx950)
 The sections above describe the NVIDIA CUDA path. This section covers the AMD
 ROCm path; use it instead of the CUDA commands when running on AMD Instinct
 GPUs.
 MiniMax H3 runs on AMD Instinct GPUs (gfx942 / gfx950) in BF16. Use
 `--diffusion-attention-backend FLASH_ATTN`, which resolves to AITER packed varlen
 attention on both architectures.
 Select the device with `HIP_VISIBLE_DEVICES` (not `CUDA_VISIBLE_DEVICES`), drop the
 CUDA-only `FLASHINFER_DISABLE_VERSION_CHECK`, and install without the `[fa4]` extra
 (FA4 is CUDA-only). The VAE uses AITER GroupNorm on ROCm.
 Install (ROCm wheel + source vLLM-Omni):
 ```bash
 pip install "vllm==0.26.0+rocm723" \
  --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
 VLLM_OMNI_TARGET_DEVICE=rocm pip install -e . --no-build-isolation
 ```
 Prebuilt image: `vllm/vllm-omni-rocm:minimax-h3`. All tasks work out of the box:
 the image bundles TorchCodec (for image+audio Ref2VA) and `ffmpeg` (for
 video-reference Ref2VA).
 ### ROCm single GPU
 Single GPU with model-level CPU offload keeps the Qwen3-VL encoder and DiT from
 being co-resident:
 ```bash
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8091
 HIP_VISIBLE_DEVICES=0 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni --host 0.0.0.0 --port "${PORT}" --trust-remote-code \
  --num-gpus 1 --enable-cpu-offload \
  --diffusion-attention-backend FLASH_ATTN
 ```
 ### ROCm four GPUs
 The best-practice CUDA four-GPU configuration works on ROCm with the changes above.
 It mirrors the CUDA command including Ulysses sequence parallelism, VAE patch
 parallelism, and text-encoder tensor parallelism, with no CPU offload when the model
 shards across the GPUs:
 ```bash
 export MODEL="${MODEL_ROOT}/FL2VA"
 export PORT=8091
 HIP_VISIBLE_DEVICES=0,1,2,3 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
 VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
 vllm serve "${MODEL}" \
  --omni \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --trust-remote-code \
  --num-gpus 4 \
  --usp 4 \
  --ring 1 \
  --text-encoder-tp-size 4 \
  --vae-patch-parallel-size 4 \
  --vae-parallel-mode tile \
  --vae-use-tiling \
  --diffusion-attention-backend FLASH_ATTN
 ```
 As on CUDA, H3 is CFG-distilled, so keep `--cfg-parallel-size` at 1, and the H3 VAE
 supports only its native `tile` mode. `--text-encoder-tp-size` is validated on
 gfx942; on gfx950 it was not exercised.
 ### Validated ROCm evidence
 vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
 `vllm/vllm-omni-rocm:minimax-h3` image; gfx950 rows measured with the
 `0.26.0+rocm723` wheel (HIP 7.2).
 | Workload | Configuration | Observed result |
 |----------|---------------|-----------------|
 | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
 gfx942 figures are the mean of three requests after one excluded warmup
 (external evidence: vllm-project/recipes#732). gfx950 figures are
 functional-correctness validations, not tuned throughput; the first request
 includes lazy regional compilation. MI325X (gfx942) and other MI355X SKUs are not
 listed until their own evidence is added.
 ## HTTP API examples
 The following requests use the synchronous endpoint so the returned body can
@@ -673,43 +574,10 @@ default. `short_edge` controls the 768-pixel canvas and must be `768`.
 outputs; the synchronous raw-MP4 endpoint returns the first output when more
 than one is requested.
 ## Request-scoped quality
 Add one of these fields to any HTTP request above. No Cache-DiT startup option
 is required; H3 installs its conservative profile when a `high` request
 arrives and removes it for a `lossless` request.
 ```bash
 -F 'quality=lossless'  # Native reference path for this request.
 -F 'quality=high'      # H3's conservative Cache-DiT profile.
 ```
 Omitting `quality` preserves the startup default: it uses the reference path
 normally, or the server-configured profile when the server was started with
 `--cache-backend cache_dit`.
 The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
 TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
 `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
 balanced switch order.
 | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
 |---|---:|---:|---:|---:|---|
 | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
 | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
 > [!NOTE]
 > `high` selects a fixed Cache-DiT profile, but its cache hit rate and
 > resulting latency/quality trade-off may vary by hardware, topology, and
 > workload. The values above apply to this deployment and are not universal
 > guarantees. `lossless` remains the exact reference path.
 ## Key parameters
 | Parameter | Recommended value | Notes |
 |-----------|-------------------|-------|
 | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
 | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
 | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
 | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
 | `fps` | `24` | H3 output FPS is fixed |
@@ -750,98 +618,19 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
 resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
 reduction.
 ## TeaCache acceleration
 TeaCache reuses DiT block residuals across denoising steps when consecutive
 timestep embeddings are similar. MiniMax-H3 TeaCache is currently calibrated
 only for the FL2VA partition. In combined serving, FL2VA requests use TeaCache
 while Ref2VA requests run uncached; Ref2VA-only serving rejects TeaCache.
 TeaCache and Cache-DiT are mutually exclusive; pick one cache backend per server.
 The model-specific default and examples use `rel_l1_thresh=0.17`, which
 provided the best conservative speed/quality balance in the validated 107-frame
 T2VA workload. Lower values
 may produce few or no cache hits, while higher values can improve performance
 at the cost of output quality. Validate the threshold on representative
 prompts and generation settings before changing it.
 ### Offline (Python API)
 Export the directory containing the `FL2VA` and `Ref2VA` subdirectories before
 running the Python example, for example `export MODEL_ROOT=/models/MiniMax-H3`.
 ```python
 import os
 from vllm_omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 omni = Omni(
    model=os.path.join(os.environ["MODEL_ROOT"], "FL2VA"),
    cache_backend="tea_cache",
    cache_config={"rel_l1_thresh": 0.17},
    trust_remote_code=True,
    enable_cpu_offload=True,
 )
 outputs = omni.generate(
    "A quiet cinematic night scene with matching ambient sound.",
    OmniDiffusionSamplingParams(
        height=256,
        width=448,
        num_frames=29,
        fps=24,
        num_inference_steps=50,
        seed=42,
        extra_args={
            "task": "t2va",
            "duration": 4.0,
            "aspect_ratio": "16:9",
            "flow_shift": 12.0,
            "audio_flow_shift": 3.0,
        },
    ),
 )
 ```
 ### Online serving
 ```bash
 vllm serve "${MODEL_ROOT}/FL2VA" \
  --omni \
  --trust-remote-code \
  --cache-backend tea_cache \
  --cache-config '{"rel_l1_thresh":0.17}'
 ```
 ## Known limitations
 - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
 - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
  both task-specific DiTs, and loads shared components once from `FL2VA`.
 - H3 currently executes one generation request per diffusion batch.
 - The first regional-compile request is a warmup and should not be included in
  steady-state performance measurements.
 - The serving path accepts fewer references than the model supports. H3 documents up
  to 9 images, 3 video clips, and 3 audio clips (12 files) per Omni Reference
  request; the current vLLM-Omni path takes exactly one image plus one audio
  reference, or one or more videos with no separate `audio_reference` (it uses the
  source soundtracks).
 - The 768 px short-edge mode is available for T2VA and FL2VA; 1344x768 is the
  documented 16:9 request shape.
 - `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch).
 - VAE patch parallelism requires size 1 or the full DiT group size and supports the
  H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
  pure Ulysses.
 - Online FP8 is not compatible with layerwise offload.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone
  audio reference.
- Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
+- VAE patch parallelism requires size 1 or the full DiT group size and supports
-  cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
+  the H3 native `tile` mode only.
  or model-level CPU offload; text-encoder TP alone is not sufficient.
 ## Additional resources
@@ -19,11 +19,6 @@ add files for this repo. Use one Markdown file per model family by default.
 > Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
 > e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
 The `Recipe` column in the Supported Models documentation links to the
 published recipe when available and otherwise to the corresponding recipe in
 this directory. The recipe files remain the source of truth for commands and
 validation details.
 Example layout:
 ```text
@@ -1805,7 +1805,6 @@ class TestQwen2_5OmniPipeline:
        assert s.owns_tokenizer is True
        assert s.engine_output_type == "latent"
        assert s.requires_multimodal_data is True
        assert s.hf_config_name == "thinker_config"
    def test_talker(self):
        p = resolve_pipeline_config("qwen2_5_omni")
@@ -1814,7 +1813,6 @@ class TestQwen2_5OmniPipeline:
        s = p.get_stage(1)
        assert isinstance(s, StagePipelineConfig)
        assert s.input_sources == (0,)
        assert s.hf_config_name == "talker_config"
        assert s.sampling_constraints["stop_token_ids"] == [8294]
        # thinker2talker was removed: qwen2_5_omni has no async_chunk support,
        # so sync_process_input_func always wins and custom_process_input_func
@@ -1831,7 +1829,6 @@ class TestQwen2_5OmniPipeline:
        assert s.execution_type == StageExecutionType.LLM_GENERATION
        assert s.final_output_type == "audio"
        assert s.engine_output_type == "audio"
        assert s.hf_config_name == "thinker_config"
 class TestQwen3TTSPipeline:
@@ -39,7 +39,6 @@ from vllm_omni.config.stage_config import (
    load_deploy_config,
    merge_pipeline_deploy,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -967,24 +966,20 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
        diffusion_attention_backend="flash_attn",
        fa_deterministic=True,
        kv_cache_dtype="fp8",
        kv_cache_skip_steps="0-1",
        kv_cache_skip_layers=[2],
        static_lora_scale=0.25,
        diffusion_kv_mode="paged_scheduler",
        diffusers_load_kwargs=None,
        diffusers_call_kwargs=None,
    )
    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
    assert cfg.fa_deterministic is True
    assert cfg.diffusion_kv_cache_dtype == "fp8"
    assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
    assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
    assert cfg.lora_scale == 0.25
    assert cfg.cache_backend == "tea_cache"
    assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
    assert cfg.diffusers_load_kwargs == {}
    assert cfg.diffusers_call_kwargs == {}
@@ -999,8 +994,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                "stages:",
                "  - stage_id: 0",
                "    diffusion_attention_backend: flash_attn",
                "    fa_deterministic: true",
                "    diffusion_kv_mode: paged_scheduler",
            ]
        )
    )
@@ -1012,29 +1005,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
    assert isinstance(stage, VllmOmniDiffusionStageConfig)
    assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
    assert stage.diffusion_config.fa_deterministic is True
    assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
    deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
    deploy_path.write_text(
        "\n".join(
            [
                "pipeline: dreamzero",
                "async_chunk: false",
                "stages:",
                "  - stage_id: 0",
                "    diffusion_kv_mode: paged_worker_local",
            ]
        )
    )
    with pytest.raises(ValueError, match="reserved but not implemented"):
        _from_pipeline_key(
            "dreamzero",
            deploy_config_path=str(deploy_path),
        )
 def test_diffusion_config_field_classification_covers_current_fields():
@@ -1055,7 +1025,6 @@ def test_diffusion_config_field_classification_covers_current_fields():
        "enable_session_state_manager",
        "prompt_embed_cache_size",
        "diffusion_kv_cache_dtype",
        "diffusion_kv_mode",
    } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
    assert {
        "revision",
@@ -1,192 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Equivalence and session-isolation tests for Cosmos3StateAdapter (RFC #4480).
 These drive the adapter directly with tiny CPU tensors -- no model, no GPU --
 and assert it stores/returns the per-branch UND K/V exactly and isolates
 durable state between sessions. They do not cover the transformer's mutable
 load-to-forward handoff.
 """
 from __future__ import annotations
 import pytest
 import torch
 from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import (
    Cosmos3StateAdapter,
 )
 from vllm_omni.experimental.world_models.session_state import SessionStateManager
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 LAYERS = 3
 SEQ, DIM = 3, 4
 def _fake_cached_kv(seed: int) -> list[tuple[torch.Tensor, torch.Tensor]]:
    """A list of per-layer (K, V); contents seeded so branches/sessions differ."""
    g = torch.Generator().manual_seed(seed)
    return [(torch.randn(1, SEQ, DIM, generator=g), torch.randn(1, SEQ, DIM, generator=g)) for _ in range(LAYERS)]
 def _adapter(session_id: str, manager: SessionStateManager | None = None) -> Cosmos3StateAdapter:
    # `manager is not None`, not `manager or ...`: an empty SessionStateManager is
    # falsy (it defines __len__), so `or` would silently swap in a fresh manager.
    return Cosmos3StateAdapter(session_id, manager if manager is not None else SessionStateManager())
 def _assert_kv_equal(got: list | None, expected: list) -> None:
    assert got is not None and len(got) == len(expected) == LAYERS
    for (k, v), (ke, ve) in zip(got, expected):
        torch.testing.assert_close(k, ke)
        torch.testing.assert_close(v, ve)
 class _FakeTransformer:
    cached_kv = None
    cached_freqs_gen = None
 # ----------------------------- equivalence -----------------------------
 def test_branch_starts_uninitialized() -> None:
    a = _adapter("s0")
    for is_neg in (False, True):
        assert a.is_branch_initialized(is_neg) is False
        assert a.get_branch_kv(is_neg) is None
 def test_set_then_get_roundtrip() -> None:
    a = _adapter("s0")
    kv = _fake_cached_kv(1)
    a.set_branch_kv(False, kv)
    assert a.is_branch_initialized(False) is True
    _assert_kv_equal(a.get_branch_kv(False), kv)
 def test_cond_uncond_branches_isolated() -> None:
    a = _adapter("s0")
    kv_pos, kv_neg = _fake_cached_kv(1), _fake_cached_kv(2)
    a.set_branch_kv(False, kv_pos)
    a.set_branch_kv(True, kv_neg)
    _assert_kv_equal(a.get_branch_kv(False), kv_pos)
    _assert_kv_equal(a.get_branch_kv(True), kv_neg)
 def test_encode_once_capture_is_noop_when_initialized() -> None:
    a = _adapter("s0")
    kv = _fake_cached_kv(1)
    a.set_branch_kv(False, kv)
    tr = _FakeTransformer()
    tr.cached_kv = _fake_cached_kv(999)  # different content
    a.capture_from_transformer(tr, False)
    _assert_kv_equal(a.get_branch_kv(False), kv)  # still the first one
 def test_capture_writes_when_uninitialized() -> None:
    a = _adapter("s0")
    tr = _FakeTransformer()
    tr.cached_kv = _fake_cached_kv(5)
    tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2))
    a.capture_from_transformer(tr, False)
    _assert_kv_equal(a.get_branch_kv(False), _fake_cached_kv(5))
 def test_capture_without_freqs_does_not_partially_initialize() -> None:
    a = _adapter("s0")
    tr = _FakeTransformer()
    tr.cached_kv = _fake_cached_kv(5)
    tr.cached_freqs_gen = None
    with pytest.raises(RuntimeError, match="did not produce GEN RoPE frequencies"):
        a.capture_from_transformer(tr, False)
    assert a.is_branch_initialized(False) is False
 def test_load_into_transformer_without_freqs_sets_kv() -> None:
    a = _adapter("s0")
    kv = _fake_cached_kv(1)
    a.set_branch_kv(False, kv)
    tr = _FakeTransformer()
    a.load_into_transformer(tr, False)
    _assert_kv_equal(tr.cached_kv, kv)
    assert tr.cached_freqs_gen is None
 def test_set_then_get_preserves_per_layer_values() -> None:
    # The stored value must reflect the tensors passed in (per-layer).
    a = _adapter("s0")
    kv = _fake_cached_kv(3)
    a.set_branch_kv(False, kv)
    _assert_kv_equal(a.get_branch_kv(False), kv)
 def test_load_into_transformer_restores_kv_and_freqs_gen() -> None:
    a = _adapter("s0")
    tr = _FakeTransformer()
    tr.cached_kv = _fake_cached_kv(3)
    tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2))
    a.capture_from_transformer(tr, False)
    tr.cached_kv = None
    tr.cached_freqs_gen = None
    a.load_into_transformer(tr, False)
    _assert_kv_equal(tr.cached_kv, _fake_cached_kv(3))
    assert tr.cached_freqs_gen is not None
    torch.testing.assert_close(tr.cached_freqs_gen[0], torch.ones(2))
 def test_partial_branch_raises() -> None:
    a = _adapter("s0")
    a.set_branch_kv(False, _fake_cached_kv(1))
    a._session.get("und_kv/1").reset()  # type: ignore[union-attr]
    with pytest.raises(RuntimeError, match="partially initialized"):
        a.get_branch_kv(False)
 # ----------------------------- reset -----------------------------
 def test_reset_clears_both_branches() -> None:
    a = _adapter("s0")
    a.set_branch_kv(False, _fake_cached_kv(1))
    a.set_branch_kv(True, _fake_cached_kv(2))
    a.reset()
    for is_neg in (False, True):
        assert a.is_branch_initialized(is_neg) is False
        assert a.get_branch_kv(is_neg) is None
 # ----------------------- concurrency isolation (core) -----------------------
 def test_two_sessions_do_not_clobber() -> None:
    manager = SessionStateManager()
    a, b = _adapter("A", manager), _adapter("B", manager)
    kv_a, kv_b = _fake_cached_kv(10), _fake_cached_kv(20)
    a.set_branch_kv(False, kv_a)
    assert b.is_branch_initialized(False) is False  # B not seeing A's state
    b.set_branch_kv(False, kv_b)
    _assert_kv_equal(a.get_branch_kv(False), kv_a)
    _assert_kv_equal(b.get_branch_kv(False), kv_b)
 def test_fresh_adapter_same_session_sees_prior_state() -> None:
    manager = SessionStateManager()
    kv = _fake_cached_kv(7)
    _adapter("s", manager).set_branch_kv(False, kv)
    again = Cosmos3StateAdapter("s", manager)
    _assert_kv_equal(again.get_branch_kv(False), kv)
 def test_adapter_survives_session_lru_eviction() -> None:
    # An in-use adapter keeps its state even after the manager evicts its
    # session id from the lookup table (LRU bounded by max_sessions).
    manager = SessionStateManager(max_sessions=2)
    a = _adapter("active", manager)
    kv = _fake_cached_kv(11)
    a.set_branch_kv(False, kv)
    for i in range(3):
        _adapter(f"other{i}", manager)
    assert "active" not in manager
    _assert_kv_equal(a.get_branch_kv(False), kv)
@@ -43,9 +43,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.0361,
+            "throughput_qps": 0.0434,
-            "latency_mean": 27.6206,
+            "latency_mean": 23.0654,
-            "peak_memory_mb_mean": 80548.0
+            "peak_memory_mb_mean": 76360.0
          }
        }
      }
@@ -100,9 +100,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.05263,
+            "throughput_qps": 0.0606,
-            "latency_mean": 18.9415,
+            "latency_mean": 16.4928,
-            "peak_memory_mb_mean": 50053.7
+            "peak_memory_mb_mean": 47164.8
          }
        }
      },
@@ -127,9 +127,9 @@
        ],
        "baseline": {
          "H100": {
-            "throughput_qps": 0.009595,
+            "throughput_qps": 0.0124,
-            "latency_mean": 104.1058,
+            "latency_mean": 80.7784,
-            "peak_memory_mb_mean": 59718.0
+            "peak_memory_mb_mean": 55374.3429
          }
        }
      }
@@ -31,7 +31,6 @@ import tempfile
 import pytest
 import torch
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.allgather_kv import (
@@ -410,7 +409,6 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
    torch.testing.assert_close(v_full, expected_value)
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize(
    "test_model_cls",
    [
@@ -418,12 +416,10 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
    ],
 )
@pytest.mark.parametrize(
-    ("ulysses_degree", "ring_degree", "allgather_degree", "num_kv_heads"),
+    ("ulysses_degree", "ring_degree", "allgather_degree"),
    [
-        pytest.param(2, 2, 1, None, id="ulysses-ring"),
+        pytest.param(2, 2, 1, id="ulysses-ring"),
-        pytest.param(1, 1, 2, None, id="allgather-kv"),
+        pytest.param(1, 1, 2, id="allgather-kv"),
        pytest.param(1, 2, 1, 2, id="ring-gqa"),
        pytest.param(1, 2, 1, 1, id="ring-mqa"),
    ],
 )
@pytest.mark.parametrize("batch_size", [2])
@@ -449,7 +445,6 @@ def test_sequence_parallel(
    seq_len: int,
    num_heads: int,
    head_size: int,
    num_kv_heads: int | None,
 ):
    """Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
    sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
@@ -481,7 +476,6 @@ def test_sequence_parallel(
                seq_len,
                num_heads,
                head_size,
                num_kv_heads,
                dtype,
                causal,
                use_sync,
@@ -513,7 +507,6 @@ def test_sequence_parallel(
                seq_len,
                num_heads,
                head_size,
                num_kv_heads,
                dtype,
                causal,
                use_sync,
@@ -610,7 +603,6 @@ def ulysses_attention_on_test_model(
    seq_len: int,
    num_heads: int,
    head_size: int,
    num_kv_heads: int | None,
    dtype: torch.dtype,
    causal: bool,
    use_sync: bool,
@@ -666,13 +658,10 @@ def ulysses_attention_on_test_model(
        cfg_parallel_size=1,
    )
-    od_config = OmniDiffusionConfig.from_kwargs(
+    od_config = OmniDiffusionConfig(
        model="test_model",
        dtype=dtype,
        parallel_config=parallel_config,
        # This regression targets pytorch_attn_forward(). Do not let an
        # installed FA/FA3 backend silently bypass the SDPA Ring path.
        diffusion_attention_backend="TORCH_SDPA",
    )
    # Initialize model parallel
@@ -698,7 +687,7 @@ def ulysses_attention_on_test_model(
            "head_size": head_size,
            "hidden_size": hidden_size,
            "causal": causal,
-            "num_kv_heads": num_kv_heads,
+            "num_kv_heads": None,
            "scatter_idx": 2,
            "gather_idx": 1,
            "use_sync": use_sync,
@@ -6,14 +6,12 @@ from __future__ import annotations
 import os
 import socket
 import tempfile
 from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -41,33 +39,6 @@ def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
    os.environ["MASTER_PORT"] = str(master_port)
@pytest.mark.core_model
@pytest.mark.cpu
 def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
    monkeypatch.setattr(
        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
        lambda default: "advanced_uaa",
    )
    sp_group = SimpleNamespace(
        ulysses_group=None,
        ulysses_world_size=2,
        ulysses_rank=0,
        ring_world_size=2,
    )
    strategy = UlyssesParallelAttention(
        sp_group=sp_group,
        scatter_idx=2,
        gather_idx=1,
        use_sync=False,
    )
    query = torch.zeros(1, 2, 8, 4)
    key = torch.zeros(1, 2, 1, 4)
    value = torch.zeros_like(key)
    with pytest.raises(ValueError, match="does not support K/V head padding"):
        strategy.pre_attention(query, key, value, None)
 def _run_attention_case(
    local_rank: int,
    world_size: int,
@@ -270,9 +241,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
    batch_size = 2
    head_size = 8
    seq_len = 10
-    # This is the positive Hybrid case: K/V heads must not require padding,
+    num_heads = 3  # head_cnt not divisible by ulysses_degree=2 -> triggers head padding
    # which is unsupported by advanced_uaa when Ring is also enabled.
    num_heads = 4
    # Ensure ring ranks see equal post-Ulysses seq_len:
    # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
@@ -11,7 +11,6 @@ This module tests the cache backend implementations:
 - DiffusionCacheConfig: configuration dataclass
 """
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 import cache_dit
@@ -22,7 +21,6 @@ from vllm_omni.diffusion.cache.cachedit import (
    CacheDiTAdapterConfig,
    CacheDiTBackend,
    CacheDiTConfig,
    cache_summary,
 )
 from vllm_omni.diffusion.cache.magcache import MagCacheBackend
 from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -93,8 +91,8 @@ class TestCacheDiTBackend:
    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
-    def test_lifecycle_applies_to_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
+    def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
-        """Combined pipelines manage every exact Cache-DiT target."""
+        """Combined pipelines enable and refresh Cache-DiT for every DiT."""
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline._dit_modules = ["transformer", "transformers_ref"]
@@ -104,67 +102,18 @@ class TestCacheDiTBackend:
            transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
                block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
            )
        installed_adapters = [Mock(), Mock()]
        mock_block_adapter.side_effect = installed_adapters
        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
        backend.enable(pipeline)
        backend.refresh(pipeline, num_inference_steps=20)
        backend.disable(pipeline)
        assert len(backend._refresh_funcs) == 2
        assert mock_cache_dit.enable_cache.call_count == 2
        assert mock_cache_dit.refresh_context.call_count == 2
        assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
            pipeline.transformer,
            pipeline.transformers_ref,
        }
        assert [call.args[0] for call in mock_cache_dit.disable_cache.call_args_list] == installed_adapters
        assert backend._refresh_funcs == []
        assert backend._cache_targets == []
        assert not backend.is_enabled()
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_summary_skips_uncached_nested_dit(self, mock_cache_dit):
        """Only DiT modules with an active Cache-DiT context are summarized."""
        language_model = Mock()
        language_model._is_cached = False
        transformer = Mock()
        transformer._is_cached = True
        transformer.language_model = language_model
        pipeline = SimpleNamespace(
            _dit_modules=["transformer.language_model", "transformer"],
            transformer=transformer,
        )
        cache_summary(pipeline)
        mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_enable_and_refresh_nested_declared_dit(self, mock_cache_dit, mock_block_adapter):
        """Dotted component paths resolve to their nested DiT module."""
        transformer = Mock()
        transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
            block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
        )
        pipeline = SimpleNamespace(
            _dit_modules=["language_model.model"],
            language_model=SimpleNamespace(model=transformer),
        )
        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
        backend.enable(pipeline)
        backend.refresh(pipeline, num_inference_steps=20)
        cache_summary(pipeline)
        mock_cache_dit.enable_cache.assert_called_once()
        mock_cache_dit.refresh_context.assert_called_once_with(
            transformer,
            num_inference_steps=20,
            verbose=True,
        )
        mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
    @patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
@@ -297,32 +246,6 @@ class TestCacheDiTBackend:
        assert mock_cache_dit.refresh_context.call_count == 2
    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_refresh_rearms_force_refresh_hint(self, mock_cache_dit, mock_block_adapter):
        """A once hint must be restored for every repeated request."""
        pipeline = Mock()
        pipeline.__class__.__name__ = "DiTPipeline"
        pipeline.transformer = Mock()
        pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
            block_forward_patterns={"layers": ForwardPattern.Pattern_0}
        )
        config = DiffusionCacheConfig(
            force_refresh_step_hint=7,
            force_refresh_step_policy="once",
        )
        backend = CacheDiTBackend(config)
        backend.enable(pipeline)
        backend.refresh(pipeline, num_inference_steps=20)
        backend.refresh(pipeline, num_inference_steps=20)
        assert mock_cache_dit.refresh_context.call_count == 2
        for call in mock_cache_dit.refresh_context.call_args_list:
            refresh_config = call.kwargs["cache_config"]
            assert refresh_config.force_refresh_step_hint == 7
            assert refresh_config.force_refresh_step_policy == "once"
    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -471,66 +394,6 @@ class TestTeaCacheBackend:
        assert backend.enabled is True
        mock_apply_hook.assert_called_once()
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_enable_uses_generic_default_threshold(self, mock_apply_hook):
        pipeline = Mock()
        pipeline.__class__.__name__ = "QwenImagePipeline"
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "QwenImageTransformer2DModel"
        TeaCacheBackend(DiffusionCacheConfig()).enable(pipeline)
        assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.2
    @pytest.mark.parametrize("partition", ["fl2va", "combined"])
    @pytest.mark.parametrize(
        ("configured_threshold", "expected_threshold"),
        [(None, 0.17), (0.2, 0.2)],
    )
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_minimax_h3_only_enables_fl2va_teacache(
        self,
        mock_apply_hook,
        partition,
        configured_threshold,
        expected_threshold,
    ):
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline.partition = partition
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
        pipeline.transformers_ref = Mock()
        config = (
            DiffusionCacheConfig()
            if configured_threshold is None
            else DiffusionCacheConfig(rel_l1_thresh=configured_threshold)
        )
        backend = TeaCacheBackend(config)
        backend.enable(pipeline)
        mock_apply_hook.assert_called_once()
        transformer, teacache_config = mock_apply_hook.call_args.args
        assert transformer is pipeline.transformer
        assert transformer is not pipeline.transformers_ref
        assert teacache_config.rel_l1_thresh == expected_threshold
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_minimax_h3_rejects_ref2va_teacache(self, mock_apply_hook):
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline.partition = "ref2va"
        pipeline.transformer = Mock()
        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
        backend = TeaCacheBackend(DiffusionCacheConfig())
        with pytest.raises(ValueError, match="only supports the MiniMax-H3 FL2VA partition"):
            backend.enable(pipeline)
        assert backend.enabled is False
        mock_apply_hook.assert_not_called()
    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
    def test_enable_with_coefficients(self, mock_apply_hook):
        """Test enabling TeaCache with custom coefficients."""
@@ -81,8 +81,6 @@ def test_cachedit_public_api_is_explicit():
        "CacheDiTAdapterConfig",
        "CacheDiTBackend",
        "CacheDiTConfig",
        "CacheDiTRequestSpec",
        "RequestScopedCacheDiTRuntime",
        "SensenovaCachedAdapter",
        "cache_summary",
        "enable_cache_for_dit",
@@ -1,101 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
 from vllm_omni.diffusion.cache.cachedit.runtime import (
    CacheDiTRequestSpec,
    RequestScopedCacheDiTRuntime,
 )
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
    return CacheDiTRequestSpec(
        installation_key=key,
        cache_config=DiffusionCacheConfig(),
        num_inference_steps=steps,
    )
 def _recording_backend(events):
    class FakeBackend:
        def __init__(self, config):
            events.append(("create", config))
            self.enabled = False
        def enable(self, pipeline):
            events.append(("enable", pipeline))
            self.enabled = True
        def refresh(self, pipeline, num_inference_steps):
            events.append(("refresh", pipeline, num_inference_steps))
        def disable(self, pipeline):
            events.append(("disable", pipeline))
            self.enabled = False
        def is_enabled(self):
            return self.enabled
    return FakeBackend
 def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
    events = []
    FakeBackend = _recording_backend(events)
    monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
    pipeline = object()
    config = DiffusionCacheConfig()
    runtime = RequestScopedCacheDiTRuntime(pipeline)
    startup_backend = FakeBackend(config)
    startup_backend.enabled = True
    runtime.adopt(startup_backend, installation_key="generic")
    runtime.prepare(_spec("generic", 50))  # omitted: keep startup cache
    runtime.prepare(None)  # lossless: uninstall
    runtime.prepare(None)  # repeated lossless: no-op
    runtime.prepare(_spec("high", 40))  # high: install H3 profile
    runtime.prepare(_spec("high", 30))  # repeated high: refresh only
    runtime.prepare(_spec("generic", 20))  # omitted: restore startup profile
    assert [event[0] for event in events] == [
        "create",
        "refresh",
        "disable",
        "create",
        "enable",
        "refresh",
        "refresh",
        "disable",
        "create",
        "enable",
        "refresh",
    ]
    assert events[1][-1] == 50
    assert events[5][-1] == 40
    assert events[6][-1] == 30
    assert events[10][-1] == 20
    assert runtime.is_enabled
 def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
    events = []
    monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
    pipeline = object()
    runtime = RequestScopedCacheDiTRuntime(pipeline)
    runtime.prepare(_spec("high", 40))
    runtime.prepare(None)
    assert [event[0] for event in events] == [
        "create",
        "enable",
        "refresh",
        "disable",
    ]
    assert events[2][-1] == 40
    assert not runtime.is_enabled
@@ -15,23 +15,19 @@ Currently implemented:
 - TestFlux2KleinExtractor: Flux2Klein model extractor
 - TestFlux2Extractor: Flux2 model extractor
 - TestFluxExtractor: Flux model extractor
 - TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
 """
 from abc import ABC, abstractmethod
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 import torch.nn as nn
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import (
    extract_flux2_context,
    extract_flux2_klein_context,
    extract_flux_context,
    extract_minimax_h3_context,
 )
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
@@ -427,280 +423,3 @@ class TestFluxExtractor(BaseExtractorTest):
                img_ids=torch.randint(0, 64, (1, 16, 3)),
                txt_ids=torch.randint(0, 64, (1, 8, 3)),
            )
 class _MiniMaxH3FakeLinear(nn.Module):
    def __init__(
        self,
        in_features,
        out_features,
        *,
        bias=False,
        params_dtype=None,
        quant_config=None,
        prefix="",
        total_num_heads=None,
        total_num_kv_heads=None,
        **kwargs,
    ):
        del kwargs
        super().__init__()
        self.prefix = prefix
        self.quant_config = quant_config
        self.weight = nn.Parameter(torch.randn(out_features, in_features, dtype=params_dtype or torch.float32))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=params_dtype or torch.float32))
        else:
            self.register_parameter("bias", None)
        self.num_heads = total_num_heads
        self.num_kv_heads = total_num_kv_heads
    def forward(self, x):
        return torch.nn.functional.linear(x, self.weight, self.bias), None
 class _MiniMaxH3FakeMergedLinear(_MiniMaxH3FakeLinear):
    def __init__(self, in_features, out_features, **kwargs):
        super().__init__(in_features, sum(out_features), **kwargs)
 class _MiniMaxH3FakeQKVLinear(_MiniMaxH3FakeLinear):
    def __init__(
        self,
        *,
        hidden_size,
        head_size,
        total_num_heads,
        total_num_kv_heads,
        **kwargs,
    ):
        qkv_features = (total_num_heads + 2 * total_num_kv_heads) * head_size
        super().__init__(
            hidden_size,
            qkv_features,
            total_num_heads=total_num_heads,
            total_num_kv_heads=total_num_kv_heads,
            **kwargs,
        )
 class _MiniMaxH3FakeAttention(nn.Module):
    class FakeBackend:
        supports_prefix_kv_slicing = False
    def __init__(self, **kwargs):
        del kwargs
        super().__init__()
        self.attn_backend = self.FakeBackend
        self.use_ring = False
    def forward(self, q, k, v, metadata=None):
        del k, v, metadata
        return q
 def _minimax_h3_small_od_config():
    arch = {
        "num_layers": 1,
        "token_refiner_num_layers": 1,
        "hidden_size": 8,
        "num_attention_heads": 1,
        "attention_head_dim": 8,
        "ffn_hidden_size": 16,
        "latents_dim": 2,
        "audio_latents_dim": 2,
        "patch_size": (1, 2, 2),
        "text_dim": 6,
        "timestep_input_dim": 4,
        "time_embed_hidden_size": 8,
        "time_embed_dim": 4,
        "adaln_out_features": 18 * 8,
        "final_adaln_out_features": 2 * 8,
        "rope_inv_freq_len": 1,
    }
    return SimpleNamespace(
        tf_model_config=arch,
        parallel_config=SimpleNamespace(ulysses_degree=1),
    )
 def _minimax_h3_sample_kwargs(*, seq_len: int = 8, device: torch.device | str = "cpu") -> dict:
    video_row_dim = 2 * 1 * 2 * 2
    audio_row_dim = 2
    text_len = 2
    img_rows = 2
    img_pos = torch.tensor([2, 3], dtype=torch.long, device=device)
    audio_pos = torch.tensor([4, 5], dtype=torch.long, device=device)
    text_pos = torch.tensor([0, 1], dtype=torch.long, device=device)
    infer_out_pos = img_pos.clone()
    token_tags = torch.full((seq_len,), -1, dtype=torch.long, device=device)
    token_tags[text_pos] = 1
    token_tags[img_pos] = 0
    token_tags[audio_pos] = 2
    inverse_indices = torch.zeros(seq_len, dtype=torch.long, device=device)
    return {
        "x": torch.randn(1, seq_len, video_row_dim, device=device),
        "audio_x": torch.randn(1, seq_len, audio_row_dim, device=device),
        "img_position_ids": torch.zeros(1, seq_len, 3, dtype=torch.float64, device=device),
        "unique_timesteps": torch.tensor([0.5], dtype=torch.float32, device=device),
        "inverse_indices": inverse_indices,
        "update_mask": torch.ones(img_rows, dtype=torch.bool, device=device),
        "token_tags": token_tags,
        "prompt_embeds": torch.randn(text_len, 6, device=device),
        "img_pos_info": {"position_ids": img_pos},
        "audio_pos_info": {"position_ids": audio_pos},
        "text_pos_info": {"position_ids": text_pos},
        "img_pos_for_infer_output_info": {"position_ids": infer_out_pos},
        "packed_seq_params": {
            "cu_seqlens_q": torch.tensor([0, seq_len], dtype=torch.int32, device=device),
            "max_seqlen_q": seq_len,
        },
        "refiner_packed_seq_params": {
            "cu_seqlens_q": torch.tensor([0, text_len], dtype=torch.int32, device=device),
            "max_seqlen_q": text_len,
        },
    }
@pytest.mark.cpu
 class TestMiniMaxH3Extractor(BaseExtractorTest):
    """Test extract_minimax_h3_context function."""
    def get_extractor(self):
        return extract_minimax_h3_context
    @pytest.fixture
    def minimax_h3_module(self, monkeypatch):
        """Create a minimal MiniMaxH3DiTModel with fake parallel linears for CPU tests."""
        from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
        monkeypatch.setattr(h3, "ColumnParallelLinear", _MiniMaxH3FakeLinear)
        monkeypatch.setattr(h3, "MergedColumnParallelLinear", _MiniMaxH3FakeMergedLinear)
        monkeypatch.setattr(h3, "QKVParallelLinear", _MiniMaxH3FakeQKVLinear)
        monkeypatch.setattr(h3, "RowParallelLinear", _MiniMaxH3FakeLinear)
        monkeypatch.setattr(h3, "Attention", _MiniMaxH3FakeAttention)
        monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
        model = h3.MiniMaxH3DiTModel(_minimax_h3_small_od_config(), quant_config=None)
        for submodule in model.modules():
            if isinstance(submodule, h3.MiniMaxH3Attention):
                submodule.rope._forward_method = submodule.rope.forward_native
        model.eval()
        return model
    def get_module(self, minimax_h3_module):
        return minimax_h3_module
    @pytest.fixture
    def sample_inputs(self):
        """Create sample packed forward kwargs for MiniMaxH3DiTModel."""
        return _minimax_h3_sample_kwargs(seq_len=8)
    def get_sample_inputs(self, sample_inputs):
        return sample_inputs
    def test_modulated_input_shape(self, minimax_h3_module, sample_inputs):
        """Test that modulated_input matches packed sequence shape [T, H]."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        seq_len = sample_inputs["x"].shape[1]
        assert context.modulated_input.shape == (seq_len, minimax_h3_module.hidden_size)
    def test_run_transformer_blocks_callable(self, minimax_h3_module, sample_inputs):
        """Test that run_transformer_blocks is callable."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        assert callable(context.run_transformer_blocks)
    def test_postprocess_callable(self, minimax_h3_module, sample_inputs):
        """Test that postprocess is callable."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        assert callable(context.postprocess)
    def test_postprocess_output_shape(self, minimax_h3_module, sample_inputs):
        """Test that postprocess returns video/audio logits for infer rows."""
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        (hidden,) = context.run_transformer_blocks()
        video_logits, audio_logits = context.postprocess(hidden)
        num_video_rows = sample_inputs["img_pos_info"]["position_ids"].shape[0]
        num_audio_rows = sample_inputs["audio_pos_info"]["position_ids"].shape[0]
        assert video_logits.shape[0] == num_video_rows
        assert audio_logits.shape[0] == num_audio_rows
    def test_extractor_matches_native_forward(self, minimax_h3_module, sample_inputs):
        """The uncached extractor path must preserve the model's forward result."""
        expected_video, expected_audio = minimax_h3_module(**sample_inputs)
        context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
        (hidden,) = context.run_transformer_blocks()
        actual_video, actual_audio = context.postprocess(hidden)
        torch.testing.assert_close(actual_video, expected_video)
        torch.testing.assert_close(actual_audio, expected_audio)
    def test_run_transformer_blocks_forwards_packed_layout(
        self,
        minimax_h3_module,
        sample_inputs,
        monkeypatch,
    ):
        """The extractor block call must mirror the native packed contract."""
        from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
        video_layout = VideoTokenLayout(prefix_len=2, latent_grid=(1, 1, 2))
        inputs = {**sample_inputs, "video_token_layout": video_layout}
        block = minimax_h3_module.blocks[0]
        original_forward = block.forward
        received: list[tuple[int, VideoTokenLayout | None]] = []
        def capture_forward(*args, **kwargs):
            received.append((kwargs["packed_total"], kwargs["video_layout"]))
            return original_forward(*args, **kwargs)
        monkeypatch.setattr(block, "forward", capture_forward)
        context = extract_minimax_h3_context(minimax_h3_module, **inputs)
        context.run_transformer_blocks()
        assert received == [(sample_inputs["x"].shape[1], video_layout)]
    def test_teacache_reuses_h3_block_residual(self, minimax_h3_module, sample_inputs, monkeypatch):
        """A guaranteed cache hit must skip the H3 transformer block."""
        from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
        from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
        block = minimax_h3_module.blocks[0]
        original_forward = block.forward
        block_calls = 0
        def counted_forward(*args, **kwargs):
            nonlocal block_calls
            block_calls += 1
            return original_forward(*args, **kwargs)
        monkeypatch.setattr(block, "forward", counted_forward)
        hook = TeaCacheHook(
            TeaCacheConfig(
                transformer_type="MiniMaxH3DiTModel",
                rel_l1_thresh=0.3,
            )
        )
        hook.initialize_hook(minimax_h3_module)
        hook.new_forward(minimax_h3_module, **sample_inputs)
        video_logits, audio_logits = hook.new_forward(minimax_h3_module, **sample_inputs)
        assert block_calls == 1
        assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0]
        assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0]
    def test_invalid_module_raises_error(self):
        """Test that invalid module without blocks raises ValueError."""
        invalid_module = Mock()
        invalid_module.blocks = []
        with pytest.raises(ValueError, match="Module must have blocks"):
            extract_minimax_h3_context(invalid_module, **_minimax_h3_sample_kwargs())
@@ -1,52 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
 def _fixed_master_port(monkeypatch) -> None:
    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
 def test_dense_legacy_is_default() -> None:
    config = OmniDiffusionConfig.from_kwargs()
    assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
 def test_paged_worker_local_is_rejected_until_implemented() -> None:
    with pytest.raises(ValueError, match="reserved but not implemented"):
        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")
 def test_unknown_cache_mode_is_rejected() -> None:
    with pytest.raises(ValueError, match="Unsupported Diffusion KV diffusion_kv_mode"):
        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="unknown")
 def test_non_mapping_omni_kv_config_is_rejected() -> None:
    with pytest.raises(TypeError, match="omni_kv_config must be a mapping"):
        OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
 def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
    with pytest.raises(ValueError, match="does not support imported AR KV"):
        OmniDiffusionConfig.from_kwargs(
            diffusion_kv_mode="paged_scheduler",
            omni_kv_config={"need_recv_cache": True},
        )
 def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
    config = OmniDiffusionConfig.from_kwargs(
        model_class_name="FutureDiffusionModel",
        diffusion_kv_mode="paged_scheduler",
    )
    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
@@ -1,200 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from vllm.v1.core.kv_cache_manager import KVCacheManager
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 BLOCK_SIZE = 4
 def _request(**overrides) -> DiffusionKVRequest:
    values = dict(
        request_id="req/diffusion-kv/0",
        sequence_id=0,
        prefix_len=4,
        target_len=4,
        seq_len=8,
    )
    values.update(overrides)
    return DiffusionKVRequest(**values)
 def test_request_exposes_native_request_and_diffusion_semantics() -> None:
    request = _request()
    assert request.sequence_id == 0
    assert request.prefix_len == 4
    assert request.target_len == 4
    assert request.num_tokens == request.seq_len == 8
    assert request.num_prompt_tokens == 4
    assert request.num_computed_tokens == 0
    assert request.block_hashes == []
    assert request.kv_contexts == ()
    assert request.skip_reading_prefix_cache is True
    assert request.shared_prefix_boundary == 0
    assert request.status is RequestStatus.WAITING
@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"sequence_id": -1}, "sequence_id"),
        ({"prefix_len": -1}, "prefix_len"),
        ({"target_len": 0}, "target_len"),
        ({"seq_len": 0}, "seq_len"),
        ({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"),
    ],
 )
 def test_request_rejects_invalid_lengths(overrides, message: str) -> None:
    with pytest.raises(ValueError, match=message):
        _request(**overrides)
@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"context_id": ""}, "context_id"),
        ({"cache_role": ""}, "cache_role"),
        ({"num_tokens": 0}, "num_tokens"),
    ],
 )
 def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None:
    values = dict(context_id="text", cache_role="cross.text", num_tokens=8)
    values.update(overrides)
    with pytest.raises(ValueError, match=message):
        DiffusionKVContext(**values)
 def test_request_owns_independent_kv_contexts() -> None:
    text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
    image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4)
    request = _request(kv_contexts=(text, image))
    assert request.kv_contexts == (text, image)
 def test_request_rejects_duplicate_context_ids() -> None:
    contexts = (
        DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8),
        DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4),
    )
    with pytest.raises(ValueError, match="unique context_id"):
        _request(kv_contexts=contexts)
 def test_request_rejects_non_context_values() -> None:
    with pytest.raises(TypeError, match="DiffusionKVContext"):
        _request(kv_contexts=(object(),))
 def _manager(*, enable_caching: bool) -> KVCacheManager:
    # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
    # already registered an out-of-tree spec, making the registry non-empty
    # before FullAttentionSpec is registered. Initialize the native mappings
    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
    if register_all is not None:
        register_all(None)
    spec = FullAttentionSpec(
        block_size=BLOCK_SIZE,
        num_kv_heads=2,
        head_size=8,
        dtype=torch.bfloat16,
    )
    group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
    config = KVCacheConfig(
        num_blocks=8,
        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
        kv_cache_groups=[group],
    )
    return KVCacheManager(
        config,
        max_model_len=64,
        scheduler_block_size=BLOCK_SIZE,
        hash_block_size=BLOCK_SIZE,
        enable_caching=enable_caching,
    )
 def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None:
    manager = _manager(enable_caching=True)
    free_before = manager.block_pool.get_num_free_blocks()
    request = _request()
    # vLLM 0.26 adds a shared-prefix-boundary return value.
    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert blocks is not None
    assert len(manager.get_block_ids(request.request_id)[0]) == 2
    # Empty hashes intentionally mean request-local residency only. Calling
    # native cache_blocks() would require a canonical hash for the full prefix
    # block and is therefore outside this request's contract.
    manager.free(request)
    assert manager.block_pool.get_num_free_blocks() == free_before
 def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None:
    manager = _manager(enable_caching=True)
    prefix_hash = BlockHash(b"prefix".ljust(32, b"\0"))
    target_hash = BlockHash(b"target".ljust(32, b"\0"))
    request = _request(block_hashes=[prefix_hash, target_hash])
    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert blocks is not None
    manager.cache_blocks(request, request.prefix_len)
    manager.free(request)
    cached_request = _request(
        request_id="req/diffusion-kv/cached",
        block_hashes=[prefix_hash, target_hash],
    )
    cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request)
    # Only the stable prefix was published. The dynamic target remains a miss.
    assert num_computed_tokens == cached_request.prefix_len
    assert len(cached_blocks.blocks[0]) == 1
    cached_allocation = manager.allocate_slots(
        cached_request,
        num_new_tokens=cached_request.num_tokens - num_computed_tokens,
        num_new_computed_tokens=num_computed_tokens,
        new_computed_blocks=cached_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert cached_allocation is not None
    manager.free(cached_request)
@@ -244,7 +244,7 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
    captured: dict[str, object] = {}
    def rotary_position_embedding(x, cos, sin, **kwargs):
-        captured.update(cos=cos, sin=sin, kwargs=kwargs)
+        captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs)
        return x
    monkeypatch.setitem(
@@ -256,8 +256,10 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
    cos = torch.randn(11, 96, dtype=torch.bfloat16)
    sin = torch.randn(11, 96, dtype=torch.bfloat16)
-    apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
+    actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
    assert actual is x
    assert captured["x"] is x
    assert captured["cos"] is cos
    assert captured["sin"] is sin
    assert captured["kwargs"] == {
@@ -265,31 +267,3 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
        "head_first": False,
        "fused": True,
    }
 def test_mindie_rope_restores_3d_layout(monkeypatch: pytest.MonkeyPatch) -> None:
    """mindiesd rotary_position_embedding only accepts 4D ``x``."""
    captured: dict[str, torch.Tensor] = {}
    def rotary_position_embedding(x, cos, sin, **kwargs):
        # Mirror mindiesd's own input guard.
        assert x.dim() == 4, f"mindiesd requires 4D x, got {x.dim()}"
        captured.update(x=x)
        return x
    monkeypatch.setitem(
        sys.modules,
        "mindiesd",
        types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
    )
    S, H, D = 11, 3, 96
    x = torch.randn(S, H, D, dtype=torch.bfloat16)
    cos = torch.randn(S, D, dtype=torch.bfloat16)
    sin = torch.randn(S, D, dtype=torch.bfloat16)
    actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
    assert captured["x"].shape == (1, S, H, D)
    assert actual.shape == (S, H, D)
    torch.testing.assert_close(actual, x, atol=0, rtol=0)
@@ -16,7 +16,6 @@ from PIL import Image
 from torch import nn
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.experimental.world_models.session_state import SessionStateManager
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -426,29 +425,6 @@ def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> N
        pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0)
 def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None:
    pipeline = make_cosmos3_pipeline()
    robolab_inputs = object()
    captured: dict[str, Any] = {}
    expected = object()
    pipeline._build_robolab_policy_inputs = lambda sp, prompt, request_id: robolab_inputs
    def fake_forward_robolab(sp, inputs, pipeline_start, session_id=None):
        del sp, inputs, pipeline_start
        captured["session_id"] = session_id
        return expected
    pipeline._forward_robolab_policy = fake_forward_robolab
    request = SimpleNamespace(
        prompts=["policy"],
        sampling_params=make_sampling_params(),
        request_id="robolab-request-7",
    )
    assert pipeline.forward(request) is expected
    assert captured["session_id"] == "robolab-request-7"
@pytest.mark.parametrize(
    ("prompt", "sampling_params", "message"),
    [
@@ -619,7 +595,6 @@ def _make_od_config(
    return SimpleNamespace(
        enable_cpu_offload=False,
        enable_diffusion_pipeline_profiler=False,
        enable_session_state_manager=False,
        model="/nonexistent/model/path",
        dtype=torch.float32,
        flow_shift=None,
@@ -1698,32 +1673,6 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
    torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
 def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
    pipeline = make_cosmos3_pipeline()
    pipeline._use_session_state = True
    pipeline._memory_manager = SessionStateManager()
    class FailingProgress:
        def __iter__(self):
            raise RuntimeError("progress failed before first step")
    pipeline.progress_bar = lambda timesteps: FailingProgress()
    with pytest.raises(RuntimeError, match="progress failed"):
        pipeline.diffuse(
            latents=torch.zeros(1, 2, 1, 1, 1),
            timesteps=torch.tensor([7]),
            cond_ids=_ids(2),
            cond_mask=_mask(),
            uncond_ids=_ids(1),
            uncond_mask=_mask(),
            guidance_scale=1.0,
            shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
            session_id="request-that-fails",
        )
    assert "request-that-fails" not in pipeline._memory_manager
 def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
    pipeline = make_cosmos3_pipeline()
    latents = torch.zeros(1, 2, 1, 1, 1)
@@ -1753,28 +1702,6 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
    torch.testing.assert_close(result, torch.full_like(latents, 254.0))
 def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
    pipeline = make_cosmos3_pipeline()
    pipeline._use_session_state = True
    with pytest.raises(NotImplementedError, match="richer cache keys"):
        pipeline.diffuse_transfer(
            latents=torch.zeros(1, 2, 1, 1, 1),
            timesteps=torch.tensor([7]),
            cond_ids=_ids(2),
            cond_mask=_mask(),
            uncond_ids=_ids(1),
            uncond_mask=_mask(),
            guidance_scale=1.0,
            control_guidance=1.0,
            control_guidance_interval=None,
            control_latents=[],
            shared_kwargs={},
            velocity_mask=torch.ones(1, 1, 1, 1, 1),
            condition_latents=torch.zeros(1, 2, 1, 1, 1),
        )
 def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
    latents = torch.zeros(1, 2, 1, 1, 1)
    velocity_mask = torch.ones(1, 1, 1, 1, 1)
@@ -1,447 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import pytest
 import torch
 from transformers.utils.generic import ModelOutput
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
    ImageInfo,
    JointImageInfo,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    HunyuanImage3Pipeline,
    get_hunyuan_image_3_pre_process_func,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
    HunyuanPreparedLayout,
    build_hunyuan_diffusion_kv_requests,
    extract_hunyuan_prompt_inputs,
    hunyuan_num_image_tokens,
    normalize_hunyuan_cot_text,
    prepare_hunyuan_layout,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def test_dense_legacy_preprocess_does_not_initialize_layout_tokenizer(monkeypatch) -> None:
    hf_config = SimpleNamespace(vae_downsample_factor=(8, 8), patch_size=2)
    image_processor = SimpleNamespace(vision_encoder_processor=SimpleNamespace(patch_size=1))
    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
    monkeypatch.setattr(
        pipeline_module,
        "TokenizerWrapper",
        lambda _model: pytest.fail("dense_legacy must not initialize the Diffusion KV tokenizer"),
    )
    get_hunyuan_image_3_pre_process_func(
        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
    )
 class _FakeTokenizerWrapper:
    def __init__(self, *, prefix_lens: list[int]) -> None:
        self.prefix_lens = prefix_lens
        self.calls: list[dict] = []
        self.pad_token_id = 0
        self.eos_token_id = 2
        self.boi_token_id = 3
        self.end_recaption_token_id = 4
        self.end_answer_token_id = 5
        self.special_token_map = {f"<img_ratio_{index}>": 100 + index for index in range(33)}
    def apply_chat_template(self, **kwargs):
        self.calls.append(kwargs)
        rows = len(self.prefix_lens)
        has_cond_image = kwargs.get("batch_cond_image_info") is not None
        seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
        padded_seq_len = max(seq_lens)
        positions = torch.arange(padded_seq_len)
        joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
        gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
        all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
        cond_vae_image_mask = (
            torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
        )
        cond_vit_image_mask = (
            torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
        )
        sections = []
        for _ in range(rows):
            row_sections = []
            if has_cond_image:
                row_sections.append(
                    {
                        "type": "joint_image",
                        "token_height": [1, 1],
                        "token_width": [2, 2],
                    }
                )
            row_sections.append({"type": "gen_image", "token_height": 4, "token_width": 4})
            sections.append(row_sections)
        return {
            "output": TokenizerEncodeOutput(
                tokens=torch.arange(rows * padded_seq_len, dtype=torch.long).reshape(rows, padded_seq_len),
                gen_timestep_scatter_index=torch.tensor(self.prefix_lens, dtype=torch.long).reshape(rows, 1),
                cond_timestep_scatter_index=(torch.zeros(rows, 1, dtype=torch.long) if has_cond_image else None),
                all_image_slices=all_image_slices,
                gen_image_mask=torch.stack(
                    [(positions >= length + 1) & (positions < length + 17) for length in self.prefix_lens]
                ),
                cond_vae_image_mask=cond_vae_image_mask,
                cond_vit_image_mask=cond_vit_image_mask,
                joint_image_slices=joint_image_slices,
                gen_image_slices=gen_image_slices,
                real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
            ),
            "sections": sections,
        }
 class _FakeImageProcessor:
    def __init__(self, image_token_length: int = 16) -> None:
        self.image_token_length = image_token_length
        self.image_sizes: list[tuple[int, int]] = []
        self.vision_encoder_processor = SimpleNamespace(patch_size=1)
    def build_image_info(self, image_size):
        self.image_sizes.append(image_size)
        return ImageInfo(
            image_type="gen_image",
            image_width=image_size[1],
            image_height=image_size[0],
            token_width=4,
            token_height=4,
            image_token_length=self.image_token_length,
            base_size=1024,
            ratio_index=0,
        )
 def _components(prefix_lens: list[int]):
    return _FakeTokenizerWrapper(prefix_lens=prefix_lens), _FakeImageProcessor()
 def _prepare(
    request: OmniDiffusionRequest,
    tokenizer: _FakeTokenizerWrapper,
    image_processor: _FakeImageProcessor,
 ) -> HunyuanPreparedLayout:
    prepared_layout = prepare_hunyuan_layout(
        request,
        tokenizer_wrapper=tokenizer,
        image_processor=image_processor,
        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
        image_base_size=1024,
    )
    request.prepared_layout = prepared_layout
    return prepared_layout
 def _request(
    *,
    guidance_scale: float,
    prompt="draw a cat",
    request_id: str = "req",
 ) -> OmniDiffusionRequest:
    return OmniDiffusionRequest(
        prompt=prompt,
        sampling_params=OmniDiffusionSamplingParams(
            height=768,
            width=1024,
            guidance_scale=guidance_scale,
            num_inference_steps=4,
        ),
        request_id=request_id,
    )
 def test_builds_kv_request_lengths_without_model_execution() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=1.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
    assert all(isinstance(item, DiffusionKVRequest) for item in kv_requests)
    assert [(item.prefix_len, item.target_len, item.seq_len) for item in kv_requests] == [(12, 17, 32)]
    assert image_processor.image_sizes == [(768, 1024)]
    assert tokenizer.calls[0]["sequence_template"] == "instruct"
    assert tokenizer.calls[0]["cfg_factor"] == 1
    assert kv_requests[0].block_hashes == []
    assert kv_requests[0].kv_contexts == ()
    assert kv_requests[0].skip_reading_prefix_cache is True
 def test_builds_one_kv_request_per_cfg_row() -> None:
    tokenizer, image_processor = _components([12, 14])
    request = _request(guidance_scale=5.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert [item.sequence_id for item in kv_requests] == [0, 1]
    assert [item.prefix_len for item in kv_requests] == [12, 14]
    assert [item.seq_len for item in kv_requests] == [32, 34]
    assert tokenizer.calls[0]["cfg_factor"] == 2
 def _reference_image() -> JointImageInfo:
    return JointImageInfo(
        vae_image_info=ImageInfo(
            image_type="vae",
            token_width=8,
            token_height=8,
            image_token_length=64,
        ),
        vision_image_info=ImageInfo(
            image_type="siglip2",
            token_width=4,
            token_height=4,
            image_token_length=16,
        ),
        vision_encoder_kwargs={
            "spatial_shapes": torch.tensor([4, 4]),
            "pixel_attention_mask": torch.ones(4, dtype=torch.bool),
        },
    )
 def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
    tokenizer, image_processor = _components([20])
    joint_image = _reference_image()
    prompt = {
        "prompt": "edit this image",
        "additional_information": {"batch_cond_image_info": [joint_image]},
    }
    request = _request(guidance_scale=1.0, prompt=prompt)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert tokenizer.calls[0]["batch_cond_image_info"] == [[joint_image]]
    assert kv_requests[0].prefix_len == 20
    assert kv_requests[0].target_len == 17
    assert kv_requests[0].seq_len == 40
    assert kv_requests[0].kv_contexts == ()
 def _assert_nested_equal(actual, expected) -> None:
    if isinstance(expected, torch.Tensor):
        torch.testing.assert_close(actual, expected)
    elif isinstance(expected, dict):
        assert actual.keys() == expected.keys()
        for key in expected:
            _assert_nested_equal(actual[key], expected[key])
    elif isinstance(expected, list | tuple):
        assert isinstance(actual, type(expected))
        assert len(actual) == len(expected)
        for actual_item, expected_item in zip(actual, expected):
            _assert_nested_equal(actual_item, expected_item)
    else:
        assert actual == expected
@pytest.mark.parametrize(
    ("prefix_lens", "guidance_scale", "with_reference_image"),
    [
        pytest.param([12], 1.0, False, id="cfg-1"),
        pytest.param([12, 14], 5.0, False, id="cfg-2"),
        pytest.param([20], 1.0, True, id="cfg-1-reference-image"),
        pytest.param([20, 22], 5.0, True, id="cfg-2-reference-image"),
    ],
 )
 def test_prepared_model_inputs_match_local_tokenization(
    monkeypatch,
    prefix_lens: list[int],
    guidance_scale: float,
    with_reference_image: bool,
 ) -> None:
    tokenizer, image_processor = _components(prefix_lens)
    prompt: str | dict = "draw a cat"
    if with_reference_image:
        prompt = {
            "prompt": "edit this image",
            "additional_information": {"batch_cond_image_info": [_reference_image()]},
        }
    request = _request(guidance_scale=guidance_scale, prompt=prompt)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    prompts, cot_texts, system_prompt, batch_cond_image_info, bot_task = extract_hunyuan_prompt_inputs(
        [request.prompt],
        request.sampling_params.extra_args or {},
        request_id=request.request_id,
        allow_cond_image=True,
    )
    cot_text = (
        [normalize_hunyuan_cot_text(text) for text in cot_texts]
        if any(text is not None for text in cot_texts)
        else None
    )
    pipeline = object.__new__(HunyuanImage3Pipeline)
    pipeline._tkwrapper = tokenizer
    pipeline.image_processor = image_processor
    pipeline.generation_config = SimpleNamespace(sequence_template="instruct", drop_think=False)
    pipeline.config = SimpleNamespace(
        image_base_size=1024,
        attention_head_dim=2,
        rope_theta=10000.0,
    )
    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda _self: torch.device("cpu")))
    def fake_encode_cond_image(batch_cond_image_info_list, cfg_factor=1, generator=None):
        del generator
        rows = len(batch_cond_image_info_list) * cfg_factor
        return (
            torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 1, 2, 2),
            torch.arange(rows, dtype=torch.float32),
            torch.arange(rows * 3, dtype=torch.float32).reshape(rows, 3),
        )
    pipeline._encode_cond_image = fake_encode_cond_image
    def fake_build_batch_2d_rope(*, image_infos, seq_len, **_kwargs):
        geometry_sum = sum(
            int(image_slice.start or 0) + int(image_slice.stop or 0) + height + width
            for row in image_infos
            for image_slice, (height, width) in row
        )
        values = torch.tensor([len(image_infos), seq_len, geometry_sum], dtype=torch.float32)
        return values, values + 1
    monkeypatch.setattr(pipeline_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
    common_kwargs = dict(
        prompt=prompts,
        cot_text=cot_text,
        system_prompt=system_prompt,
        mode="gen_image",
        guidance_scale=guidance_scale,
        image_size=(request.sampling_params.height, request.sampling_params.width),
        generator=[torch.Generator().manual_seed(0)],
        batch_cond_image_info=batch_cond_image_info,
        bot_task=bot_task,
    )
    local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
    prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
    comparable_fields = (
        "input_ids",
        "position_ids",
        "custom_pos_emb",
        "image_mask",
        "gen_timestep_scatter_index",
        "cond_vae_images",
        "cond_timestep",
        "cond_vae_image_mask",
        "cond_vit_images",
        "cond_vit_image_mask",
        "vit_kwargs",
        "cond_timestep_scatter_index",
    )
    for field in comparable_fields:
        _assert_nested_equal(prepared_inputs[field], local_inputs[field])
    tokenizer_fields = (
        "tokens",
        "gen_image_mask",
        "cond_vae_image_mask",
        "cond_vit_image_mask",
        "gen_timestep_scatter_index",
        "cond_timestep_scatter_index",
        "real_pos",
        "all_image_slices",
        "joint_image_slices",
        "gen_image_slices",
    )
    for field in tokenizer_fields:
        _assert_nested_equal(
            getattr(prepared_inputs["tokenizer_output"], field),
            getattr(local_inputs["tokenizer_output"], field),
        )
    local_mask = pipeline._prepare_attention_mask_for_generation(
        local_inputs["input_ids"],
        pipeline.generation_config,
        local_inputs,
    )
    prepared_mask = pipeline._prepare_attention_mask_for_generation(
        prepared_inputs["input_ids"],
        pipeline.generation_config,
        prepared_inputs,
    )
    torch.testing.assert_close(prepared_mask, local_mask)
    assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
    num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
    local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
    prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
    local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
    prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
    for field in ("position_ids", "attention_mask", "gen_timestep_scatter_index", "full_attn_spans"):
        _assert_nested_equal(prepared_step_inputs[field], local_step_inputs[field])
 def test_rejects_tokenizer_cfg_row_mismatch() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=5.0)
    prepared_layout = _prepare(request, tokenizer, image_processor)
    with pytest.raises(ValueError, match="sequence count does not match"):
        build_hunyuan_diffusion_kv_requests(request, prepared_layout)
 def test_sender_endpoint_does_not_change_local_kv_request() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=1.0)
    request.kv_sender_info = {"host": "127.0.0.1", "zmq_port": 5000}
    prepared_layout = _prepare(request, tokenizer, image_processor)
    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
    assert kv_requests[0].seq_len == 32
    assert kv_requests[0].num_computed_tokens == 0
 def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch) -> None:
    tokenizer, image_processor = _components([12])
    hf_config = SimpleNamespace(
        vae_downsample_factor=(8, 8),
        patch_size=2,
        image_base_size=1024,
    )
    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
    monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
    monkeypatch.setattr(
        pipeline_module.GenerationConfig,
        "from_pretrained",
        lambda _model: SimpleNamespace(sequence_template="instruct", drop_think=False),
    )
    preprocess = get_hunyuan_image_3_pre_process_func(
        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
    )
    request = _request(guidance_scale=1.0)
    prepared_request = preprocess(request)
    assert prepared_request is request
    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
    assert request.diffusion_kv_requests is not None
    assert [(item.prefix_len, item.target_len, item.seq_len) for item in request.diffusion_kv_requests] == [
        (12, 17, 32)
    ]
    assert len(tokenizer.calls) == 1
@@ -7,10 +7,7 @@ import pytest
 import torch
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module
 import vllm_omni.diffusion.models.hunyuan_image3.request_layout as hy3_layout_module
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageInfo
 from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    _STEP_AR_KV,
    _STEP_CFG_FACTOR,
@@ -21,7 +18,6 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
    _STEP_PROMPT_KV,
    HunyuanImage3Pipeline,
 )
 from vllm_omni.diffusion.models.hunyuan_image3.request_layout import HunyuanPreparedLayout
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import StepRequestState
@@ -80,72 +76,6 @@ def _sampling_params(**extra_args):
    )
 def _prepared_layout() -> HunyuanPreparedLayout:
    return HunyuanPreparedLayout(
        tokenizer_output=TokenizerEncodeOutput(
            tokens=torch.arange(21).reshape(1, 21),
            gen_image_mask=torch.ones(1, 21, dtype=torch.bool),
            gen_timestep_scatter_index=torch.tensor([[4]]),
            all_image_slices=[[slice(5, 21)]],
            joint_image_slices=[[]],
            gen_image_slices=[[slice(5, 21)]],
            real_pos=torch.tensor([[21]]),
        ),
        rope_image_info=[[(slice(5, 21), (4, 4))]],
        generated_image_info=ImageInfo(
            image_type="gen_image",
            image_width=512,
            image_height=512,
            token_width=4,
            token_height=4,
            image_token_length=16,
        ),
    )
 def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
    pipeline = _pipeline()
    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
    rope_kwargs = {}
    def fake_build_batch_2d_rope(**kwargs):
        rope_kwargs.update(kwargs)
        return torch.zeros(1), torch.zeros(1)
    monkeypatch.setattr(hy3_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
    def fail_apply_chat_template(**_kwargs):
        pytest.fail("prepared Hunyuan execution must not apply the chat template again")
    pipeline._tkwrapper = SimpleNamespace(
        apply_chat_template=fail_apply_chat_template,
        eos_token_id=2,
        boi_token_id=3,
        end_recaption_token_id=4,
        end_answer_token_id=5,
    )
    pipeline.config = SimpleNamespace(
        attention_head_dim=2,
        rope_theta=10000.0,
    )
    prepared = _prepared_layout()
    model_inputs = pipeline.prepare_model_inputs(
        prompt="prompt",
        mode="gen_image",
        guidance_scale=1.0,
        image_size=(512, 512),
        device=torch.device("cpu"),
        generator=[torch.Generator().manual_seed(0)],
        prepared_layout=prepared,
    )
    assert model_inputs["tokenizer_output"] is prepared.tokenizer_output
    assert model_inputs["batch_gen_image_info"] == [prepared.generated_image_info]
    assert rope_kwargs["image_infos"] is prepared.rope_image_info
    torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
    pipeline = _pipeline()
    states = [_state("req-0", 1), _state("req-1", 3)]
@@ -201,13 +131,12 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
        captured.update(kwargs)
        raise RuntimeError("stop after prepare_model_inputs")
-    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
    pipeline.prepare_model_inputs = fake_prepare_model_inputs
    state = StepRequestState(
        request_id="req-bot-task",
        sampling=sampling,
        prompt=prompt_item,
        prepared_layout=_prepared_layout(),
    )
    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
@@ -215,7 +144,6 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
    assert captured["bot_task"] == expected_model_bot_task
    assert captured["system_prompt_bot_task"] == expected_system_bot_task
    assert captured["prepared_layout"] is state.prepared_layout
 def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
@@ -231,22 +159,23 @@ def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
        captured.update(kwargs)
        raise RuntimeError("stop after prepare_model_inputs")
-    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
    pipeline.prepare_model_inputs = fake_prepare_model_inputs
-    request = SimpleNamespace(
+    req = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="req-forward-bot-task",
                sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
                prompt={"prompt": "prompt", "bot_task": "vanilla"},
        prepared_layout=_prepared_layout(),
            )
-    req = DiffusionRequestBatch(requests=[request])
+        ]
    )
    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
        pipeline.forward(req)
    assert captured["bot_task"] == "think"
    assert captured["system_prompt_bot_task"] == "think"
    assert captured["prepared_layout"] is request.prepared_layout
 def test_grouped_denoise_rejects_non_sdpa_attention_backend():
@@ -1,133 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 from vllm_omni.diffusion.models.lingbot_video.image_condition import (
    apply_clean_prefix,
    geometry_align_image,
    prepare_ti2v_image_condition,
 )
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _gradient_image(width: int, height: int) -> Image.Image:
    x = np.linspace(0, 255, width, dtype=np.uint8)[None, :]
    y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
    rgb = np.stack(
        [
            np.broadcast_to(x, (height, width)),
            np.broadcast_to(y, (height, width)),
            np.full((height, width), 127, dtype=np.uint8),
        ],
        axis=-1,
    )
    return Image.fromarray(rgb, mode="RGB")
@pytest.mark.parametrize("source_size", [(160, 64), (64, 160), (96, 96)])
 def test_geometry_align_image_supports_source_aspect_ratios(source_size):
    pixel = geometry_align_image(_gradient_image(*source_size), 64, 96)
    assert pixel.shape == (1, 3, 1, 64, 96)
    assert pixel.dtype == torch.float32
    assert 0 <= pixel.min() <= pixel.max() <= 1
@pytest.mark.parametrize(
    ("source_size", "target_size"),
    [((160, 64), (64, 96)), ((64, 160), (96, 64))],
 )
 def test_geometry_align_image_center_crops_without_stretching(source_size, target_size):
    width, height = source_size
    target_height, target_width = target_size
    rgb = np.full((height, width, 3), (255, 0, 0), dtype=np.uint8)
    if width > target_width:
        left = (width - target_width) // 2
        rgb[:, left : left + target_width] = (0, 255, 0)
    else:
        top = (height - target_height) // 2
        rgb[top : top + target_height] = (0, 255, 0)
    pixel = geometry_align_image(
        Image.fromarray(rgb, mode="RGB"),
        target_height,
        target_width,
    )
    assert pixel.shape == (1, 3, 1, target_height, target_width)
    expected_green = torch.tensor((0.0, 1.0, 0.0)).view(1, 3, 1, 1, 1)
    assert torch.equal(pixel, expected_green.expand_as(pixel))
 class _LatentDistribution:
    def __init__(self, latent: torch.Tensor):
        self.latent = latent
    def sample(self, generator=None):
        del generator
        return self.latent
 class _CapturingVAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.anchor = nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(
            latents_mean=[0.0, 0.0, 0.0],
            latents_std=[1.0, 1.0, 1.0],
        )
        self.encoded_pixel = None
    def encode(self, pixel):
        self.encoded_pixel = pixel.detach().clone()
        return SimpleNamespace(latent_dist=_LatentDistribution(pixel))
 def test_ti2v_vlm_and_vae_paths_share_geometry_aligned_image():
    vae = _CapturingVAE()
    condition = prepare_ti2v_image_condition(
        _gradient_image(160, 64),
        height=64,
        width=96,
        vae=vae,
        vision_patch_size=16,
        device=torch.device("cpu"),
    )
    vae_pixel = (vae.encoded_pixel + 1.0) / 2.0
    vae_image = vae_pixel[0, :, 0].permute(1, 2, 0).mul(255).byte().numpy()
    assert condition.vlm_image.size == (96, 64)
    np.testing.assert_allclose(
        np.asarray(condition.vlm_image),
        vae_image,
        rtol=0,
        atol=1,
    )
 def test_apply_clean_prefix_only_overwrites_condition_frames():
    latents = torch.arange(20, dtype=torch.float32).reshape(1, 1, 5, 2, 2)
    original_tail = latents[:, :, 2:].clone()
    clean = torch.full((1, 1, 2, 2, 2), -3.0)
    returned = apply_clean_prefix(latents, clean)
    assert returned is latents
    assert torch.equal(latents[:, :, :2], clean)
    assert torch.equal(latents[:, :, 2:], original_tail)
 def test_apply_clean_prefix_rejects_incompatible_shapes():
    with pytest.raises(ValueError, match="shape must match"):
        apply_clean_prefix(
            torch.zeros(1, 2, 3, 4, 4),
            torch.zeros(1, 3, 1, 4, 4),
        )
@@ -6,10 +6,9 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
-from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -23,8 +22,6 @@ def _make_pipeline():
    nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.default_negative_prompt = "default negative"
    pipeline.default_image_negative_prompt = "default image negative"
    pipeline.img_prompt_template = "<image>"
    pipeline.od_config = SimpleNamespace(flow_shift=None)
    return pipeline
@@ -40,24 +37,17 @@ def test_lingbot_video_pipeline_import_and_registry():
        LingBotVideoPipeline,
        LingBotVideoTransformer3DModel,
        get_lingbot_video_post_process_func,
        get_lingbot_video_pre_process_func,
    )
    from vllm_omni.diffusion.registry import (
        _DIFFUSION_MODELS,
        _DIFFUSION_POST_PROCESS_FUNCS,
        _DIFFUSION_PRE_PROCESS_FUNCS,
    )
    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
    assert LingBotVideoPipeline is not None
    assert LingBotVideoTransformer3DModel is not None
    assert get_lingbot_video_pre_process_func is not None
    assert get_lingbot_video_post_process_func is not None
    assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
        "lingbot_video",
        "pipeline_lingbot_video",
        "LingBotVideoPipeline",
    )
    assert _DIFFUSION_PRE_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_pre_process_func"
    assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
@@ -68,7 +58,6 @@ def test_component_discovery_declarations():
    assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
    assert LingBotVideoPipeline._vae_modules == ["vae"]
    assert LingBotVideoPipeline.supports_step_execution is False
    assert LingBotVideoPipeline.max_outputs_per_prompt == 1
 def test_extra_body_params_include_video_flow_shift_alias():
@@ -80,47 +69,17 @@ def test_extra_body_params_include_video_flow_shift_alias():
    assert "batch_cfg" in params
-def test_preprocess_preserves_geometry_and_converts_rgb(tmp_path):
+def test_extract_prompt_string_and_mapping():
-    from vllm_omni.diffusion.models.lingbot_video import (
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
        get_lingbot_video_pre_process_func,
    )
-    image_path = tmp_path / "reference.png"
+    string_req = SimpleNamespace(prompt="a robot waves")
-    Image.new("L", (48, 32), color=127).save(image_path)
+    assert _extract_prompt(string_req) == ("a robot waves", None)
    request = OmniDiffusionRequest(
        prompt={
            "prompt": "a robot walks",
            "modalities": ["video"],
            "multi_modal_data": {"image": str(image_path)},
        },
        sampling_params=OmniDiffusionSamplingParams(),
        request_id="lingbot-preprocess-test",
    )
-    result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
+    mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
    assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
-    image = result.prompt["multi_modal_data"]["image"]
+    with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
-    assert image.mode == "RGB"
+        _extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
    assert image.size == (48, 32)
 def test_preprocess_marks_internal_image_warmup_as_video():
    from vllm_omni.diffusion.models.lingbot_video import (
        get_lingbot_video_pre_process_func,
    )
    request = OmniDiffusionRequest(
        prompt={
            "prompt": "dummy run",
            "multi_modal_data": {"image": Image.new("RGB", (32, 32))},
        },
        sampling_params=OmniDiffusionSamplingParams(),
        request_id=DUMMY_DIFFUSION_REQUEST_ID,
    )
    result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
    assert result.prompt["modalities"] == ["video"]
 def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
@@ -145,59 +104,6 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
    array = post(frames, SimpleNamespace(output_type="np"))
    assert isinstance(array, np.ndarray)
    assert array.shape == (2, 4, 4, 3)
    video = post({"video": frames}, SimpleNamespace(output_type="pt"))
    assert set(video) == {"video"}
    assert video["video"] is frames
    image = post(
        {"image": torch.ones(4, 4, 3)},
        SimpleNamespace(output_type="pt"),
    )
    assert set(image) == {"image"}
    assert isinstance(image["image"], np.ndarray)
    assert image["image"].shape == (4, 4, 3)
    latent = torch.ones(1, 4, 1, 2, 2)
    latent_output = post({"image": latent}, SimpleNamespace(output_type="latent"))
    assert set(latent_output) == {"image"}
    assert latent_output["image"] is latent
 def test_postprocess_preserves_image_contract_for_serving():
    from vllm_omni.diffusion.data import DiffusionOutput
    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
    from vllm_omni.diffusion.output_formatter import (
        format_diffusion_outputs,
        normalize_diffusion_postprocess_output,
    )
    from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result
    sampling = OmniDiffusionSamplingParams(
        num_inference_steps=1,
        output_type="pt",
        resolution=64,
    )
    request = OmniDiffusionRequest(
        prompt={"prompt": "a robot", "modalities": ["image"]},
        sampling_params=sampling,
        request_id="lingbot-image-contract-test",
    )
    raw_output = {"image": torch.ones(8, 12, 3)}
    processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
    normalized = normalize_diffusion_postprocess_output(processed)
    assert normalized.primary_key == "image"
    [result] = format_diffusion_outputs(
        request=request,
        od_config=SimpleNamespace(model_class_name="LingBotVideoPipeline"),
        diffusion_output=DiffusionOutput(output=raw_output),
        output_data=raw_output,
        postprocess_output=normalized,
    )
    assert len(result.images) == 1
    assert isinstance(result.images[0], np.ndarray)
    endpoint_images = _extract_images_from_result(result)
    assert len(endpoint_images) == 1
    assert endpoint_images[0].size == (12, 8)
 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
@@ -210,11 +116,7 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
    pipeline._generate = fake_generate
    req = _make_request_batch(
-        {
+        {"prompt": "a robot arm", "negative_prompt": "bad hands"},
            "prompt": "a robot arm",
            "negative_prompt": "bad hands",
            "modalities": ["video"],
        },
        height=192,
        width=320,
        num_frames=9,
@@ -227,12 +129,10 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
    out = pipeline.forward(req)
-    assert torch.equal(out.output["video"], torch.zeros(9, 192, 320, 3))
+    assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
    assert len(calls) == 1
    call = calls[0]
    assert call["prompt"] == "a robot arm"
    assert call["mode"].value == "t2v"
    assert call["input_image"] is None
    assert call["negative_prompt"] == "bad hands"
    assert call["height"] == 192
    assert call["width"] == 320
@@ -294,222 +194,6 @@ def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
    assert calls[0]["negative_prompt"] == "default negative"
@pytest.mark.parametrize(
    ("prompt", "num_frames", "expected_mode", "output_key", "default_negative"),
    [
        (
            {"prompt": "a still robot", "modalities": ["image"]},
            1,
            "t2i",
            "image",
            "default image negative",
        ),
        (
            {
                "prompt": "a moving robot",
                "modalities": ["video"],
                "multi_modal_data": {"image": Image.new("RGB", (48, 32), color="blue")},
            },
            9,
            "ti2v",
            "video",
            "default negative",
        ),
    ],
 )
 def test_forward_dispatches_generation_mode_and_output(
    prompt,
    num_frames,
    expected_mode,
    output_key,
    default_negative,
 ):
    pipeline = _make_pipeline()
    calls = []
    def fake_generate(**kwargs):
        calls.append(kwargs)
        return torch.zeros(192, 320, 3) if output_key == "image" else torch.zeros(num_frames, 192, 320, 3)
    pipeline._generate = fake_generate
    output = pipeline.forward(
        _make_request_batch(
            prompt,
            height=192,
            width=320,
            num_frames=num_frames,
            num_inference_steps=2,
            guidance_scale=3.0,
            seed=42,
            output_type="pt",
        )
    )
    assert set(output.output) == {output_key}
    assert calls[0]["mode"].value == expected_mode
    assert calls[0]["negative_prompt"] == default_negative
    if expected_mode == "ti2v":
        assert isinstance(calls[0]["input_image"], Image.Image)
 def test_batched_cfg_pads_image_conditioned_embeddings_and_masks():
    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
        _batch_cfg_prompt_inputs,
    )
    positive = torch.ones(1, 7, 4)
    positive_mask = torch.ones(1, 7, dtype=torch.long)
    negative = torch.full((1, 9, 4), 2.0)
    negative_mask = torch.ones(1, 9, dtype=torch.long)
    embeds, mask = _batch_cfg_prompt_inputs(
        positive,
        positive_mask,
        negative,
        negative_mask,
    )
    assert embeds.shape == (2, 9, 4)
    assert mask.shape == (2, 9)
    assert torch.equal(mask[0], torch.tensor([1, 1, 1, 1, 1, 1, 1, 0, 0]))
    assert torch.equal(mask[1], torch.ones(9, dtype=torch.long))
 class _RecordingTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.anchor = nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(in_channels=1)
        self.latent_inputs = []
    def forward(
        self,
        latents,
        timestep,
        prompt_embeds,
        *,
        encoder_attention_mask,
        return_dict,
    ):
        del timestep, prompt_embeds, encoder_attention_mask, return_dict
        self.latent_inputs.append(latents.detach().clone())
        return (torch.zeros_like(latents),)
 class _CorruptingScheduler:
    sigma_max = 1.0
    sigma_min = 0.0
    def set_timesteps(self, num_inference_steps, *, device, shift, **kwargs):
        del num_inference_steps, shift, kwargs
        self.timesteps = torch.tensor([1000.0, 500.0], device=device)
    def step(
        self,
        noise_pred,
        timestep,
        latents,
        *,
        return_dict,
        generator,
    ):
        del noise_pred, timestep, return_dict, generator
        return (latents + 7.0,)
@pytest.mark.parametrize(
    ("guidance_scale", "batch_cfg", "expected_prompts", "expected_transformer_calls"),
    [
        pytest.param(2.0, False, ["positive", "negative"], 4, id="sequential-cfg"),
        pytest.param(2.0, True, ["positive", "negative"], 2, id="batched-cfg"),
        pytest.param(1.0, False, ["positive"], 2, id="cfg-disabled"),
    ],
 )
 def test_ti2v_reinjects_clean_prefix_across_cfg_modes(
    guidance_scale,
    batch_cfg,
    expected_prompts,
    expected_transformer_calls,
 ):
    from vllm_omni.diffusion.models.lingbot_video import (
        LingBotGenerationMode,
        LingBotImageCondition,
    )
    pipeline = _make_pipeline()
    pipeline.transformer = _RecordingTransformer()
    pipeline.scheduler = _CorruptingScheduler()
    pipeline.progress_bar = lambda timesteps: timesteps
    pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 3, 1, 1)
    vlm_image = Image.new("RGB", (32, 32), color="red")
    condition = LingBotImageCondition(
        vlm_image=vlm_image,
        clean_latent=torch.full((1, 1, 1, 1, 1), 5.0),
    )
    pipeline.prepare_ti2v_image_condition = lambda *args, **kwargs: condition
    encode_calls = []
    def fake_encode(prompt, *, images, device):
        del device
        encode_calls.append((prompt, images))
        return torch.ones(1, 2, 4), torch.ones(1, 2, dtype=torch.long)
    pipeline.encode_prompt = fake_encode
    latents = pipeline._generate(
        prompt="positive",
        mode=LingBotGenerationMode.TI2V,
        input_image=vlm_image,
        negative_prompt="negative",
        height=16,
        width=16,
        num_frames=9,
        num_inference_steps=2,
        guidance_scale=guidance_scale,
        shift=3.0,
        output_type="latent",
        batch_cfg=batch_cfg,
    )
    assert [call[0] for call in encode_calls] == expected_prompts
    assert all(call[1] == [vlm_image] for call in encode_calls)
    assert len(pipeline.transformer.latent_inputs) == expected_transformer_calls
    for item in pipeline.transformer.latent_inputs:
        expected_prefix = condition.clean_latent.expand(item.shape[0], -1, -1, -1, -1)
        assert torch.equal(item[:, :, :1], expected_prefix)
    assert torch.equal(latents[:, :, :1], condition.clean_latent)
    assert torch.equal(latents[:, :, 1:], torch.full((1, 1, 2, 1, 1), 14.0))
 def test_t2i_decodes_and_returns_the_unique_frame():
    from vllm_omni.diffusion.models.lingbot_video import LingBotGenerationMode
    pipeline = _make_pipeline()
    pipeline.transformer = _RecordingTransformer()
    pipeline.scheduler = _CorruptingScheduler()
    pipeline.progress_bar = lambda timesteps: timesteps
    pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 1, 2, 2)
    pipeline.encode_prompt = lambda *args, **kwargs: (
        torch.ones(1, 2, 4),
        torch.ones(1, 2, dtype=torch.long),
    )
    pipeline._decode_latents = lambda latents: torch.ones(1, 16, 16, 3)
    image = pipeline._generate(
        prompt="a still robot",
        mode=LingBotGenerationMode.T2I,
        height=16,
        width=16,
        num_frames=1,
        num_inference_steps=2,
        guidance_scale=1.0,
        shift=3.0,
        output_type="pt",
    )
    assert image.shape == (16, 16, 3)
 def test_forward_rejects_multi_request_batches():
    pipeline = _make_pipeline()
    first = OmniDiffusionRequest(
@@ -527,45 +211,6 @@ def test_forward_rejects_multi_request_batches():
        pipeline.forward(DiffusionRequestBatch([first, second]))
 def test_forward_marks_invalid_request_contract_as_client_error():
    from vllm_omni.errors import OmniClientError
    pipeline = _make_pipeline()
    request = _make_request_batch(
        {
            "prompt": "unsupported image edit",
            "modalities": ["image"],
            "multi_modal_data": {"image": Image.new("RGB", (32, 32), color="blue")},
        },
        height=192,
        width=320,
        num_frames=1,
    )
    with pytest.raises(OmniClientError, match="does not accept") as exc_info:
        pipeline.forward(request)
    assert exc_info.value.status_code == 400
 def test_forward_marks_unsupported_output_count_as_client_error():
    from vllm_omni.errors import OmniClientError
    pipeline = _make_pipeline()
    request = _make_request_batch(
        {"prompt": "motion", "modalities": ["video"]},
        height=192,
        width=320,
        num_frames=9,
        num_outputs_per_prompt=2,
    )
    with pytest.raises(OmniClientError, match="one output per prompt") as exc_info:
        pipeline.forward(request)
    assert exc_info.value.status_code == 400
 def test_load_weights_rejects_external_weight_stream():
    pipeline = _make_pipeline()
@@ -1,350 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import pytest
 from PIL import Image
 from vllm_omni.diffusion.models.lingbot_video.request_utils import (
    LingBotGenerationMode,
    caption_from_lingbot_prompt,
    normalize_lingbot_num_frames,
    normalize_lingbot_request,
    resolve_lingbot_mode,
    resolve_lingbot_num_frames,
    resolve_lingbot_output_dimensions,
    resolve_lingbot_size,
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _request(prompt, **sampling_overrides):
    return SimpleNamespace(
        prompt=prompt,
        sampling_params=OmniDiffusionSamplingParams(**sampling_overrides),
    )
 def _normalize(prompt, **sampling_overrides):
    return normalize_lingbot_request(
        _request(prompt, **sampling_overrides),
        default_negative_prompt="video default",
        default_image_negative_prompt="image default",
    )
@pytest.mark.parametrize(
    ("prompt", "expected"),
    [
        ({"prompt": "still", "modalities": ["image"]}, LingBotGenerationMode.T2I),
        ({"prompt": "motion", "modalities": ["video"]}, LingBotGenerationMode.T2V),
        (
            {
                "prompt": "motion",
                "modalities": ["video"],
                "multi_modal_data": {"image": Image.new("RGB", (8, 8))},
            },
            LingBotGenerationMode.TI2V,
        ),
    ],
 )
 def test_resolve_lingbot_mode(prompt, expected):
    assert resolve_lingbot_mode(prompt) is expected
@pytest.mark.parametrize(
    ("prompt", "message"),
    [
        (
            {"prompt": "x", "modalities": ["image"], "multi_modal_data": {"image": Image.new("RGB", (8, 8))}},
            "does not accept",
        ),
        (
            {
                "prompt": "x",
                "modalities": ["video"],
                "multi_modal_data": {"image": [Image.new("RGB", (8, 8)), Image.new("RGB", (8, 8))]},
            },
            "exactly one",
        ),
        ({"prompt": "x", "modalities": ["video"], "multi_modal_data": {"video": [1]}}, "reference modalities"),
        ({"prompt": "x", "multi_modal_data": {"image": Image.new("RGB", (8, 8))}}, "explicitly set"),
        ({"prompt": "x", "modalities": ["audio"]}, "output modality"),
    ],
 )
 def test_resolve_lingbot_mode_rejects_invalid_media_combinations(prompt, message):
    with pytest.raises(ValueError, match=message):
        resolve_lingbot_mode(prompt)
 def test_caption_from_lingbot_prompt_preserves_text_and_compact_json():
    assert caption_from_lingbot_prompt("plain text") == "plain text"
    assert caption_from_lingbot_prompt({"caption": "caption only", "duration": 4}) == "caption only"
    assert caption_from_lingbot_prompt({"caption": {"中文": True, "nested": {"value": 2}}}) == (
        '{"中文":true,"nested":{"value":2}}'
    )
    assert (
        caption_from_lingbot_prompt(
            {"scene": "海边", "camera": {"move": "pan"}, "duration": 5, "fps": 24, "size": "320x192"}
        )
        == '{"scene":"海边","camera":{"move":"pan"}}'
    )
@pytest.mark.parametrize("prompt", [{"duration": 4, "fps": 24}, {"caption": ""}, ""])
 def test_caption_from_lingbot_prompt_rejects_empty_caption(prompt):
    with pytest.raises(ValueError, match="empty|must not be empty"):
        caption_from_lingbot_prompt(prompt)
@pytest.mark.parametrize(
    ("num_frames", "expected"),
    [(1, 1), (2, 5), (4, 5), (5, 5), (96, 97), (98, 101), (121, 121)],
 )
 def test_normalize_lingbot_num_frames_rounds_up_to_4n_plus_1(num_frames, expected):
    assert normalize_lingbot_num_frames(num_frames) == expected
@pytest.mark.parametrize("num_frames", [0, -1, 1.5, "invalid"])
 def test_normalize_lingbot_num_frames_rejects_invalid_values(num_frames):
    with pytest.raises(ValueError, match="num_frames"):
        normalize_lingbot_num_frames(num_frames)
@pytest.mark.parametrize(
    ("duration", "fps", "expected"),
    [(4, 24, 97), (5, 24, 121), (8, 24, 193), (4, 12, 49), (4, 30, 121)],
 )
 def test_resolve_lingbot_num_frames_rounds_up(duration, fps, expected):
    assert resolve_lingbot_num_frames(duration, fps) == expected
@pytest.mark.parametrize(("duration", "fps"), [(0, 24), (-1, 24), (0.01, 24), (1, 0), (1, -1)])
 def test_resolve_lingbot_num_frames_rejects_non_positive_or_empty_results(duration, fps):
    with pytest.raises(ValueError, match="duration|fps|at least one"):
        resolve_lingbot_num_frames(duration, fps)
@pytest.mark.parametrize(
    ("resolution", "ratio", "expected"),
    [
        ("192p", "1:1", (192, 192)),
        ("192p", "9:16", (192, 320)),
        ("192p", "16:9", (320, 192)),
        ("192p", "3:4", (192, 256)),
        ("192p", "4:3", (256, 192)),
        ("480p", "1:1", (480, 480)),
        ("480p", "9:16", (480, 832)),
        ("480p", "16:9", (832, 480)),
        ("480p", "3:4", (480, 640)),
        ("480p", "4:3", (640, 480)),
        ("720p", "1:1", (736, 736)),
        ("720p", "9:16", (736, 1280)),
        ("720p", "16:9", (1280, 736)),
        ("720p", "3:4", (736, 960)),
        ("720p", "4:3", (960, 736)),
        ("1080p", "1:1", (1088, 1088)),
        ("1080p", "9:16", (1088, 1920)),
        ("1080p", "16:9", (1920, 1088)),
        ("1080p", "3:4", (1088, 1440)),
        ("1080p", "4:3", (1440, 1088)),
        ("2k", "1:1", (1440, 1440)),
        ("2k", "9:16", (1440, 2560)),
        ("2k", "16:9", (2560, 1440)),
        ("2k", "3:4", (1440, 1920)),
        ("2k", "4:3", (1920, 1440)),
        ("4k", "1:1", (2176, 2176)),
        ("4k", "9:16", (2176, 3840)),
        ("4k", "16:9", (3840, 2176)),
        ("4k", "3:4", (2176, 2880)),
        ("4k", "4:3", (2880, 2176)),
    ],
 )
 def test_resolve_lingbot_size_supports_all_official_presets(resolution, ratio, expected):
    assert resolve_lingbot_size(resolution=resolution, ratio=ratio) == expected
 def test_resolve_lingbot_size_validates_sources_and_alignment():
    assert resolve_lingbot_size(size="320x192") == (192, 320)
    assert resolve_lingbot_size(width=320, height=192) == (192, 320)
    with pytest.raises(ValueError, match="provided together"):
        resolve_lingbot_size(width=320)
    with pytest.raises(ValueError, match="provided together"):
        resolve_lingbot_size(resolution="192p")
    with pytest.raises(ValueError, match="mutually exclusive"):
        resolve_lingbot_size(width=320, height=192, resolution="192p", ratio="9:16")
    with pytest.raises(ValueError, match="multiples of 16"):
        resolve_lingbot_size(width=321, height=192)
 def test_resolve_lingbot_output_dimensions_matches_request_priority():
    assert resolve_lingbot_output_dimensions(
        sampling_width=480,
        sampling_height=480,
        prompt_fields={"resolution": "720p", "ratio": "16:9"},
    ) == (736, 1280)
    assert resolve_lingbot_output_dimensions(
        sampling_width=480,
        sampling_height=480,
        prompt_fields={"width": 320, "height": 192},
    ) == (320, 192)
    assert resolve_lingbot_output_dimensions(
        sampling_width=320,
        sampling_height=192,
    ) == (320, 192)
    assert resolve_lingbot_output_dimensions() == (480, 480)
 def test_normalize_lingbot_request_uses_mode_defaults_and_preserves_empty_negative():
    t2i = _normalize(
        {"prompt": "still", "modalities": ["image"], "negative_prompt": ""},
        width=320,
        height=192,
        num_frames=1,
    )
    assert t2i.mode is LingBotGenerationMode.T2I
    assert t2i.num_frames == 1
    assert t2i.negative_prompt == ""
    t2v = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=9,
    )
    assert t2v.mode is LingBotGenerationMode.T2V
    assert t2v.negative_prompt == "video default"
 def test_normalize_lingbot_request_rejects_explicit_zero_flow_shift():
    with pytest.raises(ValueError, match="shift.*positive"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=1,
            extra_args={"flow_shift": 0},
        )
 def test_normalize_lingbot_request_rejects_explicit_zero_width():
    with pytest.raises(ValueError, match="width"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=0,
            height=192,
            num_frames=1,
        )
 def test_normalize_lingbot_request_rounds_all_video_frame_counts():
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=96,
        fps=24,
        frame_rate=24,
    )
    assert config.num_frames == 97
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=121,
        fps=24,
        frame_rate=24,
    )
    assert config.num_frames == 121
 def test_normalize_lingbot_request_prefers_explicit_frames_over_duration():
    online = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
        num_frames=9,
        fps=24,
        extra_args={"duration": 5},
    )
    assert online.num_frames == 9
    offline = _normalize(
        {
            "prompt": {"caption": "motion", "duration": 5, "num_frames": 9},
            "modalities": ["video"],
        },
        width=320,
        height=192,
        fps=24,
    )
    assert offline.num_frames == 9
 def test_normalize_lingbot_request_uses_duration_then_model_default():
    duration = _normalize(
        {"prompt": {"caption": "motion", "duration": 5}, "modalities": ["video"]},
        width=320,
        height=192,
        fps=24,
    )
    assert duration.num_frames == 121
    default = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=320,
        height=192,
    )
    assert default.num_frames == 81
 def test_normalize_lingbot_request_rejects_contract_conflicts():
    with pytest.raises(ValueError, match="text-to-image.*duration"):
        _normalize(
            {"prompt": {"caption": "still", "duration": 4}, "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=1,
        )
    with pytest.raises(ValueError, match="text-to-image requires"):
        _normalize(
            {"prompt": "still", "modalities": ["image"]},
            width=320,
            height=192,
            num_frames=5,
        )
    with pytest.raises(ValueError, match="seconds.*duration"):
        _normalize(
            {"prompt": "motion", "modalities": ["video"]},
            width=320,
            height=192,
            num_frames=96,
            extra_args={"seconds": 4, "duration": 4},
        )
 def test_normalize_lingbot_request_model_size_overrides_sampling_dimensions():
    # A prompt-provided size overrides the sampling width/height.
    config = _normalize(
        {"prompt": {"caption": "motion", "size": "320x192"}, "modalities": ["video"]},
        width=480,
        height=480,
        num_frames=81,
    )
    assert (config.height, config.width) == (192, 320)
    # Without a prompt size, the sampling dimensions are used.
    config = _normalize(
        {"prompt": "motion", "modalities": ["video"]},
        width=480,
        height=480,
        num_frames=81,
    )
    assert (config.height, config.width) == (480, 480)
@@ -9,75 +9,10 @@ from unittest.mock import Mock
 import numpy as np
 import pytest
 import torch
 import torch.nn as nn
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.request import OmniDiffusionRequest
    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.partition = "fl2va"
    pipeline.supported_tasks = frozenset({"t2va"})
    pipeline.default_video_shift = 12.0
    pipeline.default_audio_shift = 3.0
    pipeline.device = torch.device("cpu")
    pipeline.od_config = SimpleNamespace()
    cache_spec = object()
    quality_plan = SimpleNamespace(cache_dit=cache_spec)
    pipeline._quality_policy = Mock()
    pipeline._quality_policy.resolve.return_value = quality_plan
    events = []
    pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
    pipeline.encode_prompt = Mock(
        return_value=(
            torch.ones(1, 2),
            torch.ones(1, dtype=torch.long),
        )
    )
    def diffuse(**kwargs):
        events.append(("diffuse", kwargs))
        return torch.zeros(1), torch.zeros(1)
    pipeline.diffuse = diffuse
    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
    sampling = OmniDiffusionSamplingParams(
        quality="high",
        width=1344,
        height=768,
        fps=24,
        num_frames=124,
        num_inference_steps=50,
        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
    )
    batch = DiffusionRequestBatch(
        [
            OmniDiffusionRequest(
                prompt="quality boundary",
                sampling_params=sampling,
                request_id="quality-boundary",
            )
        ]
    )
    output = pipeline.forward(batch)
    assert events[0] == ("prepare", cache_spec)
    assert events[1][0] == "diffuse"
    pipeline._quality_policy.resolve.assert_called_once_with(
        quality="high",
        num_inference_steps=50,
        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
    )
    assert output.output == pipeline.decode.return_value
 def test_pipeline_import_registry_and_component_discovery():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.registry import (
@@ -1466,149 +1401,3 @@ def test_g4_reference_video_metadata_validation(field, value, message, tmp_path)
    metadata[field] = value
    with pytest.raises(ValueError, match=message):
        _validate_reference_video_metadata(metadata, index=0, source=str(tmp_path / "reference.mp4"))
 _ENCODER_HIDDEN = 4
 _ENCODER_HEAD_DIM = 2
 _ENCODER_NUM_HEADS = 2
 _ENCODER_NUM_KV_HEADS = 1
 _ENCODER_INTERMEDIATE = 4
 # One distinct value per checkpoint tensor, none of them a parameter initializer
 # (fused weights start as `torch.empty`, norms as ones). A uniform fill would
 # prove only that every row was written, not that each source landed in its own
 # row range, which is the invariant the fused loaders have to get right.
 _SOURCE_FILL = {
    "self_attn.q_proj.weight": 3.0,
    "self_attn.k_proj.weight": 4.0,
    "self_attn.v_proj.weight": 5.0,
    "mlp.gate_proj.weight": 6.0,
    "mlp.up_proj.weight": 7.0,
    "input_layernorm.weight": 8.0,
 }
 _QKV_TARGET = "text_model.layers.0.self_attn.qkv_proj.weight"
 _GATE_UP_TARGET = "text_model.layers.0.mlp.gate_up_proj.weight"
 _NORM_TARGET = "text_model.layers.0.input_layernorm.weight"
 def _one_layer_text_encoder():
    """Encoder stub holding both fused weights plus one plain parameter."""
    from vllm_omni.diffusion.models.minimax_h3.encoder import (
        MiniMaxH3Qwen3VLEncoder,
        MiniMaxH3Qwen3VLMergedColumnParallelLinear,
        MiniMaxH3Qwen3VLQKVParallelLinear,
        MiniMaxH3Qwen3VLRMSNorm,
    )
    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _SingleRankEncoderGroup
    group = _SingleRankEncoderGroup(0)
    layer = nn.Module()
    layer.self_attn = nn.Module()
    layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
        group,
        hidden_size=_ENCODER_HIDDEN,
        num_heads=_ENCODER_NUM_HEADS,
        num_kv_heads=_ENCODER_NUM_KV_HEADS,
        head_dim=_ENCODER_HEAD_DIM,
        dtype=torch.float32,
    )
    layer.mlp = nn.Module()
    layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
        group,
        input_size=_ENCODER_HIDDEN,
        intermediate_size=_ENCODER_INTERMEDIATE,
        dtype=torch.float32,
    )
    layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)
    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
    nn.Module.__init__(encoder)
    encoder.text_model = nn.Module()
    encoder.text_model.layers = nn.ModuleList([layer])
    return encoder
 def _write_text_encoder_checkpoint(path, *, omit=()):
    """Write a one-layer Qwen3-VL-named checkpoint, one fill value per source."""
    import json
    import safetensors.torch
    prefix = "model.language_model.layers.0"
    q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
    kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM
    shapes = {
        "self_attn.q_proj.weight": (q_rows, _ENCODER_HIDDEN),
        "self_attn.k_proj.weight": (kv_rows, _ENCODER_HIDDEN),
        "self_attn.v_proj.weight": (kv_rows, _ENCODER_HIDDEN),
        "mlp.gate_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN),
        "mlp.up_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN),
        "input_layernorm.weight": (_ENCODER_HIDDEN,),
    }
    tensors = {f"{prefix}.{source}": torch.full(shape, _SOURCE_FILL[source]) for source, shape in shapes.items()}
    for source in omit:
        del tensors[f"{prefix}.{source}"]
    safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
    index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
    (path / "model.safetensors.index.json").write_text(json.dumps(index))
@pytest.mark.parametrize(
    ("omitted_sources", "expected_detail", "unreported_target"),
    [
        (("self_attn.q_proj.weight",), f"{_QKV_TARGET}: ['q']", _GATE_UP_TARGET),
        (("self_attn.k_proj.weight",), f"{_QKV_TARGET}: ['k']", _GATE_UP_TARGET),
        (("self_attn.v_proj.weight",), f"{_QKV_TARGET}: ['v']", _GATE_UP_TARGET),
        (("mlp.gate_proj.weight",), f"{_GATE_UP_TARGET}: ['gate']", _QKV_TARGET),
        (("mlp.up_proj.weight",), f"{_GATE_UP_TARGET}: ['up']", _QKV_TARGET),
        (
            ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"),
            f"{_QKV_TARGET}: ['q', 'k', 'v']",
            _GATE_UP_TARGET,
        ),
    ],
 )
 def test_encoder_load_reports_missing_fused_source_shards(
    tmp_path, omitted_sources, expected_detail, unreported_target
 ):
    encoder = _one_layer_text_encoder()
    _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
    with pytest.raises(RuntimeError) as excinfo:
        encoder._load_weights(str(tmp_path))
    message = str(excinfo.value)
    assert expected_detail in message
    assert unreported_target not in message
 def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
    encoder = _one_layer_text_encoder()
    _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
    with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
        encoder._load_weights(str(tmp_path))
 def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
    encoder = _one_layer_text_encoder()
    _write_text_encoder_checkpoint(tmp_path)
    encoder._load_weights(str(tmp_path))
    params = dict(encoder.named_parameters())
    q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
    kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM
    qkv = params[_QKV_TARGET]
    gate_up = params[_GATE_UP_TARGET]
    slices = {
        "self_attn.q_proj.weight": qkv[:q_rows],
        "self_attn.k_proj.weight": qkv[q_rows : q_rows + kv_rows],
        "self_attn.v_proj.weight": qkv[q_rows + kv_rows :],
        "mlp.gate_proj.weight": gate_up[:_ENCODER_INTERMEDIATE],
        "mlp.up_proj.weight": gate_up[_ENCODER_INTERMEDIATE:],
        "input_layernorm.weight": params[_NORM_TARGET],
    }
    for source, rows in slices.items():
        assert torch.equal(rows, torch.full_like(rows, _SOURCE_FILL[source])), source
@@ -19,7 +19,7 @@ MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
 def _assert_joint_output(outputs):
    assert len(outputs) == 1
    frames = np.asarray(outputs[0].images[0])
-    assert frames.shape == (107, 256, 448, 3)
+    assert frames.shape == (39, 256, 448, 3)
    multimodal = outputs[0].multimodal_output
    assert multimodal is not None
    assert np.asarray(multimodal["audio"]).shape[1] == 2
@@ -122,58 +122,6 @@ def test_minimax_h3_t2va_fp8_single_gpu_smoke():
                output_type="np",
                extra_args={
                    "task": "t2va",
                    "duration": 4.0,
                    "aspect_ratio": "16:9",
                    "flow_shift": 12.0,
                    "audio_flow_shift": 3.0,
                },
            ),
            use_tqdm=False,
        )
    finally:
        engine.close()
    _assert_joint_output(outputs)
@pytest.mark.skipif(
    not os.environ.get(MODEL_ENV),
    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
 )
 def test_minimax_h3_t2va_teacache_single_gpu_smoke():
    import torch
    from vllm_omni.diffusion.data import DiffusionParallelConfig
    from vllm_omni.entrypoints.omni import Omni
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
    if torch.accelerator.device_count() < 1:
        pytest.skip("MiniMax H3 TeaCache smoke requires a CUDA device")
    engine = Omni(
        model=os.environ[MODEL_ENV],
        cache_backend="tea_cache",
        cache_config={"rel_l1_thresh": 0.17},
        parallel_config=DiffusionParallelConfig(ulysses_degree=1),
        trust_remote_code=True,
        enable_cpu_offload=True,
        enforce_eager=True,
    )
    try:
        outputs = engine.generate(
            "A quiet cinematic night scene with matching ambient sound.",
            OmniDiffusionSamplingParams(
                height=256,
                width=448,
                num_frames=29,
                fps=24,
                num_inference_steps=2,
                seed=42,
                output_type="np",
                extra_args={
                    "task": "t2va",
                    "duration": 4.0,
                    "aspect_ratio": "16:9",
                    "flow_shift": 12.0,
                    "audio_flow_shift": 3.0,
                },
@@ -69,7 +69,6 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    attention = object.__new__(MiniMaxH3Attention)
    nn.Module.__init__(attention)
    attention.rot_dim = 96
    attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
    attention.rope._forward_method = attention.rope.forward_native
@@ -1,212 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 import pytest
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 from vllm_omni.errors import OmniClientError
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 def _quality_od_config(*, cache_backend: str = "cache_dit", **overrides):
    values = {
        "cache_backend": cache_backend,
        "cache_config": DiffusionCacheConfig(),
        "diffusion_attention_config": SimpleNamespace(default=None, per_role={}),
        "enable_cpu_offload": False,
        "enable_layerwise_offload": False,
        "enable_distributed_layerwise_offload": False,
        "enforce_eager": True,
        "num_gpus": 4,
        "quantization_config": None,
        "parallel_config": SimpleNamespace(
            allgather_degree=1,
            cfg_parallel_size=1,
            data_parallel_size=1,
            pipeline_parallel_size=1,
            ring_degree=1,
            sequence_parallel_size=4,
            tensor_parallel_size=1,
            ulysses_degree=4,
            use_hsdp=False,
        ),
    }
    values.update(overrides)
    return SimpleNamespace(**values)
 def _resolve_quality(policy, **overrides):
    values = {
        "quality": "high",
        "num_inference_steps": 50,
        "extra_args": None,
    }
    values.update(overrides)
    return policy.resolve(**values)
 def test_high_quality_policy_emits_cache_profile_without_deployment_gates():
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
        MiniMaxH3QualityPolicy,
    )
    policy = MiniMaxH3QualityPolicy(
        _quality_od_config(
            cache_backend="none",
            enable_layerwise_offload=True,
            enforce_eager=False,
            num_gpus=1,
        )
    )
    plan = _resolve_quality(policy, num_inference_steps=17)
    spec = plan.cache_dit
    assert spec is not None
    assert spec.installation_key == "minimax_h3.high"
    assert spec.num_inference_steps == 17
    assert spec.cache_config.Fn_compute_blocks == 1
    assert spec.cache_config.Bn_compute_blocks == 0
    assert spec.cache_config.max_warmup_steps == 4
    assert spec.cache_config.residual_diff_threshold == pytest.approx(0.04)
    assert spec.cache_config.max_continuous_cached_steps == 1
    assert spec.cache_config.enable_taylorseer is False
    assert spec.cache_config.scm_steps_mask_policy is None
 def test_h3_force_refresh_is_model_owned_and_request_scoped():
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
    policy = MiniMaxH3QualityPolicy(_quality_od_config())
    plan = _resolve_quality(
        policy,
        quality="high",
        num_inference_steps=17,
        extra_args={"force_refresh_step_hint": 7},
    )
    assert plan.cache_dit is not None
    assert plan.cache_dit.installation_key == "minimax_h3.high:force_refresh=7:once"
    assert plan.cache_dit.cache_config.force_refresh_step_hint == 7
    assert plan.cache_dit.cache_config.force_refresh_step_policy == "once"
 def test_h3_force_refresh_policy_supports_repeat_and_does_not_mutate_generic_config():
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
    config = _quality_od_config()
    plan = _resolve_quality(
        MiniMaxH3QualityPolicy(config),
        quality=None,
        num_inference_steps=20,
        extra_args={"force_refresh_step_hint": 5, "force_refresh_step_policy": "repeat"},
    )
    assert plan.cache_dit is not None
    assert plan.cache_dit.installation_key == "minimax_h3.generic:force_refresh=5:repeat"
    assert plan.cache_dit.cache_config is not config.cache_config
    assert plan.cache_dit.cache_config.force_refresh_step_hint == 5
    assert plan.cache_dit.cache_config.force_refresh_step_policy == "repeat"
    assert config.cache_config.force_refresh_step_hint is None
    assert config.cache_config.force_refresh_step_policy == "once"
    restored = _resolve_quality(MiniMaxH3QualityPolicy(config), quality=None)
    assert restored.cache_dit is not None
    assert restored.cache_dit.installation_key == "minimax_h3.generic"
    assert restored.cache_dit.cache_config is config.cache_config
@pytest.mark.parametrize(
    ("extra_args", "message"),
    [
        ({"force_refresh_step_hint": True}, "positive integer"),
        ({"force_refresh_step_hint": 0}, "between 1"),
        ({"force_refresh_step_hint": 51}, "between 1"),
        ({"force_refresh_step_hint": 5, "force_refresh_step_policy": "always"}, "one of"),
        ({"force_refresh_step_policy": "repeat"}, "requires force_refresh_step_hint"),
    ],
 )
 def test_h3_force_refresh_rejects_invalid_values(extra_args, message):
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
    with pytest.raises(OmniClientError, match=message):
        _resolve_quality(
            MiniMaxH3QualityPolicy(_quality_od_config()),
            extra_args=extra_args,
        )
@pytest.mark.parametrize(
    ("cache_backend", "quality"),
    [
        ("cache_dit", "lossless"),
        ("none", None),
    ],
 )
 def test_h3_force_refresh_requires_active_cache_target(cache_backend, quality):
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
    policy = MiniMaxH3QualityPolicy(_quality_od_config(cache_backend=cache_backend))
    with pytest.raises(OmniClientError, match="require an active Cache-DiT"):
        _resolve_quality(
            policy,
            quality=quality,
            extra_args={"force_refresh_step_hint": 7},
        )
@pytest.mark.parametrize(
    ("cache_backend", "quality", "expected_key"),
    [
        ("cache_dit", None, "minimax_h3.generic"),
        ("cache_dit", "lossless", None),
        ("cache_dit", "high", "minimax_h3.high"),
        ("none", None, None),
        ("none", "lossless", None),
        ("none", "high", "minimax_h3.high"),
    ],
 )
 def test_model_policy_owns_request_cache_target(cache_backend, quality, expected_key):
    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
        MiniMaxH3QualityPolicy,
    )
    config = _quality_od_config(cache_backend=cache_backend)
    policy = MiniMaxH3QualityPolicy(config)
    plan = _resolve_quality(
        policy,
        quality=quality,
    )
    if expected_key is None:
        assert plan.cache_dit is None
    else:
        assert plan.cache_dit is not None
        assert plan.cache_dit.installation_key == expected_key
        if quality is None:
            assert plan.cache_dit.cache_config is config.cache_config
 def test_h3_adopts_runner_installed_cache_dit_backend():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    backend = object()
    pipeline = object.__new__(MiniMaxH3Pipeline)
    pipeline._cache_dit_runtime = SimpleNamespace(
        adopt=lambda adopted, *, installation_key: setattr(
            pipeline,
            "adopted",
            (adopted, installation_key),
        ),
        is_enabled=True,
    )
    pipeline.adopt_cache_dit_backend(backend)
    assert pipeline.adopted == (backend, "minimax_h3.generic")
    assert pipeline.is_cache_dit_enabled()
@@ -178,9 +178,10 @@ class TestAsyncOutputLoopLogic:
        t = threading.Thread(target=proc._async_output_loop, daemon=True)
        t.start()
-        # Give the loop time to process, then send the shutdown sentinel.
+        # Give the loop time to process, then unblock and stop.
        time.sleep(0.5)
-        proc._async_output_queue.put(None)
+        proc._running = False
        proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
        t.join(timeout=2.0)
        mock_pack.assert_called()
@@ -190,43 +191,3 @@ class TestAsyncOutputLoopLogic:
            for c in proc.result_mq.enqueue.call_args_list
        )
        assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
 class TestWorkerProcShutdown:
    def test_shutdown_stops_async_thread_and_releases_queues(self):
        import threading
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()
        proc.mq = MagicMock()
        broadcast_mq = proc.mq
        result_mq = proc.result_mq
        thread = threading.Thread(target=proc._async_output_queue.get)
        proc._async_output_thread = thread
        thread.start()
        proc.shutdown()
        assert not thread.is_alive()
        proc.worker.shutdown.assert_called_once_with()
        broadcast_mq.shutdown.assert_called_once_with()
        result_mq.shutdown.assert_called_once_with()
        assert proc.mq is None
        assert proc.result_mq is None
    def test_shutdown_releases_queues_when_worker_shutdown_fails(self):
        proc = _make_worker_proc(step_execution=True)
        proc.worker = MagicMock()
        proc.worker.shutdown.side_effect = RuntimeError("shutdown failed")
        proc.mq = MagicMock()
        broadcast_mq = proc.mq
        result_mq = proc.result_mq
        with pytest.raises(RuntimeError, match="shutdown failed"):
            proc.shutdown()
        broadcast_mq.shutdown.assert_called_once_with()
        result_mq.shutdown.assert_called_once_with()
        assert proc.mq is None
        assert proc.result_mq is None
@@ -15,11 +15,7 @@ from vllm_omni.diffusion.data import (
    DiffusionParallelConfig,
    OmniDiffusionConfig,
 )
-from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.model_metadata import (
    HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
    QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
 )
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -125,11 +121,6 @@ class TestCreateDefaultDiffusion:
        assert ea["enforce_eager"] is True
        assert ea["lora_path"] == "/tmp/lora"
    def test_diffusion_kv_mode_roundtrip(self):
        od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
        assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
@@ -145,30 +136,6 @@ def test_task_type_roundtrip():
    assert od.task_type == "model-defined-task"
 def test_architecture_name_resolves_via_pipeline_class_fallback():
    # Direct hit: the architecture name itself is the metadata table key.
    qwen_od_config = OmniDiffusionConfig(
        model="Qwen/Qwen-Image-Edit-2511",
        model_class_name="QwenImageEditPlusPipeline",
    )
    qwen_od_config.update_multimodal_support()
    assert qwen_od_config.supports_multimodal_inputs is True
    assert qwen_od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
    # Fallback: the HF architecture name differs from the internal pipeline class
    # name, so ``get_diffusion_model_metadata`` must consult ``_DIFFUSION_MODELS``
    # and resolve the metadata via the pipeline class name.
    hunyuan_od_config = OmniDiffusionConfig(
        model="tencent/HunyuanImage-3.0-Instruct",
        model_class_name="HunyuanImage3ForCausalMM",
    )
    hunyuan_od_config.update_multimodal_support()
    assert hunyuan_od_config.supports_multimodal_inputs is True
    assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
 def test_additional_config_roundtrip():
    additional_config = {"torchair_graph_config": {"enabled": True}}
    od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)
@@ -25,7 +25,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
    CachedRequestData,
    NewRequestData,
    _AdmissionWaitDecision,
 )
 from vllm_omni.diffusion.sched.interface import (
    DiffusionSchedulerOutput as RealDiffusionSchedulerOutput,
@@ -76,34 +75,6 @@ class MockScheduler:
    def num_running_requests(self) -> int:
        return 0
    def get_admission_wait_decision(
        self,
        *,
        now: float,
        dp_concurrent: bool = False,
    ) -> _AdmissionWaitDecision:
        del dp_concurrent
        return _AdmissionWaitDecision(
            should_wait=True,
            deadline=now + 0.5,
            stable_window_s=0.05,
            max_batch=self.max_num_running_reqs,
        )
    def should_end_admission_wait(
        self,
        decision: _AdmissionWaitDecision,
        *,
        now: float,
        stable_since: float,
    ) -> bool:
        waiting = self.num_waiting_requests()
        return (
            waiting >= decision.max_batch
            or (waiting > 0 and now - stable_since >= decision.stable_window_s)
            or (decision.deadline is not None and now >= decision.deadline)
        )
    def schedule(self) -> DiffusionSchedulerOutput:
        if not self._waiting_queue:
            return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -137,17 +108,6 @@ class MockScheduler:
        pass
 def _make_admission_engine(pre_process_func) -> DiffusionEngine:
    engine = object.__new__(DiffusionEngine)
    engine.scheduler = MockScheduler()
    engine.pre_process_func = pre_process_func
    engine._out_streams = {}
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._closed = False
    return engine
 class _BatchCapablePipeline:
    supports_request_batch = True
@@ -371,55 +331,6 @@ class TestRequestBatchCapability:
            DiffusionEngine(od_config)
        fake_executor_cls.assert_not_called()
    def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args=None,
            streaming_output=False,
            max_num_seqs=2,
            request_batch_max_wait_ms=0,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        fake_executor = SimpleNamespace(
            execute_request=mocker.Mock(return_value="per-request"),
            execute_batch=mocker.Mock(return_value="batch"),
            execute_step=mocker.Mock(return_value="step"),
        )
        fake_executor_cls = mocker.Mock(return_value=fake_executor)
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )
        engine = DiffusionEngine(od_config)
        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
        assert engine.supports_request_batch is False
        assert engine.dp_concurrent is True
        assert engine.scheduler.max_num_running_reqs == 2
        assert od_config.request_batch_max_wait_ms == 0
        fake_executor_cls.assert_called_once_with(od_config)
    @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
    def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
        self,
@@ -565,11 +476,6 @@ class TestRequestBatchAdmission:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0)
    @pytest.mark.parametrize("bad_wait", [float("nan"), float("inf"), float("-inf")])
    def test_config_rejects_nonfinite_request_batch_max_wait_ms(self, bad_wait: float) -> None:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=bad_wait)
    def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None:
        config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5)
@@ -637,7 +543,7 @@ class TestRequestBatchAdmission:
        start = time.monotonic()
        with engine._cv:
-            engine._wait_for_admission_if_needed_locked()
+            engine._wait_for_request_batch_admission_locked()
        waited_s = time.monotonic() - start
        # Stable-window exit (~50ms), not the full 1000ms deadline.
@@ -728,34 +634,6 @@ async def _consume_final_output(generator):
    return final_output
@pytest.mark.cpu
@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"])
 def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None:
    raw_request = OmniDiffusionRequest(
        prompt="raw",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
    )
    prepared_request = OmniDiffusionRequest(
        prompt="prepared",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
        prepared_layout=object(),
    )
    preprocess_calls = []
    def preprocess(request):
        preprocess_calls.append(request)
        return prepared_request
    engine = _make_admission_engine(preprocess)
    getattr(engine, entrypoint)(raw_request)
    assert preprocess_calls == [raw_request]
    assert engine.scheduler._waiting_queue == [prepared_request]
@pytest.mark.cpu
@pytest.mark.asyncio
 async def test_async_add_req_and_stream_response():
@@ -127,7 +127,6 @@ def _make_request(tag: str):
        request_id=tag,
        prompt=f"prompt_{tag}",
        sampling_params=SimpleNamespace(**sampling_params),
        diffusion_kv_requests=None,
    )
@@ -146,7 +145,7 @@ def _make_engine_with_loop(
    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
    sched = RequestScheduler()
-    sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+    sched.initialize(SimpleNamespace(max_num_seqs=1))
    engine.scheduler = sched
    engine.step_execution = False
    engine.supports_request_batch = False
@@ -2,19 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
 import multiprocessing as mp
 import queue
 import threading
 import time
 from multiprocessing import shared_memory
 import numpy as np
 import pytest
 import torch
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import (
    _SHM_TENSOR_THRESHOLD,
    DIFFUSION_RPC_RESULT_ENVELOPE,
@@ -38,129 +30,6 @@ def _cleanup_shm_handle(value: object) -> None:
            _unpack_if_shm_handle(value)
 def _result_queue_worker(connection, rank: int) -> None:
    """Send one nested NumPy result through a worker-owned MessageQueue."""
    result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
    try:
        connection.send(result_mq.export_handle())
        assert connection.recv() == "reader-ready"
        frames = np.full(300_000, rank, dtype=np.float32)
        audio = np.arange(300_000, dtype=np.float32) + rank
        output = DiffusionOutput(
            output={
                "rank": rank,
                "nested": {"frames": frames},
                "audio": [audio],
            }
        )
        pack_diffusion_output_shm(output)
        connection.send(
            [
                output.output["nested"]["frames"]["name"],
                output.output["audio"][0]["name"],
            ]
        )
        result_mq.enqueue(output)
        assert connection.recv() == "shutdown"
    finally:
        result_mq.shutdown()
        connection.close()
 def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
    """Exercise the production per-worker queue/pump/SHM lifecycle."""
    ctx = mp.get_context("spawn")
    parent_connections = []
    processes = []
    result_mqs = []
    shm_names = []
    executor = None
    try:
        for rank in range(2):
            parent_connection, child_connection = ctx.Pipe()
            process = ctx.Process(
                target=_result_queue_worker,
                args=(child_connection, rank),
                name=f"IPCResultWorker-{rank}",
            )
            process.start()
            child_connection.close()
            parent_connections.append(parent_connection)
            processes.append(process)
        for connection in parent_connections:
            handle = connection.recv()
            result_mqs.append(MessageQueue.create_from_handle(handle, 0))
            connection.send("reader-ready")
            shm_names.extend(connection.recv())
        executor = object.__new__(MultiprocDiffusionExecutor)
        executor._result_mqs = result_mqs
        executor._result_mq = result_mqs[0]
        executor._pump_running = False
        executor._pump_stop = threading.Event()
        executor._sync_result_buffer = queue.Queue()
        executor._is_failed = False
        executor._futures_lock = threading.RLock()
        executor._rpc_futures = {}
        executor._output_futures = {}
        executor._completed_outputs = {}
        executor._batch_split_map = {}
        executor._start_result_pump()
        deadline = time.monotonic() + 10.0
        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
            time.sleep(0.01)
        assert executor._sync_result_buffer.qsize() == 2
        outputs = []
        for _ in range(2):
            output = executor._sync_result_buffer.get_nowait()
            unpack_diffusion_output_shm(output)
            outputs.append(output)
        assert sorted(output.output["rank"] for output in outputs) == [0, 1]
        for output in outputs:
            rank = output.output["rank"]
            assert output.output["nested"]["frames"][0] == rank
            assert output.output["audio"][0][0] == rank
        executor._pump_stop.set()
        for thread in executor._result_pump_threads:
            thread.join(timeout=3.0)
        assert all(not thread.is_alive() for thread in executor._result_pump_threads)
    finally:
        if executor is not None:
            executor._pump_stop.set()
            for thread in getattr(executor, "_result_pump_threads", []):
                thread.join(timeout=3.0)
        for connection in parent_connections:
            with contextlib.suppress(BrokenPipeError, EOFError, OSError):
                connection.send("shutdown")
        for process in processes:
            process.join(timeout=10.0)
            if process.is_alive():
                process.terminate()
                process.join(timeout=5.0)
        for result_mq in result_mqs:
            with contextlib.suppress(Exception):
                result_mq.shutdown()
        for connection in parent_connections:
            connection.close()
    assert all(process.exitcode == 0 for process in processes)
    for name in shm_names:
        try:
            leaked = shared_memory.SharedMemory(name=name)
        except FileNotFoundError:
            continue
        leaked.close()
        pytest.fail(f"shared memory segment {name} still exists after unpack")
 def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
    image = torch.arange(300_000, dtype=torch.float32)
    video = torch.arange(300_000, dtype=torch.float32) * 2
@@ -214,29 +83,6 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
    torch.testing.assert_close(output.output[1], frames[1])
 def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
    frames = np.arange(300_000, dtype=np.float32)
    output = DiffusionOutput(output=frames)
    pack_diffusion_output_shm(output)
    assert output.output["__ndarray_shm__"] is True
    unpack_diffusion_output_shm(output)
    assert isinstance(output.output, np.ndarray)
    np.testing.assert_array_equal(output.output, frames)
 def test_numpy_object_array_stays_inline() -> None:
    values = np.empty(200_000, dtype=object)
    values[:] = "safe-to-pickle"
    packed = _pack_value_if_large(values)
    assert packed is values
 def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
    tensor = torch.arange(300_000, dtype=torch.float32)
    envelope = {
@@ -262,25 +108,6 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
    assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
 def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
    frames = np.arange(300_000, dtype=np.float32)
    envelope = {
        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
        "result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
        "rank_statuses": [{"rank": 1, "ok": True}],
    }
    packed = pack_diffusion_output_shm(envelope)
    tagged = packed["result"]
    assert tagged["dp_rank"] == 1
    assert tagged["output"].output["__ndarray_shm__"] is True
    unpacked = unpack_diffusion_output_shm(packed)
    tagged = unpacked["result"]
    assert tagged["dp_rank"] == 1
    np.testing.assert_array_equal(tagged["output"].output, frames)
 def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
    first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
    second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1
@@ -162,21 +162,6 @@ def _fake_platform_for_peak_memory():
    )
 def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
    events = []
    plain_pipeline = object()
    request_scoped_pipeline = SimpleNamespace(
        adopt_cache_dit_backend=lambda backend: events.append(backend),
        is_cache_dit_enabled=lambda: True,
    )
    backend = object()
    assert not model_runner_module.adopt_request_scoped_cache_dit(plain_pipeline, backend)
    assert model_runner_module.adopt_request_scoped_cache_dit(request_scoped_pipeline, backend)
    assert model_runner_module.is_request_scoped_cache_dit_enabled(request_scoped_pipeline)
    assert events == [backend]
 def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
    runner = object.__new__(DiffusionModelRunner)
    runner.vllm_config = object()
@@ -218,25 +203,6 @@ def _make_compile_runner(
    return runner
@pytest.mark.core_model
@pytest.mark.cpu
 def test_update_states_carries_prepared_layout() -> None:
    runner = _make_runner(cache_backend=None, cache_backend_name=None)
    request = _make_request()
    prepared_layout = object()
    request.prepared_layout = prepared_layout
    scheduler_output = SimpleNamespace(
        finished_req_ids=set(),
        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request)],
        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
    )
    states, new_request_ids = DiffusionModelRunner._update_states(runner, scheduler_output)
    assert new_request_ids == [request.request_id]
    assert states[0].prepared_layout is prepared_layout
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("use_hsdp", [False, True])
@@ -453,33 +419,6 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
    assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)]
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("cache_dit_enabled", [False, True])
 def test_execute_model_cache_summary_follows_pipeline_owned_cache_dit_state(monkeypatch, cache_dit_enabled):
    runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
    runner.pipeline.adopt_cache_dit_backend = lambda backend: None
    runner.pipeline.is_cache_dit_enabled = lambda: cache_dit_enabled
    req = _make_request()
    cache_summary_calls = []
    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
    monkeypatch.setattr(
        model_runner_module,
        "cache_summary",
        lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
    )
    output = DiffusionModelRunner.execute_model(runner, req)
    assert output.output == "ok"
    assert cache_summary_calls == ([(runner.pipeline, True)] if cache_dit_enabled else [])
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):
@@ -12,8 +12,6 @@ from pytest_mock import MockerFixture
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
    BaseScheduler,
@@ -244,86 +242,7 @@ class TestGetRequestBatchSamplingParamsKey:
 class TestRequestScheduler:
    def setup_method(self) -> None:
        self.scheduler: RequestScheduler = RequestScheduler()
-        self.scheduler.initialize(SimpleNamespace(request_batch_max_wait_ms=0.0))
+        self.scheduler.initialize(SimpleNamespace())
    def test_admission_wait_disabled_with_zero_max_wait(self) -> None:
        self.scheduler.initialize(SimpleNamespace(request_batch_max_wait_ms=0.0))
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        assert decision.should_wait is False
    @pytest.mark.parametrize(
        ("dp_concurrent", "expected_stable_window_s"),
        [(False, 0.05), (True, 0.3)],
    )
    def test_admission_wait_decision_encodes_coalescing_policy(
        self,
        dp_concurrent: bool,
        expected_stable_window_s: float,
    ) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=4,
                request_batch_max_wait_ms=1000.0,
            )
        )
        decision = self.scheduler.get_admission_wait_decision(
            now=10.0,
            dp_concurrent=dp_concurrent,
        )
        assert decision.should_wait is True
        assert decision.deadline == 11.0
        assert decision.stable_window_s == expected_stable_window_s
        assert decision.max_batch == 4
    def test_admission_wait_disabled_while_wave_is_running(self) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=1,
                request_batch_max_wait_ms=1000.0,
            )
        )
        self.scheduler.add_request(_make_request("running"))
        self.scheduler.schedule()
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        assert decision.should_wait is False
    def test_admission_wait_end_conditions(self) -> None:
        self.scheduler.initialize(
            SimpleNamespace(
                max_num_seqs=2,
                request_batch_max_wait_ms=1000.0,
            )
        )
        decision = self.scheduler.get_admission_wait_decision(now=10.0)
        self.scheduler.add_request(_make_request("a"))
        assert not self.scheduler.should_end_admission_wait(
            decision,
            now=10.01,
            stable_since=10.0,
        )
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=10.05,
            stable_since=10.0,
        )
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=11.0,
            stable_since=11.0,
        )
        self.scheduler.add_request(_make_request("b"))
        assert self.scheduler.should_end_admission_wait(
            decision,
            now=10.01,
            stable_since=10.01,
        )
    def test_single_request_success_lifecycle(self) -> None:
        req_id = self.scheduler.add_request(_make_request("a"))
@@ -340,90 +259,6 @@ class TestRequestScheduler:
        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
        assert self.scheduler.has_requests() is False
    def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        prepared_layout = object()
        request.prepared_layout = prepared_layout
        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
        kv_request = DiffusionKVRequest(
            "diffusion-kv/diffusion-kv/0",
            sequence_id=0,
            prefix_len=4,
            target_len=8,
            seq_len=16,
            kv_contexts=(text_context,),
        )
        request.diffusion_kv_requests = (kv_request,)
        self.scheduler.add_request(request)
        state = self.scheduler.get_request_state(request.request_id)
        scheduler_output = self.scheduler.schedule()
        assert state is not None
        assert state.diffusion_kv_requests == (kv_request,)
        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
        assert request.diffusion_kv_requests is None
        assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
    def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
            self.scheduler.add_request(request)
    @pytest.mark.parametrize(
        ("owner_name", "field_name"),
        [
            ("request", "past_key_values"),
            ("sampling_params", "past_key_values"),
            ("sampling_params", "cfg_text_past_key_values"),
            ("sampling_params", "cfg_img_past_key_values"),
            ("sampling_params", "cfg_branch_past_key_values"),
        ],
    )
    def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
        request = _make_request("diffusion-kv")
        owner = request if owner_name == "request" else request.sampling_params
        setattr(owner, field_name, object())
        request.diffusion_kv_requests = (
            DiffusionKVRequest(
                "diffusion-kv/diffusion-kv/0",
                sequence_id=0,
                prefix_len=4,
                target_len=8,
                seq_len=16,
            ),
        )
        with pytest.raises(ValueError, match=rf"{owner_name}\.{field_name}"):
            self.scheduler.add_request(request)
    def test_dense_scheduler_rejects_scheduler_only_kv_request(self) -> None:
        request = _make_request("dense-kv-request")
        request.diffusion_kv_requests = (
            DiffusionKVRequest(
                "dense-kv-request/diffusion-kv/0",
                sequence_id=0,
                prefix_len=4,
                target_len=8,
                seq_len=16,
            ),
        )
        with pytest.raises(ValueError, match="dense_legacy request unexpectedly contains"):
            self.scheduler.add_request(request)
    def test_dense_scheduler_accepts_legacy_dense_kv(self) -> None:
        request = _make_request("dense-kv")
        request.sampling_params.past_key_values = object()
        request_id = self.scheduler.add_request(request)
        assert request_id == request.request_id
    def test_error_output_marks_finished_error(self) -> None:
        req_id = self.scheduler.add_request(_make_request("err"))
@@ -545,63 +380,6 @@ class TestRequestScheduler:
        assert first.num_running_reqs == 1
        assert first.num_waiting_reqs == 1
    def test_batches_different_quality_levels_separately(self) -> None:
        scheduler = RequestScheduler()
        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
        high = scheduler.add_request(
            _make_step_request(
                "high",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=2,
                    quality="high",
                ),
            )
        )
        scheduler.add_request(
            _make_step_request(
                "lossless",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=2,
                    quality="lossless",
                ),
            )
        )
        first = scheduler.schedule()
        assert _new_ids(first) == [high]
        assert first.num_running_reqs == 1
        assert first.num_waiting_reqs == 1
    def test_batches_omitted_and_explicit_lossless_separately(self) -> None:
        scheduler = RequestScheduler()
        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
        omitted = scheduler.add_request(
            _make_step_request(
                "omitted",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=2,
                ),
            )
        )
        scheduler.add_request(
            _make_step_request(
                "explicit",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=2,
                    quality="lossless",
                ),
            )
        )
        first = scheduler.schedule()
        assert _new_ids(first) == [omitted]
        assert first.num_running_reqs == 1
        assert first.num_waiting_reqs == 1
    def test_batches_different_request_local_seed_together(self) -> None:
        scheduler = RequestScheduler()
        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -727,20 +505,12 @@ class TestDiffusionEngine:
        engine.abort_queue = queue.Queue()
        request = _make_request("engine")
        prepared_layout = object()
        def preprocess(req):
            req.prepared_layout = prepared_layout
            return req
        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
        runner_output = _make_request_output("engine")
        engine.execute_fn = mocker.Mock(return_value=runner_output)
        output = engine.add_req_and_wait_for_response(request)
        assert output is runner_output.result
        engine.pre_process_func.assert_called_once_with(request)
        engine.execute_fn.assert_called_once()
    def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None:
@@ -804,29 +574,17 @@ class TestDiffusionEngine:
    async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine._check_and_start_background_loop = mocker.AsyncMock()
-        request = _make_request("req-abort")
+        engine.pre_process_func = None
        prepared_layout = object()
        def preprocess(req):
            req.prepared_layout = prepared_layout
            return req
        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
        async def _stream(_request):
            yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
-        engine._add_prepared_request = mocker.Mock(return_value=request.request_id)
+        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
        engine.get_output_stream = mocker.Mock(return_value=_stream(None))
        with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            async for _ in engine.step_streaming(request):
+            async for _ in engine.step_streaming(_make_request("req-abort")):
                pass
        engine.pre_process_func.assert_called_once_with(request)
        engine._add_prepared_request.assert_called_once_with(request)
        engine.get_output_stream.assert_called_once_with(request.request_id)
    def test_abort_queue_marks_request_finished_aborted(self) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine._rpc_lock = threading.RLock()
@@ -959,32 +717,12 @@ class TestDiffusionEngine:
        with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
            engine._dummy_run()
    def test_dummy_run_delegates_preprocessing_to_sync_admission(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
        engine.od_config = SimpleNamespace(model_class_name="mock_model", diffusion_load_format="default")
        engine.pre_process_func = mocker.Mock()
        engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(output=None))
        engine._dummy_run()
        engine.pre_process_func.assert_not_called()
        admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
        assert admitted_request.request_id == "dummy_req_id"
 class TestStepScheduler:
    def setup_method(self) -> None:
        self.scheduler: StepScheduler = StepScheduler()
        self.scheduler.initialize(SimpleNamespace())
    def test_admission_wait_is_not_supported(self) -> None:
        decision = self.scheduler.get_admission_wait_decision(
            now=10.0,
            dp_concurrent=True,
        )
        assert decision.should_wait is False
    def test_single_request_step_lifecycle(self) -> None:
        request = _make_step_request("step", num_inference_steps=3)
        req_id = self.scheduler.add_request(request)
@@ -1192,35 +930,6 @@ class TestStepScheduler:
        assert second.num_running_reqs == 1
        assert second.num_waiting_reqs == 1
    def test_step_batch_rejects_different_quality_levels(self) -> None:
        scheduler = StepScheduler()
        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
        high = scheduler.add_request(
            _make_step_request(
                "high",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=4,
                    quality="high",
                ),
            )
        )
        scheduler.add_request(
            _make_step_request(
                "lossless",
                sampling_params=OmniDiffusionSamplingParams(
                    num_inference_steps=4,
                    quality="lossless",
                ),
            )
        )
        first = scheduler.schedule()
        assert _new_ids(first) == [high]
        assert first.num_running_reqs == 1
        assert first.num_waiting_reqs == 1
    def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
        """Multiple requests with the same LoRA (id + scale) co-batch."""
        from vllm_omni.lora.request import LoRARequest
@@ -75,32 +75,6 @@ async def test_inline_dispatch_request_success(client, mock_engine):
    assert output.request_id == "req-1"
@pytest.mark.asyncio
 async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
    requests = []
    async def _step_streaming(request):
        requests.append(request)
        yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
    mock_engine.step_streaming = _step_streaming
    shared = OmniDiffusionSamplingParams()
    await client.add_request_async("req-1", "First prompt", shared)
    await client.add_request_async("req-2", "Second prompt", shared)
    for _ in range(10):
        if len(requests) == 2:
            break
        await asyncio.sleep(0.01)
    assert len(requests) == 2
    assert requests[0].sampling_params is not requests[1].sampling_params
    assert requests[0].sampling_params.guidance_scale_provided is False
    assert requests[1].sampling_params.guidance_scale_provided is False
    assert shared.guidance_scale is None
@pytest.mark.asyncio
 async def test_inline_dispatch_request_streaming_success(client, mock_engine):
    """Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
@@ -163,20 +137,6 @@ def test_inline_shutdown(client, mock_engine):
    mock_engine.close.assert_called_once()
 def test_inline_shutdown_runs_after_orchestrator_premark(client, mock_engine):
    # The orchestrator pre-marks clients to suppress false worker-death errors
    # before StagePool invokes the actual teardown.
    client._shutting_down = True
    client.shutdown()
    mock_engine.close.assert_called_once()
    assert client._shutdown_complete
    client.shutdown()
    mock_engine.close.assert_called_once()
 def test_inline_registers_executor_failure_callback(client, mock_engine):
    mock_engine.executor.register_failure_callback.assert_called_once()
@@ -137,18 +137,6 @@ class TestPackDiffusionOutputShmWithD2hStream:
            d2h_stream=d2h,
        )
    @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
    def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope):
        mock_pack_fields = mocker.patch(
            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
        )
        d2h = mocker.MagicMock()
        result = DiffusionOutput()
        tagged = {"dp_rank": 1, "output": result}
        value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
        pack_diffusion_output_shm(value, d2h_stream=d2h)
        mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h)
    def test_runner_output_path_with_stream(self, mocker):
        mock_pack_fields = mocker.patch(
            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
@@ -48,7 +48,6 @@ def _mock_request(tag: str):
        request_id=tag,
        prompt=f"prompt_{tag}",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        diffusion_kv_requests=None,
    )
@@ -277,8 +276,6 @@ class TestRequestModeDispatch:
        def worker():
            request = req_q.get(timeout=2)
            assert "rpc_id" not in request
            assert request["output_rank"] is None
            assert request["collect_rank_status"] is False
            wave_id = request["wave_id"]
            for dp_rank in range(2):
                res_q.put(
@@ -300,181 +297,6 @@ class TestRequestModeDispatch:
        assert [result.error for result in results] == ["0", "1"]
    @pytest.mark.parametrize(
        ("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
        [
            (2, 1, 1, 1, [0, 2]),
            (2, 2, 1, 1, [0, 4]),
            (1, 2, 2, 1, [0, 4]),
            (1, 1, 2, 2, [0, 4]),
        ],
    )
    def test_step_execution_reads_each_dp_primary_result_queue(
        self,
        tp_size,
        sp_size,
        pp_size,
        cfg_size,
        expected_primary_ranks,
    ):
        dp_size = 2
        num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
        executor, _, _ = _make_executor(num_gpus=num_gpus)
        executor.od_config = SimpleNamespace(
            step_execution=True,
            parallel_config=SimpleNamespace(
                data_parallel_size=dp_size,
                tensor_parallel_size=tp_size,
                sequence_parallel_size=sp_size,
                pipeline_parallel_size=pp_size,
                cfg_parallel_size=cfg_size,
            ),
        )
        request: dict = {}
        executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
        def result_queue(global_rank):
            result_mq = MagicMock()
            result_mq.dequeue.side_effect = lambda timeout=None: {
                "dp_rank": global_rank,
                "output": _tagged_output(str(global_rank)),
                "wave_id": request["wave_id"],
            }
            return result_mq
        # Every queue returns a valid, rank-tagged result. If the executor
        # selects a non-primary queue, the returned rank exposes the mistake.
        executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
        executor._result_mq = executor._result_mqs[0]
        results = executor.collective_rpc(
            "execute_model",
            unique_reply_rank=None,
            exec_all_ranks=True,
        )
        assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
        for rank, result_mq in enumerate(executor._result_mqs):
            assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
    @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
    def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
        executor, _, _ = _make_executor(num_gpus=2)
        executor.od_config = SimpleNamespace(
            step_execution=False,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        executor.collective_rpc = Mock()
        scheduler_output = _make_sched_output("invalid", "valid")
        scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
        with pytest.raises(ValueError, match="non-empty prompt"):
            executor.execute_request(scheduler_output)
        executor.collective_rpc.assert_not_called()
    def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
        executor, _, _ = _make_executor(num_gpus=2)
        executor.od_config = SimpleNamespace(
            step_execution=False,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
        invalid_wave = _make_sched_output("invalid", "valid")
        invalid_wave.scheduled_new_reqs[0].req.prompt = ""
        with pytest.raises(ValueError, match="non-empty prompt"):
            executor.execute_request(invalid_wave)
        result = executor.execute_request(_make_sched_output("A", "B"))
        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
        executor.collective_rpc.assert_called_once()
    def test_dlo_dp_allows_shared_default_denoise_steps(self):
        executor, _, _ = _make_executor(num_gpus=2)
        executor.od_config = SimpleNamespace(
            step_execution=False,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
        scheduler_output = _make_sched_output("A", "B")
        for new_req in scheduler_output.scheduled_new_reqs:
            new_req.req.sampling_params.num_inference_steps = None
        result = executor.execute_request(scheduler_output)
        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
        executor.collective_rpc.assert_called_once()
    @pytest.mark.parametrize(
        ("field", "value"),
        [
            ("num_inference_steps", 2),
            ("guidance_scale", 7.5),
            ("width", 1024),
        ],
    )
    def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
        executor, _, _ = _make_executor(num_gpus=2)
        executor.od_config = SimpleNamespace(
            step_execution=False,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        executor.collective_rpc = Mock()
        scheduler_output = _make_sched_output("A", "B")
        setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
        with pytest.raises(ValueError, match="compatible shape, CFG"):
            executor.execute_request(scheduler_output)
        executor.collective_rpc.assert_not_called()
    def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
        executor, req_q, res_q = _make_executor(num_gpus=2)
        executor.od_config = SimpleNamespace(
            step_execution=False,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        executor._sync_result_buffer = res_q
        executor._fail_closed_on_dp_wave_timeout = Mock()
        monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
        def reply_from_only_one_dp_rank():
            request = req_q.get(timeout=2.0)
            res_q.put(
                {
                    "dp_rank": 0,
                    "output": _tagged_output("rank-0"),
                    "wave_id": request["wave_id"],
                }
            )
        worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
        worker.start()
        started = time.monotonic()
        result = executor.execute_request(_make_sched_output("A", "B"))
        elapsed = time.monotonic() - started
        worker.join(timeout=2.0)
        assert elapsed < 1.0
        assert all("timed out" in output.result.error for output in result.runner_outputs)
        executor._fail_closed_on_dp_wave_timeout.assert_called_once()
        assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
 # ───────────────── concurrent collective RPC ─────────────────
@@ -1126,42 +948,6 @@ class TestStageDiffusionClientErrorPropagation:
        assert rpc_id not in client._pending_rpcs
 class TestExecutorShutdownCleaner:
    def test_worker_joins_share_one_global_deadline(self, monkeypatch):
        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
        class FakeProcess:
            def __init__(self, name):
                self.name = name
                self.alive = True
                self.terminated = False
                self.join_timeouts = []
            def is_alive(self):
                return self.alive
            def join(self, timeout):
                self.join_timeouts.append(timeout)
                if self.terminated:
                    self.alive = False
            def terminate(self):
                self.terminated = True
        monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
        first = FakeProcess("worker-0")
        second = FakeProcess("worker-1")
        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
        cleaner()
        assert first.join_timeouts == [15.0, 5.0]
        assert second.join_timeouts == [5.0, 1.0]
        assert first.terminated and second.terminated
        assert not first.is_alive() and not second.is_alive()
 # ───────── monitor thread & death sentinel integration tests ─────────
@@ -35,7 +35,6 @@ def _make_executor(step_execution=False):
    executor._pump_stop = threading.Event()
    executor._sync_result_buffer = queue.Queue()
    executor._result_mq = MagicMock()
    executor._result_mqs = []
    executor._broadcast_mq = MagicMock()
    executor._closed = False
    executor._is_failed = False
@@ -158,41 +157,6 @@ class TestResultPumpDispatch:
        retrieved = executor._sync_result_buffer.get_nowait()
        assert isinstance(retrieved, DiffusionOutput)
    def test_start_result_pump_reads_every_worker_queue(self):
        executor = _make_executor()
        messages = [DiffusionOutput(output="rank0"), DiffusionOutput(output="rank1")]
        result_mqs = [MagicMock(), MagicMock()]
        def make_dequeue(message):
            emitted = False
            def dequeue(timeout=None):
                nonlocal emitted
                if not emitted:
                    emitted = True
                    return message
                executor._pump_stop.wait(0.01)
                raise TimeoutError
            return dequeue
        for result_mq, message in zip(result_mqs, messages, strict=True):
            result_mq.dequeue = make_dequeue(message)
        executor._result_mqs = result_mqs
        executor._result_mq = result_mqs[0]
        executor._start_result_pump()
        deadline = time.monotonic() + 2.0
        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
            time.sleep(0.01)
        executor._pump_stop.set()
        for thread in executor._result_pump_threads:
            thread.join(timeout=2.0)
        assert len(executor._result_pump_threads) == 2
        received = [executor._sync_result_buffer.get_nowait().output for _ in range(2)]
        assert sorted(received) == ["rank0", "rank1"]
    def test_compute_done_routes_to_rpc_future(self):
        executor = _make_executor()
        rpc_id = "42"
@@ -269,20 +233,6 @@ class TestResultPumpDispatch:
 class TestShutdownCleansUpFutures:
    """Test that shutdown cancels pending futures."""
    def test_shutdown_joins_result_pump_threads(self):
        executor = _make_executor()
        pump = threading.Thread(
            target=executor._pump_stop.wait,
            name="test-result-pump",
        )
        executor._result_pump_threads = [pump]
        pump.start()
        executor.shutdown()
        assert not pump.is_alive()
        assert executor._result_pump_threads == []
    def test_shutdown_sets_exception_on_pending_futures(self):
        executor = _make_executor()
Author	SHA1	Message	Date
Hongsheng Liu	01e7d09119	docs: define module ownership contracts Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-07 11:45:59 +08:00
hsliuustc0106	69766960f2	docs: update module design navigation Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-07-16 14:09:07 +08:00
hsliuustc0106	80fd96c1da	docs: reorganize module design documents Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-07-16 12:08:45 +08:00