Revert "[Bugfix] Use MediaConnector for image/video URL fetching to prevent S…"

This reverts commit d8edf0c9d5.

.buildkite/pipeline.yml

@@ -26,23 +26,33 @@ steps:
     agents:
       queue: "cpu_queue_premerge"
 
-  # L2 Test
+  # L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
     if: __UPLOAD_READY_IF__
     commands:
-      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
+      - |
+        EXTRA=""
+        if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
+          EXTRA="--e2e"
+        fi
+        python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-ready.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  # L3 Test
+  # L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
     if: __UPLOAD_MERGE_IF__
     commands:
-      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
+      - |
+        EXTRA=""
+        if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
+          EXTRA="--e2e"
+        fi
+        python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 

.buildkite/scripts/upload_pipeline.py

@@ -13,7 +13,7 @@ Test pipeline mode (e.g. test-merge.yml):
     recognize this uploader-only key).
 
 Usage:
-  python3 upload_pipeline.py [--upload] [--all] <pipeline.yml>
+  python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml>
 
   # Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
   python3 upload_pipeline.py --upload .buildkite/pipeline.yml
@@ -23,6 +23,9 @@ Usage:
 
   # Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
   python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
+
+  # Nightly L2 E2E only: keep the E2E Test group, ignore source_file_dependencies:
+  python3 upload_pipeline.py --upload --e2e .buildkite/test-ready.yml
 """
 
 from __future__ import annotations
@@ -47,6 +50,7 @@ LOG = "upload_pipeline"
 ROOT = Path(__file__).resolve().parent.parent.parent
 DOC_SEP = "\n---\n"
 BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
+E2E_GROUP_MARKER = "E2E Test"
 
 
 def _log(message: str) -> None:
@@ -322,18 +326,18 @@ def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
     nightly_only = (
         '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
     )
+    nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
+    ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"'
+    merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"'
+    merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
     if skip_ci:
         image_if = f"'{nightly_only}'"
-        ready_if = "'false'"
-        merge_if = "'false'"
+        ready_if = f"'({nightly_main})'"
+        merge_if = f"'({nightly_main})'"
     else:
         image_if = "'true'"
-        ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
-        merge_if = (
-            '\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
-            '&& build.env("WEEKLY") != "1") || '
-            '(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
-        )
+        ready_if = f"'(({nightly_main}) || ({ready_pr}))'"
+        merge_if = f"'(({nightly_main}) || (({merge_main}) || ({merge_pr})))'"
 
     return (
         continuation.replace("__IMAGE_BUILD_IF__", image_if)
@@ -396,6 +400,20 @@ def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
     return filtered
 
 
+def _is_e2e_group(step: dict[str, Any]) -> bool:
+    group = step.get("group")
+    return isinstance(group, str) and E2E_GROUP_MARKER in group
+
+
+def _select_e2e_group_steps(steps: list[Any]) -> list[Any]:
+    selected = [step for step in steps if isinstance(step, dict) and _is_e2e_group(step)]
+    if not selected:
+        _log(f"no group matching {E2E_GROUP_MARKER!r} found")
+    else:
+        _log(f"keep {len(selected)} group(s) matching {E2E_GROUP_MARKER!r}")
+    return selected
+
+
 def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
     """Remove uploader-only keys while keeping all steps (no diff filtering)."""
     stripped: list[Any] = []
@@ -422,11 +440,16 @@ def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
 def render_test_pipeline(
     doc: dict[str, Any],
     changed_files: list[str] | None,
+    *,
+    e2e_only: bool = False,
 ) -> dict[str, Any]:
     steps = doc.get("steps")
     if not isinstance(steps, list):
         return doc
-    if changed_files is not None:
+    if e2e_only:
+        steps = _select_e2e_group_steps(steps)
+        steps = _strip_uploader_metadata_from_steps(steps)
+    elif changed_files is not None:
         steps = _filter_steps(steps, changed_files)
     else:
         steps = _strip_uploader_metadata_from_steps(steps)
@@ -440,7 +463,12 @@ def resolve_pipeline_path(arg: str) -> Path:
     return ROOT / path
 
 
-def render_pipeline(path: Path, *, force_all: bool = False) -> str:
+def render_pipeline(
+    path: Path,
+    *,
+    force_all: bool = False,
+    e2e_only: bool = False,
+) -> str:
     text = path.read_text(encoding="utf-8")
     diff_range = resolve_diff_range()
     changed_files = resolve_changed_files()
@@ -448,7 +476,7 @@ def render_pipeline(path: Path, *, force_all: bool = False) -> str:
     # ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
     # stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
     # run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
-    if force_all:
+    if force_all or e2e_only:
         changed_files = None
 
     if BOOTSTRAP_MARKER in text:
@@ -462,7 +490,7 @@ def render_pipeline(path: Path, *, force_all: bool = False) -> str:
     if not isinstance(doc, dict):
         raise ValueError(f"invalid pipeline YAML: {path}")
 
-    doc = render_test_pipeline(doc, changed_files)
+    doc = render_test_pipeline(doc, changed_files, e2e_only=e2e_only)
 
     return yaml.safe_dump(doc, sort_keys=False)
 
@@ -489,11 +517,17 @@ def main() -> int:
         action="store_true",
         help="Pipe rendered YAML to buildkite-agent pipeline upload",
     )
-    parser.add_argument(
+    mode = parser.add_mutually_exclusive_group()
+    mode.add_argument(
         "--all",
         action="store_true",
         help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
     )
+    mode.add_argument(
+        "--e2e",
+        action="store_true",
+        help="Keep only the E2E Test group; disable diff-aware skipping for those steps",
+    )
     args = parser.parse_args()
 
     path = resolve_pipeline_path(args.pipeline)
@@ -501,7 +535,7 @@ def main() -> int:
         _log(f"missing pipeline file: {path}")
         return 1
 
-    rendered = render_pipeline(path, force_all=args.all)
+    rendered = render_pipeline(path, force_all=args.all, e2e_only=args.e2e)
     if args.upload:
         upload_to_buildkite(rendered)
     else:

.buildkite/test-weekly.yml

@@ -85,6 +85,66 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
+  - label: "Reliability Test - hunyuan-image"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 4
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Reliability Test - voxcpm2"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
   - group: "Reliability Test - Invalid parameters Test"
     depends_on: upload-weekly-pipeline

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

@@ -63,6 +63,8 @@ If multiple rows apply (e.g., a diffusion model is also a new model), union the
 
 Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
 
+**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
+
 ### Step 5: Print Report
 
 ```
@@ -74,6 +76,7 @@ Pre-check report for <branch>
   Dimension          Result
   ─────────────────  ──────
   PR title format    ✓
+  Code quality       ⚠ 1 broad except, 2 Any hints
   PR desc integrity  ✓
   Registry/config    ✓
   Dead code          ⚠ 2 warnings

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

@@ -35,6 +35,31 @@ PR title must follow the project convention documented in `docs/contributing/REA
 
 ---
 
+## All PRs: Code Quality (run on every PR, quick & full)
+
+Diff-scoped sweep for five fragility patterns that are pervasive in the existing codebase — **only lines the PR adds count**, not the pre-existing backlog. Detection commands, codebase examples, and full severity rules live in [code-quality.md](code-quality.md).
+
+- [ ] **No new `**kwargs` string-lookup plumbing:** the diff adds no `**kwargs` + `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`. ⚠ any new instance; ✗ if a new string key is duplicated across ≥2 files without a shared constant, or unknown keys are silently dropped on a fail-fast path.
+- [ ] **No new broad exception swallow:** the diff adds no `except Exception:` / bare `except:`. ⚠ any new instance; ✗ for `except: pass` / `except Exception: return None` / `continue` on a fail-fast path (init, config validation, weight loading, request handling).
+- [ ] **No new `Any` / wrong type hints:** the diff adds no `: Any` / `-> Any` / untyped production signature, and no new `SimpleNamespace` in tests faking an object that already has a real typed stub. ⚠ any new instance; ✗ for a *wrong* (actively misleading) annotation or a `SimpleNamespace` test fake of a typed object.
+- [ ] **No new hot-path copy:** the diff adds no `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path. ⚠ any new instance; ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without a comment, or `deepcopy` of a scheduler / sampling-config on the request path.
+- [ ] **No new event-loop blocking:** the diff adds no `time.sleep` / blocking HTTP (`requests`, `urllib`) / `await` inside a held lock in `async def` code. ⚠ any new instance; ✗ for blocking sleep or HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock held across an `await`. Blocking inside a dedicated worker thread is fine.
+
+Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
+
+### Conventions (eyeball — no grep)
+
+Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
+
+- **Right log level** — `info` is for end-user-visible events; routine internal events are `debug`; hot paths are `trace` or removed. Logging is not free.
+- **Structured fields over f-strings** — pass dynamic values as structured fields (`logger.info("msg", extra={...})`), not interpolated strings, so they stay greppable and parseable.
+- **No reflexive synchronization** — don't add `threading.Lock` / `asyncio.Lock` / `deepcopy` / re-wrap shared state unless there is real concurrent mutation; the GIL covers most Python code. Owners decide their own synchronization — don't pre-lock in a constructor.
+- **Context managers for cleanup** — use `with` / `contextlib` for files, locks, CUDA streams, and profilers; not raw `open()` + manual `.close()` or ad-hoc `finally`.
+- **No new dependencies without justification** — prefer the stdlib and existing vllm-omni deps; a new top-level dep needs a reason (check `pyproject.toml` / `requirements*.txt`).
+- **Naming / interface discipline** — don't rename or change interfaces just for taste in this PR, but do fix names that actively mislead (e.g. `serve` implying a long-running server, `Instance` where multiple instance kinds coexist).
+
+---
+
 ## Bug Fix PRs
 
 ### Quick
@@ -49,7 +74,7 @@ PR title must follow the project convention documented in `docs/contributing/REA
 - [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
 - [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
 - [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
-- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
+- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added. See the broad-except pattern in [code-quality.md](code-quality.md#2-broad-exception-swallow) for the full severity rule (⚠ any new broad catch; ✗ for a swallow on a fail-fast path).
 - [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
 - [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
 

.claude/skills/precheck-pr/references/code-quality.md

@@ -0,0 +1,179 @@
+# Code-Quality Patterns
+
+Five fragility patterns to sweep on **every** PR, in both quick and full mode. These are pervasive in the existing codebase (hundreds of pre-existing sites) — so the checks are **diff-scoped**: only lines the PR *adds* count. Whole-repo greps light up every PR with the backlog and are useless.
+
+All commands assume:
+```bash
+BASE="$(git merge-base HEAD origin/main)"
+```
+
+**Severity:** ⚠ for any *new* instance introduced by the diff; ✗ for the acute sub-cases defined per pattern. Acute ✗ findings belong in the report's blocking count.
+
+---
+
+## 1. `**kwargs` string-lookup passthrough
+
+### Anti-pattern
+```python
+def forward(self, **kwargs):
+    # BLOCKER-prone: raw dict + string keys duplicated across files
+    if "runtime_additional_information" in kwargs and \
+       "model_intermediate_buffer" not in kwargs:
+        ...
+```
+Real instances of this exact guard are copy-pasted across `cosyvoice3.py`, `indextts2_talker.py` (×2), `indextts2_s2mel_decoder.py`, `qwen3_tts_talker.py`, `qwen3_omni.py`. A typo in any string key silently no-ops; there is no type check, no IDE completion, and unknown keys are silently dropped.
+
+### Why dangerous
+Raw `dict` + `**kwargs` + lists of param-name strings is **fragile by construction**: the compiler can't see a missing or misspelled key, so failures are silent. Each duplicated copy is a new place for the strings to drift.
+
+### Severity
+- ⚠ any *new* `**kwargs` signature whose body does `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`.
+- ✗ when the diff introduces a **new** string key that is duplicated across ≥2 files without a shared module constant, or when `**kwargs` plumbing **silently drops unknown keys** on a fail-fast path (init, config validation, weight loading).
+
+### Detect (diff-scoped)
+```bash
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
+  | grep -E '\*\*kwargs|kwargs(\.get|\[)|"[a-z_]+"\s+in\s+kwargs'
+```
+
+### Fix
+Prefer explicit typed params. The repo already has the right primitive: `msgspec.Struct` with `forbid_unknown_fields=True` rejects unknown keys instead of silently dropping them — see `_StructBase` in `vllm_omni/data_entry_keys.py:105` and `MoriPullRequest` in `vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py:81`. A `@dataclass` or `TypedDict` works too.
+
+If `**kwargs` is unavoidable (vLLM base-class compat), centralize the key strings as module-level constants behind **one** typed accessor, instead of repeating the literal across files.
+
+---
+
+## 2. Broad exception swallow
+
+### Anti-pattern
+```python
+# BLOCKER-prone: catches everything, returns None on any failure
+try:
+    evt = build_event(...)
+    return evt
+except Exception:
+    return None
+```
+Real clusters: `vllm_omni/metrics/stats.py:240,274,367,450` and `vllm_omni/metrics/{modality,utils}.py` — all `except Exception: return None`.
+
+### Why dangerous
+A generic `except Exception:` is **too flexible**. If we catch an exception type we don't expect, it is almost always an unhandled edge case that needs to be *fixed* — not swallowed into `None`/`pass`/`continue`. Broad catches turn fixable bugs into silent wrong behavior and make debugging impossible.
+
+The repo backs this with ruff `BLE001` (broad-exception catch); it is currently suppressed via `# noqa: BLE001` in `vllm_omni/patch.py:413`. Prefer fixing over suppressing.
+
+### Severity
+- ⚠ any *new* `except Exception:` / `except BaseException:` / bare `except:`.
+- ✗ for `except:` / `except Exception: pass|return None|continue` on a **fail-fast path** (init, config validation, weight loading, request handling, connector setup). This extends the Bug-Fix "No silent failure risk" ✗ in [checklists.md](checklists.md).
+
+### Detect (diff-scoped)
+```bash
+# Any new broad catch (⚠ baseline):
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
+  | grep -E 'except\s*(Exception|BaseException)?\s*:'
+# Note: this matches only the immediate-colon swallow forms (`except :`,
+# `except Exception:`) — it deliberately excludes specific catches like
+# `except ValueError:` and the `except Exception as e:` variant. Eyeball
+# the latter; `as e` usually means it is logged, but still verify.
+```
+
+### Fix
+Catch the **specific** types you actually expect (`ValueError`, `KeyError`, `AttributeError`, `OSError`, `torch.cuda.OutOfMemoryError`, `msgspec.ValidationError`, …). At a genuine top-level / best-effort boundary (metrics, signal handlers), at minimum **log** the exception; never swallow into a bare `pass`/`return None` on a path that should fail loudly.
+
+---
+
+## 3. `Any` / wrong type annotations
+
+### Anti-pattern
+```python
+# ⚠ Any param + Any return — no contract, no narrowing
+def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]:
+    ...
+```
+
+### Why dangerous (root cause)
+`Any` disables the type checker for that position, so typos and wrong shapes pass silently. The repo has 636 `: Any` params, 244 `-> Any` returns, and **1,173 `SimpleNamespace` in `tests/`**. The `Any` leak is largely test-driven: vibe-coded unit tests fake objects with `SimpleNamespace`, which forces the production code consuming them to be typed `Any`, and the `Any` then propagates outward into public signatures.
+
+### Severity
+- ⚠ any *new* `: Any` / `-> Any` / fully-untyped signature in production code.
+- ✗ for a **wrong** annotation (actively misleading — e.g. `-> bool` that returns `Optional[bool]`), or a *new* `SimpleNamespace` in a test that mimics an object which already has a real typed stub / `@dataclass` / `TypedDict` / `Protocol`.
+
+### Detect (diff-scoped)
+```bash
+# Any leaking into new signatures:
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
+  | grep -E ':\s*Any\b|->\s*Any\b'
+# SimpleNamespace newly added in tests (the Any-leak source):
+git diff ${BASE}...HEAD -- "tests/" | grep "^+[^+]" | grep "SimpleNamespace"
+```
+
+### Fix
+Replace `Any` with the concrete type, a `Protocol`, or a `Union`. If a type is genuinely dynamic, prefer `object` + `isinstance` narrowing over `Any` — `object` keeps the checker engaged. In tests, replace `SimpleNamespace` with the real class, a small `@dataclass`/`TypedDict`, or a `Protocol` + fake, so the system under test stays fully typed.
+
+---
+
+## 4. Unnecessary copy — `.clone()` / `copy.deepcopy` in hot paths
+
+### Anti-pattern
+```python
+# Per-step in a diffusion/AR loop: clones the full latent every iteration
+for t in timesteps:
+    trajectory_latents.append(x_t.clone())   # O(N) GPU alloc + memcpy each step
+```
+Real clusters: `vllm_omni/diffusion/models/bagel/bagel_transformer.py` calls `x_t.clone()` ~9× across a trajectory loop (lines 1820–2186); the cache hooks `vllm_omni/diffusion/cache/teacache/hook.py:155` and `magcache/hook.py:182` clone `hidden_states` per step. Separately, `copy.deepcopy(self.scheduler)` / `copy.deepcopy(req.sampling_params)` is copy-pasted across diffusion pipelines — `pipeline_qwen_image.py:786`, `pipeline_hunyuan_image3.py:1904`, `pipeline_ltx2*.py` (×3), `pipeline_dreamzero.py:1085`, `diffusion_model_runner.py:398` — a full object-graph deepcopy on the request path. (Contrast `dreamzero/causal_wan_model.py:203`, which carries a `# No .clone() needed.` comment — the aware pattern.)
+
+### Why dangerous
+A `.clone()` of a GPU tensor or a `copy.deepcopy` of a config object inside a per-step / per-request path is a hidden allocator: it reads as one Python line but triggers a GPU memcpy or a deep object-graph walk every iteration. Across 50 diffusion steps or hundreds of AR steps it dominates latency and VRAM bandwidth. `deepcopy` additionally ignores any custom device / `__copy__` logic and can drag tensors through CPU.
+
+### Severity
+- ⚠ any *new* `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path.
+- ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without an explanatory comment, or `copy.deepcopy` of a scheduler / sampling-config on the request hot path.
+
+### Detect (diff-scoped)
+```bash
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
+  | grep -E '\.clone\(\)|\.copy_\(|copy\.deepcopy|deepcopy\('
+```
+
+### Fix
+Tensors: prefer a view (`x[..., :k]`, `torch.narrow`, `as_strided`) or write into a pre-allocated buffer with `dst.copy_(src)` instead of `.clone()`; if the caller no longer needs the original, `move`/`del` it rather than copy. Schedulers / configs: build a fresh lightweight instance per request, or split out the mutable per-request bits so the shared immutable part isn't deepcopied. If a copy is genuinely required (e.g. the source is mutated in place later), keep it and add a one-line comment saying why.
+
+---
+
+## 5. Blocking the asyncio event loop / lock held across `await`
+
+### Anti-pattern
+```python
+# BLOCKER-prone: blocking sleep / blocking HTTP / lock held across an await,
+# all inside an async (event-loop) function
+async def handle(self, req):
+    with self._lock:
+        resp = await self._client.call(req)   # lock held across await -> serializes all requests
+    time.sleep(0.1)                           # stalls the single loop for every concurrent request
+```
+
+### Why dangerous
+vLLM-Omni serves over a single asyncio event loop. Any synchronous blocking work on an `async def` path — `time.sleep`, blocking HTTP (`requests`/`urllib`), heavy CPU work, or a lock held across an `await` — stalls the loop for *every* concurrent request, not just the blocking one. A lock held across `await` silently serializes the whole pipeline; blocking I/O tanks throughput and inflates tail latency.
+
+### Severity
+- ⚠ any *new* `time.sleep(`, blocking HTTP call (`requests.get/post`, `urllib.request`, `urlopen`), or `await` inside a `with Lock` / `async with` critical section, in `async def` code.
+- ✗ for `time.sleep` or blocking HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock acquired and held across an `await` on the request path.
+
+### Detect (diff-scoped)
+```bash
+# New blocking primitives introduced by the diff — then eyeball whether each
+# sits in an async function on the loop vs a dedicated worker thread:
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
+  | grep -E 'time\.sleep|requests\.(get|post|put|delete)|urllib\.request|urlopen'
+# Lock acquired with an await potentially inside the critical section (manual eyeball):
+git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" | grep -E 'async with|\.acquire\(\)'
+```
+Note: blocking inside a **dedicated worker thread is correct** — the repo already does this (`diffusion/model_loader/hub_prefetch.py`, `distributed/omni_connectors/kv_transfer_manager.py`, the `mooncake`/`mori` connectors poll with `time.sleep` on their own threads). The check targets the *event-loop* path only.
+
+### Fix
+Move blocking work off the loop with `asyncio.to_thread(...)` / `run_in_executor`, or use a native async client (`httpx.AsyncClient` / `aiohttp`). Never `await` while holding a lock — drop the lock before the `await` and re-acquire after, or restructure so the critical section contains no `await`.
+
+---
+
+## Running the sweep
+
+Run all five detections against `${BASE}...HEAD`, then for each hit decide ⚠ vs ✗ using the severity rules above. Roll the results into the report as a single **Code quality** dimension row (count of ⚠ and ✗), alongside the type-specific checklist from [checklists.md](checklists.md).

apps/lvsa-omni-showcase/README.md

@@ -0,0 +1,87 @@
+# LVSA Long-Video Showcase
+
+**Training-free block-sparse attention for long-video diffusion**, running on
+vLLM-Omni via the public plugin entry points (no core changes). LVSA speeds up
+extended-horizon generation and prevents the freeze/loop failure mode dense
+attention exhibits beyond a model's training length — with no fine-tuning and no
+weight changes.
+
+Upstream: **[LongVideoSparseAttention](https://github.com/JiusiServe/LongVideoSparseAttention)**
+(algorithm, all model adapters, full docs).
+
+## Results
+
+Wan 2.1 1.3B, single A100 80 GB, mean over 5 prompts (Dense vs LVSA-FlashInfer):
+
+| Horizon | 1× | 2× | 3× | 4× | 5× | 6× |
+|---|---|---|---|---|---|---|
+| **Speedup** | 0.9× | 1.4× | 1.9× | 2.5× | 3.0× | **3.5×** |
+| **VQeval composite** (Dense→LVSA) | 63→63 | 59→63 | 61→63 | 61→64 | 59→64 | 59→63 |
+
+The speedup **grows with length** (LVSA is the dense regime at 1× by design, so
+the win starts at 2×), and quality *improves* at extension — LVSA's rotating
+keyframes prevent the looping/static output dense produces beyond its horizon
+(VQeval loop-quality +27 to +34 at ≥3×). At long horizons dense **OOMs on 80 GB**
+(HunyuanVideo ≥2×, Wan2.1-14B and Cosmos at 6×) where LVSA still generates.
+
+→ Full sweep (5 models × 6 horizons, VQeval + VBench-Long) and the
+quality-metric methodology:
+**[benchmarks/](https://github.com/JiusiServe/LongVideoSparseAttention/tree/main/benchmarks)**.
+
+## Quickstart
+
+```bash
+# 1. vLLM-Omni 0.22.0 (stable) + matching vLLM
+pip install "vllm==0.22.0"
+pip install --no-build-isolation \
+  "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@v0.22.0"
+
+# 2. LVSA core + plugin
+git clone https://github.com/JiusiServe/LongVideoSparseAttention.git
+( cd LongVideoSparseAttention && pip install -e . && pip install -e lvsa-vllm-omni/ )
+
+# 3. this showcase's client dep
+pip install requests
+```
+
+FlashInfer (used by the default `flashinfer` backend) **ships with vLLM** — step 1
+already pulls the pinned `flashinfer-python` + `flashinfer-cubin`, so the fused
+backend works with no extra install. Two optional extras:
+
+```bash
+# (optional) faster startup — prebuilt JIT cache, matched to your CUDA + flashinfer
+CU=cu128   # match your CUDA: cu121 / cu124 / cu128 / …
+FIVER=$(python -c "import importlib.metadata as m; print(m.version('flashinfer-python'))")
+pip install --extra-index-url "https://flashinfer.ai/whl/${CU}" \
+  "flashinfer-jit-cache==${FIVER}+${CU}"
+
+# (fallback) no FlashInfer in your environment? run the SDPA backend instead:
+export LVSA_BACKEND=sdpa
+```
+
+```bash
+# Serve an LVSA-enabled endpoint (Wan 2.1 1.3B at 2× horizon):
+MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 \
+  bash serve_lvsa.sh
+
+# In another shell — generate:
+python generate.py --frames 161 --height 480 --width 832 \
+  --prompt "A dog running through a sunlit forest." --out wan_lvsa.mp4
+```
+
+`serve_lvsa.sh` sets the LVSA env vars for the chosen `MODEL_FAMILY`
+(`wan` / `hunyuan`). Wan and HunyuanVideo run through the LVSA **attention
+backend** by default — it stays sparse even under sequence-parallel, unlike the
+monkey-patch hooks.
+
+## Configuration
+
+The one knob you must get right is the model's training horizon
+(`LVSA_REFERENCE_LATENT_FRAMES`: Wan = 21, HunyuanVideo = 33); `serve_lvsa.sh`
+sets it per family. For the full environment-variable reference, the multi-GPU
+support matrix, and tuning (`sparsity_scale`, rotating keyframes, non-standard
+resolutions), see the upstream docs:
+
+- **[lvsa-vllm-omni plugin reference](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/lvsa-vllm-omni/README.md)** — every `LVSA_*` env var
+- **[docs/parallelism.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/parallelism.md)** — TP / Ulysses / HSDP support per model
+- **[docs/tuning.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/tuning.md)** — picking knobs for your model and horizon

apps/lvsa-omni-showcase/generate.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+"""Generate a long video through an LVSA-enabled vLLM-Omni server.
+
+A thin client over the Omni video API — start the server first with
+``serve_lvsa.sh``. The sparse attention runs server-side; this just submits a
+job and saves the result.
+
+Uses the **asynchronous** endpoint (``POST /v1/videos`` → poll
+``GET /v1/videos/{id}`` → download ``GET /v1/videos/{id}/content``) rather than
+``POST /v1/videos/sync``. The sync endpoint is capped by a fixed server-side
+timeout (``VIDEO_SYNC_TIMEOUT_S``, 600 s) and is documented as
+"for benchmark and testing scenarios" — long videos (LVSA's use case) routinely
+exceed it. The async job has no such cap.
+
+  python generate.py --frames 481 --height 480 --width 832 \
+      --prompt "A golden retriever running through a sunlit forest." \
+      --out dog_forest_lvsa.mp4
+"""
+
+import argparse
+import os
+import time
+
+import requests
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--server", default=os.environ.get("OMNI_SERVER", "http://localhost:8000"))
+    ap.add_argument("--prompt", default="A dog running through a sunlit forest.")
+    ap.add_argument("--frames", type=int, default=161)  # Wan 2x horizon
+    ap.add_argument("--height", type=int, default=480)
+    ap.add_argument("--width", type=int, default=832)
+    ap.add_argument("--fps", type=int, default=16)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--guidance", type=float, default=5.0)
+    ap.add_argument("--out", default="lvsa_long_video.mp4")
+    ap.add_argument("--poll-interval", type=float, default=5.0, help="Seconds between status polls.")
+    ap.add_argument(
+        "--max-wait", type=float, default=7200.0, help="Give up if the job hasn't finished after this many seconds."
+    )
+    args = ap.parse_args()
+
+    base = args.server.rstrip("/")
+    form = {
+        "prompt": args.prompt,
+        "size": f"{args.width}x{args.height}",
+        "num_frames": str(args.frames),
+        "fps": str(args.fps),
+        "num_inference_steps": str(args.steps),
+        "guidance_scale": str(args.guidance),
+    }
+
+    # ── 1. submit the async job ──────────────────────────────────────────────
+    print(f"[lvsa-app] POST {base}/v1/videos  {args.width}x{args.height} {args.frames}f steps={args.steps}")
+    resp = requests.post(f"{base}/v1/videos", data=form, timeout=120)
+    resp.raise_for_status()
+    job = resp.json()
+    vid = job["id"]
+    print(f"[lvsa-app] job {vid} ({job.get('status', '?')})")
+
+    # ── 2. poll until completed / failed ─────────────────────────────────────
+    deadline = time.time() + args.max_wait
+    last = None
+    while True:
+        if time.time() > deadline:
+            raise SystemExit(f"[lvsa-app] timed out after {args.max_wait:.0f}s (job {vid} still running)")
+        time.sleep(args.poll_interval)
+        r = requests.get(f"{base}/v1/videos/{vid}", timeout=60)
+        if r.status_code != 200:
+            # retrieve_video returns a non-200 JSON when the job has FAILED
+            raise SystemExit(f"[lvsa-app] job failed: {r.text}")
+        st = r.json()
+        status = st.get("status")
+        progress = st.get("progress", 0)
+        if (status, progress) != last:
+            print(f"[lvsa-app] {status}  {progress}%")
+            last = (status, progress)
+        if status == "completed":
+            break
+        if status == "failed":
+            raise SystemExit(f"[lvsa-app] job failed: {st.get('error')}")
+
+    # ── 3. download the rendered mp4 ─────────────────────────────────────────
+    c = requests.get(f"{base}/v1/videos/{vid}/content", timeout=600)
+    c.raise_for_status()
+    with open(args.out, "wb") as f:
+        f.write(c.content)
+    print(f"[lvsa-app] saved {args.out}")
+
+
+if __name__ == "__main__":
+    main()

apps/lvsa-omni-showcase/serve_lvsa.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+# Start an LVSA-enabled vLLM-Omni server for Wan or HunyuanVideo.
+#
+# LVSA is selected as the attention backend via --diffusion-attention-config (see
+# below); the `LVSA_*` env vars only tune it. The `lvsa-vllm-omni` plugin registers
+# itself through vLLM-Omni's `general_plugins` entry points — nothing in vLLM-Omni
+# core is modified. Watch the log for "[LVSA] Geometry detected" to confirm sparse
+# attention is active for your request geometry.
+#
+#   MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan     FRAMES=161 bash serve_lvsa.sh
+#   MODEL=/path/to/HunyuanVideo-1.5-...       MODEL_FAMILY=hunyuan FRAMES=193 bash serve_lvsa.sh
+set -euo pipefail
+
+MODEL=${MODEL:?set MODEL=/path/to/Wan-or-HunyuanVideo-Diffusers}
+MODEL_FAMILY=${MODEL_FAMILY:-wan}     # wan | hunyuan
+PORT=${PORT:-8000}
+FRAMES=${FRAMES:-161}                 # target horizon in frames
+VAE_T=${VAE_T:-4}                     # VAE temporal compression (Wan/Hunyuan = 4)
+BACKEND=${LVSA_BACKEND:-flashinfer}   # flashinfer (fused, fastest) | sdpa
+
+T_LAT=$(( (FRAMES - 1) / VAE_T + 1 ))
+
+# Per-family training reference horizon (Wan train ref = 21, Hunyuan = 33).
+case "$MODEL_FAMILY" in
+  wan)     REFERENCE=${REFERENCE:-21} ;;
+  hunyuan) REFERENCE=${REFERENCE:-33} ;;
+  *) echo "MODEL_FAMILY must be 'wan' or 'hunyuan'"; exit 1 ;;
+esac
+export LVSA_BACKEND="$BACKEND"
+export LVSA_TOTAL_LATENT_FRAMES="$T_LAT"
+export LVSA_REFERENCE_LATENT_FRAMES="$REFERENCE"
+export LVSA_ROTATE_KEYFRAMES=${LVSA_ROTATE_KEYFRAMES:-1}
+
+echo "[lvsa-app] serving $MODEL ($MODEL_FAMILY) on :$PORT"
+echo "[lvsa-app]   LVSA backend=$BACKEND  T_lat=$T_LAT (frames=$FRAMES)  reference=$REFERENCE"
+echo "[lvsa-app]   confirm activation by grep '\\[LVSA\\] Geometry detected' in this log"
+
+# vLLM-Omni 0.22 selects the LVSA attention backend per-role via
+# --diffusion-attention-config. The backend is what drives sparse attention
+# here; it engages under tensor-parallel too, unlike the LVSA_*_HOOK monkey-patch
+# path (which falls back to dense under sequence-parallel), so we deliberately do
+# NOT set a hook.
+exec vllm serve "$MODEL" --omni --host 0.0.0.0 --port "$PORT" \
+     --diffusion-attention-config '{"per_role": {"self": {"backend": "LVSA"}}}' \
+     --init-timeout 1800 ${EXTRA_ARGS:-}

benchmarks/tts/model_configs.yaml

@@ -80,6 +80,17 @@ models:
     task_extra_body:
       voice_clone: {}
 
+  # MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local
+  # depth transformer (n_vq=12), paired with the 48 kHz stereo
+  # MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the
+  # other MOSS-TTS variants.
+  OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+
   # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
   # has single-speaker rows; bench piggybacks on it (both speakers share the
   # same ref_audio). For real dialogue evaluation use the ttsd dataset.

docs/contributing/ci/CI_5levels.md

@@ -155,7 +155,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         /tests/dfx/stability/tests/test_wan22.json<br>
         <strong>Reliability:</strong><br>
         tests/dfx/reliability/test_reliability_{model_key}.py<br>
-        (e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
+        (e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
       </td>
       <td>
         <a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
@@ -741,14 +741,16 @@ L5 level testing focuses on the performance of model services under ***long-runn
 ### 4.2 Testing Content and Scope
 
 -   ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
--   ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
+-   ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
 
 ### 4.3 Test Directory and Execution Files
 
 -   ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
 -   ***Reliability Test Suite*** (`tests/dfx/reliability/`):
     -   `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
-    -   `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
+    -   `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
+    -   `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
+    -   `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
     -   `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
     -   `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
     -   `README.md` — Short local run commands for this directory.
@@ -846,22 +848,33 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
 | `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
 | `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
 | `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
-| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
+| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
+| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
+| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
 | `README.md` | Minimal run / collect examples. |
 
 #### Parametrization and markers
 
 - Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
-- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
 - **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
 
 #### CI trigger
 
-Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
+Weekly Buildkite (`.buildkite/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example:
+
+| Buildkite step | Test file | CI hardware |
+| -------------- | --------- | ----------- |
+| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
+| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
+| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
+| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
 
 ```bash
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
+pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
 ```
 
 #### Local commands
@@ -870,6 +883,8 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
 pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```
 
 #### Adding a new model suite
@@ -882,7 +897,7 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
 </details>
 
 -   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
+    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/test-weekly.yml` for new models)
 
 ## Summary
 

docs/contributing/ci/test_guide.md

@@ -96,15 +96,21 @@ Our test scripts use the pytest framework. First, please use `git clone https://
     # Stability: Wan2.2 (v1/videos diffusion benchmark loop)
     pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow
 
-    # Reliability: Qwen3-Omni
+    # Reliability: Qwen3-Omni (H100 × 2)
     pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow
 
-    # Reliability: Wan2.2
+    # Reliability: Wan2.2 (H100)
     pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow
 
+    # Reliability: HunyuanImage DiT (H100 × 4)
+    pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
+
+    # Reliability: VoxCPM2 (L4)
+    pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
+
     ```
 
-    The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+    The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-weekly.yml).
 
 You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
 

docs/models/supported_models.md

@@ -39,6 +39,7 @@ th {
 | `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `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` | ✅︎ |   |   |   |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |

docs/serving/speech_api.md

@@ -233,9 +233,9 @@ Server -> Client:
 
 | Message | Description |
 |---------|-------------|
-| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
+| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
 | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
-| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
+| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
 | `{"type": "session.done", "total_sentences": N}` | Session complete |
 | `{"type": "error", "message": "..."}` | Non-fatal error |
 
@@ -245,8 +245,7 @@ All REST API parameters are supported, plus:
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
-| `split_granularity` | string | "sentence" | Text splitting granularity |
+| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
 
 
 ```bash

docs/user_guide/fault_injection_reliability_matrix.md

@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
 
 ## Source of Conclusions
 
-The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
+The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/test-weekly.yml`.
 
 ### Example: `SIGTERM` fault injection log (Qwen3-Omni)
 

examples/offline_inference/text_to_speech/moss_tts/end2end.py

@@ -1,11 +1,12 @@
 """Offline inference example for MOSS-TTS variants.
 
-Covers all five models:
-  - OpenMOSS-Team/MOSS-TTS           (8B, general TTS)
-  - OpenMOSS-Team/MOSS-TTSD-v1.0     (8B, dialogue TTS)
-  - OpenMOSS-Team/MOSS-SoundEffect   (8B, sound effect synthesis)
-  - OpenMOSS-Team/MOSS-TTS-Realtime  (1.7B, low-latency streaming)
-  - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
+Covers six models:
+  - OpenMOSS-Team/MOSS-TTS                       (8B, general TTS)
+  - OpenMOSS-Team/MOSS-TTSD-v1.0                 (8B, dialogue TTS)
+  - OpenMOSS-Team/MOSS-SoundEffect               (8B, sound effect synthesis)
+  - OpenMOSS-Team/MOSS-TTS-Realtime              (1.7B, low-latency streaming)
+  - OpenMOSS-Team/MOSS-VoiceGenerator             (1.7B, voice design)
+  - OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo)
 
 Usage
 -----
@@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None:
         "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
         "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
         "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
+        "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local",
     }
     deploy_config = deploy_map.get(args.model, "moss_tts")
 
     is_sound_effect = "SoundEffect" in args.model
     is_voice_gen = "VoiceGenerator" in args.model
     is_realtime = "Realtime" in args.model
+    # Local-v1.5 ships its own AutoProcessor (processor_config.json +
+    # processing_moss_tts.py), so it reuses the same _build_unified_codes
+    # path as the Delay variants below — just a different n_vq/codec.
+    is_local = "Local" in args.model
+    output_sr = 48000 if is_local else 24000
 
     ref_audio = None
     builder_kwargs: dict = {}
@@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None:
             wav = mm.get("model_outputs")
         if wav is None:
             continue
+        # Local-v1.5's codec emits stereo chunks shaped (C, T); keep the
+        # channel axis intact (don't flatten) so concatenation below stays
+        # along the time axis. Mono variants stay (T,) as before.
         if isinstance(wav, list):
-            non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
+            non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
             if non_empty and t_first_audio is None:
                 t_first_audio = _time.monotonic()
             chunks.extend(non_empty)
         elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
             if t_first_audio is None:
                 t_first_audio = _time.monotonic()
-            chunks.append(wav.reshape(-1))
+            chunks.append(wav)
 
     t_gen_end = _time.monotonic()
     if not chunks:
         print("No audio in output — check model loading logs.")
         return
 
-    wav_t = torch.cat(chunks).float().cpu()
-    sf.write(args.output, wav_t.numpy(), 24000)
-    print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
+    cat_dim = -1 if chunks[0].dim() == 2 else 0
+    wav_t = torch.cat(chunks, dim=cat_dim).float().cpu()
+    n_samples = wav_t.shape[-1]
+    # soundfile wants (T,) mono or (T, C) multi-channel.
+    sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr)
+    print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}")
     ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
     gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
-    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
+    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}")
 
 
 def main() -> None:

recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md

@@ -0,0 +1,109 @@
+# Wan2.2 VACE-Fun Diffusers
+
+> All-in-one video creation and editing with Wan2.2 VACE-Fun A14B in diffusers format.
+
+## Summary
+
+- Vendor: Pyros13
+- Model: `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers`
+- Base model: `alibaba-pai/Wan2.2-VACE-Fun-A14B`
+- Task: VACE video generation and editing (T2V / I2V / R2V / V2V / inpainting)
+- Mode: Offline inference and OpenAI-compatible online serving
+- Maintainer: Community
+
+## When To Use This Recipe
+
+Use this recipe to run Wan2.2-VACE-Fun-A14B through the standard diffusers-style
+`Wan22VACEPipeline` path in vLLM-Omni. The checkpoint is a diffusers conversion of
+`alibaba-pai/Wan2.2-VACE-Fun-A14B`, with `model_index.json`, `transformer/`,
+`transformer_2/`, `text_encoder/`, `tokenizer/`, `vae/`, and `scheduler/`
+subfolders. Because the model already uses the expected diffusers layout, vLLM-Omni
+does not need any original-format key remapping or `.pth` component loaders.
+
+The model is a two-expert architecture. The high-noise `transformer` and low-noise
+`transformer_2` are switched by the Wan2.2 boundary during denoising.
+
+## References
+
+- Model card: <https://huggingface.co/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers>
+- Base model card: <https://huggingface.co/alibaba-pai/Wan2.2-VACE-Fun-A14B>
+- Related example: [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md)
+- Related issue: [vllm-project/vllm-omni#4206](https://github.com/vllm-project/vllm-omni/issues/4206)
+- PR: [vllm-project/vllm-omni#4667](https://github.com/vllm-project/vllm-omni/pull/4667)
+
+## Offline Inference
+
+```bash
+python examples/offline_inference/vace/vace_video_generation.py \
+  --model Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers \
+  --mode i2v --image ./i2v_input.JPG \
+  --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \
+  --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
+  --seed 42 --guidance-scale 5.0 \
+  --vae-use-tiling --enforce-eager \
+  --output vace_fun_i2v.mp4
+```
+
+See [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md)
+for other VACE modes and their input flags.
+
+## Online Serving
+
+### Server
+
+```bash
+vllm serve Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers --omni \
+  --model-class-name Wan22VACEPipeline \
+  --vae-use-tiling \
+  --enforce-eager \
+  --port 8091
+```
+
+For multi-GPU sequence-parallel serving, add the same parallelism flags you would
+use offline, for example `--ulysses-degree 4`. The explicit `--model-class-name`
+is optional for this diffusers checkpoint, but keeping it in examples makes the
+selected pipeline obvious.
+
+### Client
+
+```bash
+no_proxy=127.0.0.1 \
+curl -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -F "prompt=Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \
+  -F "input_reference=@/absolute/path/to/i2v_input.JPG" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=81" \
+  -F "fps=16" \
+  -F "num_inference_steps=30" \
+  -F "guidance_scale=5.0" \
+  --output vace_fun_i2v_serve.mp4
+```
+
+For video-to-video or inpainting-style requests, upload the reference video with
+`input_reference=@/absolute/path/to/input.mp4` or use `video_reference` with a URL
+or JSON-safe data URI.
+
+## Cache-DiT / TaylorSeer Notes
+
+`Wan22VACEPipeline` uses the shared Wan2.2 Cache-DiT enabler for both the
+high-noise and low-noise transformers. The adapter intentionally wraps only the
+main denoising `blocks`. VACE `vace_blocks` form a conditioning branch: each step
+they combine the current latent with the VACE context to produce hints that are
+then injected into selected main blocks. Recomputing that branch preserves the
+control signal; caching the main backbone still gives the acceleration target.
+
+## Validation
+
+Validated on NVIDIA B300 with `mkt_part_000.mp4`, 736x1280, 61 frames, 30
+inference steps, `seed=1`, and `boundary_ratio=0.875`.
+
+The Pyros13 diffusers checkpoint produced a byte-identical MP4 to the
+`alibaba-pai/Wan2.2-VACE-Fun-A14B` original-format checkpoint under the same
+vLLM-Omni pipeline and inputs:
+
+- SHA256: `2e2db74a612a009b0d6ad6d739e039883f0de84ce5cdf2676324daf78fb05ef8`
+- Frame shape: 61 x 736 x 1280
+- MAE / RMSE / max pixel diff: 0 / 0 / 0
+- PSNR: infinite
+- Mean SSIM: 1.0

tests/config/test_omni_config.py

@@ -0,0 +1,752 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the additive structured Omni config."""
+
+from __future__ import annotations
+
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+from pydantic import ValidationError
+
+from vllm_omni.config import omni_config as omni_config_module
+from vllm_omni.config.omni_config import (
+    BaseVllmOmniStageConfig,
+    OmniStageCacheConfig,
+    OmniStageConnectorConfig,
+    OmniStageDiffusionParallelConfig,
+    OmniStageLoadConfig,
+    OmniStageModelConfig,
+    OmniStageParallelConfig,
+    OmniStageRuntimeConfig,
+    OmniStageSchedulerConfig,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+    VllmOmniDiffusionStageConfig,
+    VllmOmniGenerationStageConfig,
+)
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import (
+    _STAGE_DEPLOY_FIELDS,
+    PIPELINE_WIDE_ENGINE_FIELDS,
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    load_deploy_config,
+    merge_pipeline_deploy,
+)
+from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
+
+
+@pytest.fixture(autouse=True)
+def _stable_test_platform(monkeypatch):
+    from vllm_omni import platforms
+
+    platform = platforms.current_omni_platform
+    monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
+    monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
+
+
+def _load_default_deploy(model_type: str) -> DeployConfig:
+    deploy_path = _DEPLOY_DIR / f"{model_type}.yaml"
+    if deploy_path.exists():
+        return load_deploy_config(deploy_path)
+    return DeployConfig()
+
+
+def _resolve_pipeline_or_skip(model_type: str) -> PipelineConfig:
+    registered = OMNI_PIPELINES[model_type]
+    pipeline = registered(None) if callable(registered) else registered
+    if pipeline is None:
+        pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
+    return pipeline
+
+
+@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
+def test_vllm_omni_config_from_registry_matches_merge_pipeline_deploy(model_type: str):
+    pipeline = _resolve_pipeline_or_skip(model_type)
+    legacy_deploy = _load_default_deploy(model_type)
+
+    legacy_stages = merge_pipeline_deploy(pipeline, legacy_deploy)
+    omni_config = VllmOmniConfig.from_registry(model_type)
+
+    assert omni_config.pipeline_config is pipeline
+    assert len(omni_config.stage_configs) == len(legacy_stages)
+
+    for legacy_stage, omni_stage in zip(legacy_stages, omni_config.stage_configs, strict=True):
+        assert omni_config.stage_by_id(legacy_stage.stage_id) is omni_stage
+
+        assert omni_stage.stage_pipeline_config is pipeline.get_stage(legacy_stage.stage_id)
+        assert omni_stage.model_config.default_sampling_params == legacy_stage.yaml_extras.get(
+            "default_sampling_params"
+        )
+        assert omni_stage.connector_config.output_connectors == legacy_stage.yaml_extras.get("output_connectors")
+        assert omni_stage.connector_config.input_connectors == legacy_stage.yaml_extras.get("input_connectors")
+        assert omni_stage.runtime_config.devices == legacy_stage.yaml_runtime.get("devices")
+        assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
+
+        engine_args = legacy_stage.yaml_engine_args
+        assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
+        assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
+        assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
+        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90)
+        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128)
+        assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens")
+        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True)
+        legacy_parallel_config = engine_args.get("parallel_config") or {}
+        assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get(
+            "tensor_parallel_size",
+            engine_args.get("tensor_parallel_size", 1),
+        )
+
+        if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.DIFFUSION:
+            assert isinstance(omni_stage, VllmOmniDiffusionStageConfig)
+            assert omni_stage.diffusion_config is not None
+            assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
+            assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
+        elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
+            assert isinstance(omni_stage, VllmOmniARStageConfig)
+            assert not hasattr(omni_stage, "diffusion_config")
+        else:
+            assert isinstance(omni_stage, VllmOmniGenerationStageConfig)
+            assert not hasattr(omni_stage, "diffusion_config")
+
+
+def test_stage_by_id_raises_for_unknown_stage():
+    omni_config = VllmOmniConfig.from_registry("qwen3_tts")
+
+    with pytest.raises(KeyError, match="no stage 99"):
+        omni_config.stage_by_id(99)
+
+
+def test_resolve_execution_mode_rejects_unknown_execution_type():
+    with pytest.raises(ValueError, match="Unsupported stage execution type"):
+        omni_config_module._resolve_execution_mode("unknown_execution_type")
+
+
+def test_from_registry_preserves_current_pipeline_config_object():
+    omni_config = VllmOmniConfig.from_registry("minicpmo_4_5")
+    pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
+
+    assert omni_config.pipeline_config is pipeline
+    assert not hasattr(omni_config, "pipeline")
+    assert "hf_config_predicate" in {f.name for f in fields(PipelineConfig)}
+    assert omni_config.pipeline_config.hf_config_predicate is pipeline.hf_config_predicate
+
+
+def test_from_registry_normalizes_stage_engine_extras_without_expanding_stage_deploy_config():
+    assert not hasattr(StageDeployConfig, "model_config")
+    assert not hasattr(StageDeployConfig, "parallel_config")
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2").stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.tensor_parallel_size == 1
+    assert stage.parallel_config.cfg_parallel_size == 2
+    assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena"
+
+
+def test_from_registry_applies_cli_overrides_without_stage_config_runtime_bridge():
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "stage_0_max_num_seqs": 7,
+            "stage_1_tensor_parallel_size": 2,
+        },
+    )
+
+    stage0 = omni_config.stage_by_id(0)
+    stage1 = omni_config.stage_by_id(1)
+
+    assert stage0.scheduler_config.max_num_seqs == 7
+    assert stage1.parallel_config.tensor_parallel_size == 2
+    assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size
+
+
+def test_runtime_num_gpus_is_derived_from_parallel_world_size():
+    omni_config = VllmOmniConfig.from_registry("hunyuan_image3_dit")
+    stage = omni_config.stage_by_id(0)
+
+    assert stage.parallel_config.tensor_parallel_size == 4
+    assert stage.parallel_config.world_size == 4
+    assert stage.runtime_config.num_gpus == 4
+
+
+def test_runtime_num_gpus_ignores_stale_runtime_override():
+    omni_config = VllmOmniConfig.from_registry(
+        "hunyuan_image3_dit",
+        cli_overrides={
+            "stage_0_num_gpus": 1,
+        },
+    )
+    stage = omni_config.stage_by_id(0)
+
+    assert stage.parallel_config.world_size == 4
+    assert stage.runtime_config.num_gpus == 4
+
+
+def test_from_registry_does_not_route_server_cli_keys_to_diffusion_stage():
+    omni_config = VllmOmniConfig.from_registry(
+        "dreamzero",
+        deploy_config_path="dreamzero_tp1_cfg2",
+        cli_overrides={
+            "host": "0.0.0.0",
+            "port": 8000,
+            "api_key": "secret",
+            "stage_0_host": "127.0.0.1",
+            "stage_0_port": 23456,
+        },
+    )
+
+    stage = omni_config.stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.host == "127.0.0.1"
+    assert stage.diffusion_config.port == 23456
+    assert not hasattr(stage.diffusion_config, "api_key")
+
+
+def test_pipeline_deploy_cli_fields_reuse_legacy_pipeline_wide_engine_fields():
+    assert omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS is PIPELINE_WIDE_ENGINE_FIELDS
+    assert "active_stream_window" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
+    assert "custom_voice_dir" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
+
+
+def test_pipeline_wide_model_fields_are_retained_on_structured_stage_configs(tmp_path):
+    custom_voice_dir = tmp_path / "voices"
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "active_stream_window": 2,
+            "custom_voice_dir": str(custom_voice_dir),
+        },
+    )
+
+    assert {stage.model_config.active_stream_window for stage in omni_config.stage_configs} == {2}
+    assert {stage.model_config.custom_voice_dir for stage in omni_config.stage_configs} == {str(custom_voice_dir)}
+
+
+def test_stage_deploy_engine_fields_reuse_legacy_stage_deploy_fields():
+    assert omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS == tuple(_STAGE_DEPLOY_FIELDS)
+    assert "tensor_parallel_size" in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
+    assert "stage_id" not in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
+
+
+def test_public_config_exports_use_stage_specific_sub_config_names():
+    import vllm_omni.config as config_pkg
+
+    generic_names = {
+        "CacheConfig",
+        "ConnectorConfig",
+        "LoadConfig",
+        "ModelConfig",
+        "OrchestratorConfig",
+        "ParallelConfig",
+        "RuntimeConfig",
+        "SchedulerConfig",
+    }
+
+    assert generic_names.isdisjoint(config_pkg.__all__)
+    assert {
+        "OmniStageCacheConfig",
+        "OmniStageConnectorConfig",
+        "OmniStageDiffusionParallelConfig",
+        "OmniStageLoadConfig",
+        "OmniStageModelConfig",
+        "VllmOmniOrchestratorConfig",
+        "OmniStageParallelConfig",
+        "OmniStageRuntimeConfig",
+        "OmniStageSchedulerConfig",
+        "StageConfigType",
+    }.issubset(config_pkg.__all__)
+
+
+def test_from_registry_keeps_worker_backend_separate_from_distributed_executor_backend():
+    omni_config = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
+
+    stage = omni_config.stage_by_id(0)
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.distributed_executor_backend == "mp"
+    assert omni_config.orchestrator_config.worker_backend == "multi_process"
+
+
+def test_from_registry_maps_orchestrator_cli_overrides():
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "stage_init_timeout": 1200,
+            "init_timeout": 1800,
+            "worker_backend": "ray",
+            "ray_address": "ray://127.0.0.1:10001",
+            "omni_master_address": "127.0.0.1",
+            "omni_master_port": 12345,
+            "omni_dp_size_local": 2,
+            "omni_lb_policy": "round_robin",
+            "omni_heartbeat_timeout": 9.5,
+            "shm_threshold_bytes": 4096,
+            "batch_timeout": 3,
+        },
+    )
+
+    orchestrator_config = omni_config.orchestrator_config
+    assert orchestrator_config.stage_init_timeout == 1200
+    assert orchestrator_config.init_timeout == 1800
+    assert orchestrator_config.worker_backend == "ray"
+    assert orchestrator_config.ray_address == "ray://127.0.0.1:10001"
+    assert orchestrator_config.omni_master_address == "127.0.0.1"
+    assert orchestrator_config.omni_master_port == 12345
+    assert orchestrator_config.omni_dp_size_local == 2
+    assert orchestrator_config.omni_lb_policy == "round_robin"
+    assert orchestrator_config.omni_heartbeat_timeout == 9.5
+    assert orchestrator_config.shm_threshold_bytes == 4096
+    assert orchestrator_config.batch_timeout == 3
+
+
+def test_from_registry_records_loaded_deploy_path_on_orchestrator_config():
+    omni_config = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
+
+    assert omni_config.pipeline_config.model_type == "dreamzero"
+    assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / "dreamzero_tp1_cfg2.yaml")
+
+
+def test_from_registry_dispatches_async_chunk_processors_without_mutating_topology():
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+
+    async_config = VllmOmniConfig.from_registry("qwen3_tts")
+    assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
+    assert async_config.stage_by_id(1).custom_process_input_func is None
+
+    sync_config = VllmOmniConfig.from_registry("qwen3_tts", cli_overrides={"async_chunk": False})
+    assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
+    assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
+
+    assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
+    assert pipeline.get_stage(1).custom_process_input_func is None
+
+
+def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
+    assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
+    assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
+
+    public_fields = {f.name for f in fields(BaseVllmOmniStageConfig)}
+
+    assert public_fields == {
+        "stage_pipeline_config",
+        "model_config",
+        "load_config",
+        "cache_config",
+        "scheduler_config",
+        "connector_config",
+        "runtime_config",
+        "parallel_config",
+        "quantization_config",
+    }
+    assert "diffusion_config" not in public_fields
+    assert {f.name for f in fields(VllmOmniDiffusionStageConfig)} == public_fields | {"diffusion_config"}
+    assert {f.name for f in fields(VllmOmniARStageConfig)} == public_fields
+    assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields
+
+
+def test_runtime_config_fields_match_rfc_runtime_scope():
+    assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
+        "devices",
+        "num_replicas",
+        "env",
+        "num_gpus",
+        "log_level",
+        "log_stats",
+        "profiler_config",
+    }
+
+
+def test_sub_config_fields_match_rfc_scopes():
+    assert {f.name for f in fields(OmniStageModelConfig)} == {
+        "active_stream_window",
+        "enable_sleep_mode",
+        "default_sampling_params",
+        "subtalker_sampling_params",
+        "has_sampling_extra_args",
+        "custom_voice_dir",
+        "task_type",
+        "codec_frame_rate_hz",
+        "enforce_eager",
+        "enable_flashinfer_autotune",
+        "compilation_config",
+        "enable_multithread_weight_load",
+        "num_weight_load_threads",
+        "disable_autocast",
+    }
+    assert {f.name for f in fields(OmniStageLoadConfig)} == {
+        "load_format",
+        "tokenizer_mode",
+        "config_format",
+        "skip_mm_profiling",
+    }
+    assert {f.name for f in fields(OmniStageCacheConfig)} == {
+        "gpu_memory_utilization",
+        "enable_prefix_caching",
+        "disable_hybrid_kv_cache_manager",
+        "mm_processor_cache_gb",
+    }
+    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
+        "max_num_seqs",
+        "max_num_batched_tokens",
+        "max_model_len",
+        "enable_chunked_prefill",
+        "async_scheduling",
+    }
+    assert {f.name for f in fields(OmniStageConnectorConfig)} == {
+        "stage_connector",
+        "output_connectors",
+        "input_connectors",
+    }
+    assert {f.name for f in fields(OmniStageParallelConfig)} == {
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "tensor_parallel_size",
+        "enable_expert_parallel",
+        "world_size",
+    }
+    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "tensor_parallel_size",
+        "sequence_parallel_size",
+        "ulysses_degree",
+        "ring_degree",
+        "ulysses_mode",
+        "cfg_parallel_size",
+        "vae_patch_parallel_size",
+        "use_hsdp",
+        "mask_sp_padding",
+        "hsdp_shard_size",
+        "hsdp_replicate_size",
+        "enable_expert_parallel",
+        "world_size",
+    }
+
+
+def test_diffusion_parallel_config_fields_cover_legacy_surface():
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
+    structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
+    expected_upstream_fields = {"mask_sp_padding"}
+
+    assert legacy_fields | expected_upstream_fields <= structured_fields
+    assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
+
+
+def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
+    cfg = OmniStageDiffusionParallelConfig(
+        pipeline_parallel_size=2,
+        data_parallel_size=3,
+        tensor_parallel_size=4,
+        cfg_parallel_size=3,
+        mask_sp_padding=True,
+    )
+
+    assert cfg.pipeline_parallel_size == 2
+    assert cfg.data_parallel_size == 3
+    assert cfg.cfg_parallel_size == 3
+    assert cfg.mask_sp_padding is True
+    assert cfg.world_size == 72
+
+
+def test_parallel_config_derived_fields_are_not_init_inputs():
+    with pytest.raises(ValidationError):
+        OmniStageParallelConfig(world_size=4)
+
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(world_size=4)
+
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(sequence_parallel_size=2)
+
+
+def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae_patch_parallel():
+    cfg = OmniStageDiffusionParallelConfig(
+        tensor_parallel_size=2,
+        cfg_parallel_size=2,
+        vae_patch_parallel_size=4,
+    )
+
+    assert cfg.vae_patch_parallel_size == 4
+    assert cfg.world_size == 4
+
+
+def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
+    cfg = OmniStageDiffusionParallelConfig(
+        pipeline_parallel_size=2,
+        ulysses_degree=2,
+        use_hsdp=True,
+        hsdp_shard_size=-1,
+        hsdp_replicate_size=2,
+    )
+
+    assert cfg.hsdp_shard_size == 2
+    assert cfg.world_size == 4
+
+
+def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
+    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
+
+    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
+
+
+def test_from_registry_preserves_legacy_pp_dp_for_world_size():
+    cfg = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(0).parallel_config
+
+    assert cfg.pipeline_parallel_size == 1
+    assert cfg.data_parallel_size == 1
+    assert cfg.tensor_parallel_size == 4
+    assert cfg.world_size == 4
+
+
+def test_from_registry_derives_sequence_parallel_size_from_degrees(tmp_path):
+    deploy_path = tmp_path / "dreamzero_derived_parallel.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    parallel_config:",
+                "      sequence_parallel_size: 99",
+                "      ulysses_degree: 2",
+                "      ring_degree: 3",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.sequence_parallel_size == 6
+    assert stage.parallel_config.world_size == 6
+
+
+def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound():
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+
+
+def test_stage_realizations_use_stage_specific_parallel_config_types():
+    qwen_config = VllmOmniConfig.from_registry("qwen3_tts")
+    ar_stage = qwen_config.stage_by_id(0)
+    generation_stage = qwen_config.stage_by_id(1)
+    diffusion_stage = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(0)
+
+    assert isinstance(ar_stage, VllmOmniARStageConfig)
+    assert type(ar_stage.parallel_config) is OmniStageParallelConfig
+    assert not hasattr(ar_stage.parallel_config, "cfg_parallel_size")
+    assert not hasattr(ar_stage.parallel_config, "sequence_parallel_size")
+    assert not hasattr(ar_stage.parallel_config, "ulysses_degree")
+
+    assert isinstance(generation_stage, VllmOmniGenerationStageConfig)
+    assert type(generation_stage.parallel_config) is OmniStageParallelConfig
+    assert not hasattr(generation_stage.parallel_config, "cfg_parallel_size")
+    assert not hasattr(generation_stage.parallel_config, "sequence_parallel_size")
+    assert not hasattr(generation_stage.parallel_config, "ulysses_degree")
+
+    assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
+    assert isinstance(diffusion_stage.parallel_config, OmniStageDiffusionParallelConfig)
+    assert diffusion_stage.parallel_config.cfg_parallel_size == 1
+    assert diffusion_stage.parallel_config.sequence_parallel_size == 1
+    assert diffusion_stage.parallel_config.ulysses_degree == 1
+
+
+def test_from_registry_preserves_diffusion_parallel_mask_sp_padding(tmp_path):
+    deploy_path = tmp_path / "dreamzero_mask_sp_padding.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    parallel_config:",
+                "      mask_sp_padding: true",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.mask_sp_padding is True
+
+
+def test_from_registry_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("qwen3_tts"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    omni_stage = VllmOmniConfig.from_registry("qwen3_tts").stage_by_id(legacy_stage.stage_id)
+
+    assert omega_stage.stage_id == omni_stage.stage_id
+    assert omega_stage.stage_type == omni_stage.stage_type.value
+    assert omega_stage.engine_input_source == omni_stage.input_sources
+    assert omega_stage.final_output == omni_stage.final_output
+    assert omega_stage.final_output_type == omni_stage.final_output_type
+    assert omega_stage.is_comprehension == omni_stage.is_comprehension
+    assert omega_stage.engine_args.model_stage == omni_stage.model_stage
+    assert omega_stage.engine_args.worker_type == omni_stage.worker_type
+    assert omega_stage.engine_args.scheduler_cls == omni_stage.scheduler_cls
+    assert omega_stage.runtime.process is True
+    assert omega_stage.runtime.requires_multimodal_data == omni_stage.requires_multimodal_data
+
+
+def test_from_registry_matches_to_omegaconf_diffusion_parallel_config():
+    pipeline = _resolve_pipeline_or_skip("hunyuan_image3_dit")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("hunyuan_image3_dit"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    omni_stage = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(legacy_stage.stage_id)
+
+    assert (
+        omega_stage.engine_args.parallel_config.pipeline_parallel_size
+        == omni_stage.parallel_config.pipeline_parallel_size
+    )
+    assert omega_stage.engine_args.parallel_config.data_parallel_size == omni_stage.parallel_config.data_parallel_size
+    assert (
+        omega_stage.engine_args.parallel_config.tensor_parallel_size == omni_stage.parallel_config.tensor_parallel_size
+    )
+    assert (
+        omega_stage.engine_args.parallel_config.sequence_parallel_size
+        == omni_stage.parallel_config.sequence_parallel_size
+    )
+    assert omega_stage.engine_args.parallel_config.cfg_parallel_size == omni_stage.parallel_config.cfg_parallel_size
+    assert (
+        omega_stage.engine_args.parallel_config.vae_patch_parallel_size
+        == omni_stage.parallel_config.vae_patch_parallel_size
+    )
+
+
+def test_from_registry_matches_build_engine_args_dict_behavior_for_representative_stage(monkeypatch):
+    from vllm_omni.engine import stage_init_utils
+
+    monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("qwen3_tts"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    legacy_engine_args = build_engine_args_dict(
+        omega_stage,
+        model="/tmp/qwen3-tts",
+        stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}},
+    )
+    omni_stage = VllmOmniConfig.from_registry("qwen3_tts").stage_by_id(legacy_stage.stage_id)
+
+    assert legacy_engine_args["model"] == "/tmp/qwen3-tts"
+    assert legacy_engine_args["stage_id"] == omni_stage.stage_id
+    assert legacy_engine_args["model_stage"] == omni_stage.model_stage
+    assert legacy_engine_args["worker_type"] == omni_stage.worker_type
+    assert legacy_engine_args["scheduler_cls"] == omni_stage.scheduler_cls
+    assert legacy_engine_args["stage_connector_spec"] == {"name": "SharedMemoryConnector", "extra": {}}
+    assert legacy_engine_args["has_sampling_extra_args"] == bool(
+        (omni_stage.model_config.default_sampling_params or {}).get("extra_args")
+    )
+    assert omni_stage.model_config.has_sampling_extra_args == legacy_engine_args["has_sampling_extra_args"]
+
+
+def test_from_registry_derives_has_sampling_extra_args_from_stage_defaults():
+    stage = VllmOmniConfig.from_registry("voxtral_tts").stage_by_id(0)
+
+    assert (stage.model_config.default_sampling_params or {}).get("extra_args")
+    assert stage.model_config.has_sampling_extra_args is True
+
+
+def test_diffusion_config_preserves_existing_coercion_hooks():
+    import torch
+
+    from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
+
+    cfg = omni_config_module._DiffusionConfigProjection(
+        dtype="float32",
+        cache_config={"rel_l1_thresh": 0.3},
+        diffusion_attention_config={"default": "flash_attn"},
+        diffusion_kv_cache_skip_steps="0-2,4",
+        diffusion_kv_cache_skip_layers=[1, 3],
+    )
+
+    assert cfg.dtype is torch.float32
+    assert isinstance(cfg.cache_config, DiffusionCacheConfig)
+    assert isinstance(cfg.diffusion_attention_config, AttentionConfig)
+    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
+    assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1, 2, 4}
+    assert cfg.diffusion_kv_cache_skip_layer_indices == {1, 3}
+    assert cfg.max_cpu_loras == 1
+
+
+def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
+    monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
+
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        diffusion_attention_backend="flash_attn",
+        kv_cache_dtype="fp8",
+        kv_cache_skip_steps="0-1",
+        kv_cache_skip_layers=[2],
+        static_lora_scale=0.25,
+        diffusers_load_kwargs=None,
+        diffusers_call_kwargs=None,
+    )
+
+    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
+    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.diffusers_load_kwargs == {}
+    assert cfg.diffusers_call_kwargs == {}
+
+
+def test_from_registry_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
+    deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_attention_backend: flash_attn",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry(
+        "dreamzero",
+        deploy_config_path=str(deploy_path),
+    ).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
+
+
+def test_diffusion_config_field_classification_covers_current_fields():
+    classified_fields = (
+        omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
+        | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
+        | omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
+    )
+
+    assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
+    assert {
+        "enable_prompt_embed_cache",
+        "prompt_embed_cache_size",
+        "diffusion_kv_cache_dtype",
+    } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
+    assert {
+        "revision",
+        "trust_remote_code",
+        "distributed_executor_backend",
+    } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
+    assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS

tests/dfx/reliability/README.md

@@ -6,9 +6,13 @@ The detailed RFC document is maintained outside this repository (local/internal
 - Fault helpers: `conftest.py`, `helpers.py`
 - Qwen3 tests: `test_reliability_qwen3_omni.py`
 - Wan2.2 tests: `test_reliability_wan22.py`
+- HunyuanImage tests: `test_reliability_hunyuan_image.py`
+- VoxCPM2 tests: `test_reliability_voxcpm2.py`
 
 ```bash
 pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```

tests/dfx/reliability/test_reliability_hunyuan_image.py

@@ -0,0 +1,478 @@
+"""
+HunyuanImage-3.0-Instruct reliability integration tests (DiT-only deploy).
+
+Uses ``hunyuan_image3_dit.yaml`` (single diffusion stage, 4-GPU TP=4, no
+AR/Mooncake connector) for simpler local bring-up. Scenarios follow
+``docs/user_guide/fault_injection_reliability_matrix.md``.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+)
+from tests.dfx.reliability.helpers import (
+    FaultInjector,
+    assert_no_server_tree_process_residual_and_gpu_release,
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    make_server_root_kill_fault_injector,
+    make_server_tree_kill_fault_injector,
+    post_json_raw,
+    run_fault_injection_with_rate_load,
+    stop_gpu_oom_hogs,
+    worker_residual_timeout_after_kill_signal,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OpenAIClientHandler
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "hunyuan_image_dit_reliability_default",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "stage_config_name": "hunyuan_image3_dit.yaml",
+        },
+    },
+]
+
+DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
+OOM_INJECTION_CONFIG = {
+    "target_mem_ratio": 0.95,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": True,
+}
+FAULT_ERROR_KEYWORDS = (
+    "oom",
+    "out of memory",
+    "cuda",
+    "job failed",
+    "unknown error",
+    "internal server error",
+    "500 server error",
+)
+PROCESS_KILL_ERROR_KEYWORDS = (
+    "timeout",
+    "did not complete within",
+    "connection",
+    "engine",
+    "orchestrator",
+    "dead",
+    "internal",
+    "500",
+    "503",
+)
+# Matrix: worker SIGTERM/SIGKILL; DiT-only HunyuanImage workers are vLLM-Omni::.
+WORKER_KILL_PATTERNS = ("vLLM-Omni::",)
+SERVE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGINT", id="sigint"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+WORKER_SIGNAL_FAULT_PARAMS = [
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=WORKER_KILL_PATTERNS,
+            signal_name="SIGTERM",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigterm",
+    ),
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=WORKER_KILL_PATTERNS,
+            signal_name="SIGKILL",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigkill",
+    ),
+]
+
+HUNYUAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
+INFLIGHT_INJECTION_REQUEST_RATE = 0.3
+INFLIGHT_INJECTION_REQUEST_COUNT = 10
+
+
+def _image_request_config(omni_server: Any) -> dict[str, Any]:
+    return {
+        "json": {
+            "model": omni_server.model,
+            "prompt": "A simple red apple on a white background.",
+            "size": "512x512",
+            "n": 1,
+            "response_format": "b64_json",
+            "num_inference_steps": 4,
+            "guidance_scale": 7.5,
+            "seed": 42,
+            "bot_task": "none",
+            "use_system_prompt": "en_unified",
+        },
+        "timeout": 600,
+    }
+
+
+def _submit_image_generation(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
+    responses = openai_client.send_images_generations_http_request(_image_request_config(omni_server))
+    resp = responses[0]
+    if resp.status_code >= 400:
+        raise RuntimeError(
+            f"/v1/images/generations failed: status={resp.status_code}, body={getattr(resp, 'error_message', '')!r}"
+        )
+
+
+def _assert_post_fault_image_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
+    payload = {
+        "model": model,
+        "prompt": "post-fault fast-fail check",
+        "size": "512x512",
+        "n": 1,
+        "response_format": "b64_json",
+        "num_inference_steps": 4,
+        "bot_task": "none",
+        "use_system_prompt": "en_unified",
+    }
+    start = time.monotonic()
+    try:
+        status, body = post_json_raw(host, port, "/v1/images/generations", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] /v1/images/generations did not fail fast: {elapsed:.2f}s"
+        assert status >= 500, (
+            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
+        )
+    except Exception:  # noqa: BLE001
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
+
+
+def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            if status == 503:
+                return
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+            return
+        time.sleep(0.5)
+    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skip(reason="issue#4285")
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_image_large_request_failure(
+    omni_server_function,
+    openai_client_function,
+) -> None:
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        try:
+            _submit_image_generation(openai_client_function, omni_server_function)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected large /v1/images/generations request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_request_failure(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    try:
+        _submit_image_generation(openai_client_function, omni_server_after_fault_function)
+    except Exception as exc:
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected /v1/images/generations request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skip(reason="issue#4522")
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    model = omni_server_after_fault_function.model
+
+    payload = {
+        "model": model,
+        "prompt": "fast-fail check",
+        "size": "512x512",
+        "n": 1,
+        "response_format": "b64_json",
+        "num_inference_steps": 4,
+        "bot_task": "none",
+        "use_system_prompt": "en_unified",
+    }
+    start = time.monotonic()
+    try:
+        status, body = post_json_raw(host, port, "/v1/images/generations", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, (
+            f"[process_kill fast_fail] /v1/images/generations did not fail fast after fault: {elapsed:.2f}s"
+        )
+        assert status >= 500, (
+            "[process_kill fast_fail] expected server-side error after fatal fault, "
+            f"got status={status}, body={body[:200]!r}"
+        )
+    except Exception:
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, (
+            f"[process_kill fast_fail] /v1/images/generations exception was too slow after fault: {elapsed:.2f}s"
+        )
+
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(_submit_image_generation, openai_client_function, omni_server_after_fault_function)
+            for _ in range(3)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=40,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time image requests hung: pending={len(pending)}"
+    assert elapsed < 40, f"[process_kill concurrent] fault-time image request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    for future in done:
+        try:
+            future.result()
+        except Exception as exc:
+            fault_observed = True
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one /v1/images/generations request to fail after fault"
+    )
+
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    final_health_status: int | None = None
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            final_health_status = status
+            if status == 503:
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    assert final_health_status == 503, (
+        "[process_kill health] expected /health 503 after fatal fault, "
+        f"got status={final_health_status}, last_observation={last_observation}"
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skip(reason="issue#4526")
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_worker_with_load_request_failure(
+    omni_server_function,
+    openai_client_function,
+    fault_injector: FaultInjector,
+) -> None:
+    scenario = "kill_worker_with_load"
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+        inject_fault=lambda: fault_injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.skip(reason="issue#4526")
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )

tests/dfx/reliability/test_reliability_voxcpm2.py

@@ -0,0 +1,505 @@
+"""
+VoxCPM2 reliability integration tests.
+
+Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+)
+from tests.dfx.reliability.helpers import (
+    FaultInjector,
+    assert_no_server_tree_process_residual_and_gpu_release,
+    extract_openai_error_contract_from_bytes,
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    make_server_root_kill_fault_injector,
+    make_server_tree_kill_fault_injector,
+    post_json_raw,
+    run_fault_injection_with_rate_load,
+    stop_gpu_oom_hogs,
+    worker_residual_timeout_after_kill_signal,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OpenAIClientHandler
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "voxcpm2_reliability_default",
+        "server_params": {
+            "model": "openbmb/VoxCPM2",
+            "stage_config_name": "voxcpm2.yaml",
+        },
+    },
+]
+
+DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
+OOM_INJECTION_CONFIG = {
+    "target_mem_ratio": 0.96,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": False,
+}
+FAULT_ERROR_KEYWORDS = (
+    "the request failed",
+    "oom",
+    "out of memory",
+    "cuda",
+    "orchestrator",
+    "timeout",
+    "connection",
+    "500",
+    "503",
+)
+PROCESS_KILL_ERROR_KEYWORDS = (
+    "timeout",
+    "did not complete within",
+    "connection",
+    "engine",
+    "orchestrator",
+    "dead",
+    "internal",
+    "500",
+    "503",
+)
+RUNTIME_WORKER_PATTERN = "VLLM::"
+SERVE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGINT", id="sigint"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+SERVE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGINT", id="sigint"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+WORKER_SIGNAL_FAULT_PARAMS = [
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=RUNTIME_WORKER_PATTERN,
+            signal_name="SIGTERM",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigterm",
+    ),
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=RUNTIME_WORKER_PATTERN,
+            signal_name="SIGKILL",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigkill",
+    ),
+]
+
+VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
+INFLIGHT_INJECTION_REQUEST_RATE = 0.3
+INFLIGHT_INJECTION_REQUEST_COUNT = 10
+_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
+_SPEECH_FF_MAX_ELAPSED_SEC = 30
+
+
+def _speech_request_config(omni_server: Any) -> dict[str, Any]:
+    return {
+        "model": omni_server.model,
+        "input": "The weather is nice today, perfect for a walk in the park.",
+        "stream": False,
+        "timeout": 120.0,
+        "response_format": "wav",
+        "voice": "default",
+    }
+
+
+def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
+    openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
+
+
+def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
+    base = "The weather is nice today, perfect for a walk in the park. "
+    return {
+        "model": omni_server.model,
+        "input": base * 80,
+        "max_new_tokens": 4096,
+        "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
+        "stream": False,
+        "timeout": 300.0,
+        "response_format": "wav",
+        "voice": "default",
+    }
+
+
+def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
+    openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
+
+
+def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
+    payload = {
+        "model": model,
+        "input": "post-fault fast-fail check",
+        "voice": "default",
+    }
+    start = time.monotonic()
+    try:
+        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
+        assert status >= 500, (
+            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
+        )
+    except Exception:  # noqa: BLE001
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
+
+
+def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            if status == 503:
+                return
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+            return
+        time.sleep(0.5)
+    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skip(reason="issue#4285")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_speech_request_failure(
+    omni_server_function,
+    openai_client_function,
+) -> None:
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        try:
+            _submit_oom_speech_request(openai_client_function, omni_server_function)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skip(reason="issue#4285")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_speech_error_contract(
+    omni_server_function,
+) -> None:
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    oom_config = _oom_speech_request_config(omni_server_function)
+    payload = {
+        "model": oom_config["model"],
+        "input": oom_config["input"],
+        "voice": oom_config["voice"],
+        "max_new_tokens": oom_config["max_new_tokens"],
+        "ref_audio": oom_config["ref_audio"],
+    }
+    try:
+        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+    assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
+    speech_error = extract_openai_error_contract_from_bytes(body)
+    assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
+    assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_request_failure(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    try:
+        _submit_speech_request(openai_client_function, omni_server_after_fault_function)
+    except Exception as exc:
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    model = omni_server_after_fault_function.model
+    payload = {
+        "model": model,
+        "input": "Say hello in one short sentence.",
+        "voice": "default",
+    }
+
+    start = time.monotonic()
+    try:
+        status, body = post_json_raw(
+            host,
+            port,
+            "/v1/audio/speech",
+            payload,
+            timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
+        )
+        elapsed = time.monotonic() - start
+        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
+            f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
+        )
+        assert status >= 500, (
+            "[process_kill fast_fail] expected server-side failure after fault, "
+            f"got status={status}, body={body[:200]!r}"
+        )
+    except Exception:
+        elapsed = time.monotonic() - start
+        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
+            f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
+        )
+
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
+            for _ in range(3)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=30,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
+    assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    for future in done:
+        try:
+            future.result()
+        except Exception as exc:
+            fault_observed = True
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_worker_with_load_request_failure(
+    omni_server_function,
+    openai_client_function,
+    fault_injector: FaultInjector,
+) -> None:
+    scenario = "kill_worker_with_load"
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
+        inject_fault=lambda: fault_injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=300.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )

tests/dfx/stability/scripts/test_stability_voxcpm2.py

@@ -0,0 +1,69 @@
+"""
+VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
+
+Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from tests.dfx.conftest import (
+    create_benchmark_indices,
+    create_test_parameter_mapping,
+    create_unique_server_params,
+    load_configs,
+)
+from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
+
+STABILITY_DIR = Path(__file__).resolve().parent.parent
+DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
+CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
+DEFAULT_NUM_PROMPTS_PER_BATCH = 20
+STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
+
+try:
+    BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
+except FileNotFoundError:
+    BENCHMARK_CONFIGS = []
+
+test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
+benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
+def test_stability_voxcpm2(omni_server, stability_benchmark_params):
+    test_name = stability_benchmark_params["test_name"]
+    params = stability_benchmark_params["params"]
+    duration_sec = params.get("duration_sec", 300)
+    num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
+    request_rate = params.get("request_rate")
+    max_concurrency = params.get("max_concurrency")
+
+    bench_params = {
+        k: v
+        for k, v in params.items()
+        if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
+    }
+
+    result = run_stability_benchmark_loop(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=omni_server.model,
+        duration_sec=duration_sec,
+        params=bench_params,
+        request_rate=request_rate,
+        max_concurrency=max_concurrency,
+        result_dir=str(STABILITY_DIR),
+        num_prompts_per_batch=num_prompts_per_batch,
+        run_one_batch=_run_one_vllm_bench_batch,
+    )
+
+    assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
+    assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"

tests/dfx/stability/tests/test_voxcpm2.json

@@ -0,0 +1,35 @@
+[
+    {
+        "test_name": "test_voxcpm2_stability",
+        "server_params": {
+            "model": "openbmb/VoxCPM2",
+            "trust_remote_code": true
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-audio-speech",
+                "endpoint": "/v1/audio/speech",
+                "trust_remote_code": true,
+                "duration_sec": 86400,
+                "max_concurrency": 2,
+                "num_prompts_per_batch": 50,
+                "random_input_len": {
+                    "min": 0,
+                    "max": 1000
+                },
+                "random_output_len": {
+                    "min": 0,
+                    "max": 1000
+                },
+                "random_range_ratio": 0.0,
+                "extra_body": {
+                    "voice": "default",
+                    "language": "English"
+                },
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration"
+            }
+        ]
+    }
+]

tests/diffusion/cache/test_cache_dit.py

@@ -42,6 +42,10 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
+def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
+    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22
+
+
 @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
 def test_cache_dit_configs_have_separate_cfg(transformer_model):
     """Check models with separate CFG set has_separate_cfg=True in their configs."""
@@ -62,6 +66,34 @@ def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_blo
     assert adapter_kwargs["has_separate_cfg"] is True
 
 
+@pytest.mark.parametrize("has_transformer_2", [False, True])
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+def test_wan22_custom_enabler_passes_taylorseer_calibrator(
+    mock_cache_dit,
+    mock_block_adapter,
+    has_transformer_2,
+):
+    mock_pipeline = Mock()
+    mock_pipeline.transformer.blocks = [Mock()]
+    if has_transformer_2:
+        mock_pipeline.transformer_2.blocks = [Mock()]
+    else:
+        mock_pipeline.transformer_2 = None
+    cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1)
+
+    cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config)
+
+    enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs
+    calibrator_config = enable_cache_kwargs["calibrator_config"]
+    assert calibrator_config is not None
+    assert calibrator_config.taylorseer_order == 1
+
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    for modifier in adapter_kwargs["params_modifiers"]:
+        assert modifier._context_kwargs["calibrator_config"] is calibrator_config
+
+
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
 def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):

tests/diffusion/distributed/test_sequence_parallel.py

@@ -23,6 +23,7 @@ from PIL import Image
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.distributed.utils import build_local_sp_padding_mask
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -412,3 +413,85 @@ def test_sp_correctness_advanced(model_name: str):
             f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
         )
     print("=" * 70)
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A partially padded SP shard must receive a local-length mask."""
+    mask = build_local_sp_padding_mask(
+        batch_size=2,
+        local_seq_len=4,
+        sp_original_seq_len=5,
+        sp_padding_size=3,
+        sequence_parallel_rank=1,
+        device=torch.device(current_omni_platform.device_type),
+    )
+
+    expected = torch.tensor(
+        [
+            [True, False, False, False],
+            [True, False, False, False],
+        ],
+        dtype=torch.bool,
+        device=mask.device,
+    )
+    assert mask is not None
+    assert mask.shape == (2, 4)
+    assert torch.equal(mask, expected)
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A rank whose local shard contains no padding should not get a mask."""
+    mask = build_local_sp_padding_mask(
+        batch_size=2,
+        local_seq_len=4,
+        sp_original_seq_len=5,
+        sp_padding_size=3,
+        sequence_parallel_rank=0,
+        device=torch.device(current_omni_platform.device_type),
+    )
+
+    assert mask is None
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_wan_sp_plan() -> None:
+    """Wan2.2 must shard hidden states before CacheDiT-wrapped transformer blocks."""
+    try:
+        from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan
+        from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+    except ImportError as exc:
+        pytest.skip(f"WanTransformer3DModel not available: {exc}")
+
+    plan = getattr(WanTransformer3DModel, "_sp_plan", None)
+
+    assert plan is not None
+    assert "_sp_shard_point" in plan
+    assert "blocks.0" not in plan
+
+    shard_plan = plan["_sp_shard_point"]
+    assert 0 in shard_plan
+    assert shard_plan[0].split_dim == 1
+    assert shard_plan[0].expected_dims == 3
+    assert shard_plan[0].split_output is True
+
+    validate_sp_plan(plan)

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -14,6 +14,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     create_vace_transformer_from_config,
     get_wan22_vace_pre_process_func,
 )
+from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -93,6 +94,10 @@ def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch
     }
 
 
+def test_vace_transformer_layerwise_offloads_conditioning_blocks() -> None:
+    assert WanVACETransformer3DModel._layerwise_offload_blocks_attrs == ["vace_blocks", "blocks"]
+
+
 def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None:
     pipeline = _make_vace_pipeline()
     mask = torch.ones(1, 3, 5, 16, 16)

tests/diffusion/test_attention_compile.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.forward_context import set_forward_context
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
+    attention = object.__new__(Attention)
+    calls = []
+
+    def _boundary(query, key, value, attn_metadata=None):
+        calls.append("boundary")
+        return query
+
+    def _impl(query, key, value, attn_metadata=None):
+        calls.append("impl")
+        return query
+
+    attention._forward_hsdp_compile_boundary = _boundary
+    attention._forward_impl = _impl
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True))
+    query = torch.empty(1)
+    with set_forward_context(omni_diffusion_config=config):
+        assert Attention.forward(attention, query, query, query) is query
+
+    assert calls == ["boundary"]
+
+
+def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
+    attention = object.__new__(Attention)
+    calls = []
+
+    def _boundary(query, key, value, attn_metadata=None):
+        calls.append("boundary")
+        return query
+
+    def _impl(query, key, value, attn_metadata=None):
+        calls.append("impl")
+        return query
+
+    attention._forward_hsdp_compile_boundary = _boundary
+    attention._forward_impl = _impl
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False))
+    query = torch.empty(1)
+    with set_forward_context(omni_diffusion_config=config):
+        assert Attention.forward(attention, query, query, query) is query
+
+    assert calls == ["impl"]
+
+
+def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
+    attention = object.__new__(Attention)
+    calls = []
+
+    def _boundary(query, key, value, attn_metadata=None):
+        calls.append("boundary")
+        return query
+
+    def _impl(query, key, value, attn_metadata=None):
+        calls.append("impl")
+        return query
+
+    attention._forward_hsdp_compile_boundary = _boundary
+    attention._forward_impl = _impl
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    query = torch.empty(1)
+    with set_forward_context(vllm_config=None):
+        assert Attention.forward(attention, query, query, query) is query
+
+    assert calls == ["impl"]

tests/diffusion/test_compile.py

@@ -1,20 +1,28 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import pytest
 import torch.nn as nn
 
+import vllm_omni.diffusion.compile as compile_module
 from vllm_omni.diffusion.compile import regionally_compile
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 class _WrappedBlock(nn.Module):
     def __init__(self) -> None:
         super().__init__()
         self.compile_called = False
+        self.forward_compiled = False
 
     def compile(self, *args, **kwargs):
         self.compile_called = True
         return self
 
+    def forward(self, x):
+        return x
+
 
 class _ModelWithWrappedRepeatedBlocks(nn.Module):
     _repeated_blocks = ["OriginalBlock"]
@@ -26,10 +34,23 @@ class _ModelWithWrappedRepeatedBlocks(nn.Module):
         self.other_blocks = nn.ModuleList([_WrappedBlock()])
 
 
-def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr():
+def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(monkeypatch):
     model = _ModelWithWrappedRepeatedBlocks()
+    compile_calls = []
 
-    regionally_compile(model)
+    def _compile(fn, *args, **kwargs):
+        compile_calls.append((fn, args, kwargs))
 
-    assert all(block.compile_called for block in model.transformer_blocks)
+        def _compiled(*fn_args, **fn_kwargs):
+            return f"compiled:{fn(*fn_args, **fn_kwargs)}"
+
+        return _compiled
+
+    monkeypatch.setattr(compile_module.torch, "compile", _compile)
+
+    regionally_compile(model, dynamic=True)
+
+    assert len(compile_calls) == 2
+    assert all(not block.compile_called for block in model.transformer_blocks)
     assert not model.other_blocks[0].compile_called
+    assert model.transformer_blocks[0].forward("ok") == "compiled:ok"

tests/diffusion/test_diffusion_model_runner.py

@@ -80,6 +80,7 @@ def _make_request(skip_cache_refresh: bool = True):
         num_inference_steps=4,
     )
     return SimpleNamespace(
+        request_id="req-test",
         prompts=["a prompt"],
         sampling_params=sampling_params,
         skip_cache_refresh=skip_cache_refresh,
@@ -104,11 +105,43 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.kv_transfer_manager = SimpleNamespace(
         receive_kv_cache=lambda req, target_device=None: None,
         receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
-        receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+        receive_multi_kv_cache_distributed=lambda *a, **k: None,
     )
+    runner._kv_prefetch_enabled = False
     return runner
 
 
+def _make_compile_runner(*, use_hsdp: bool):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = SimpleNamespace(transformer=SimpleNamespace())
+    runner.od_config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=use_hsdp))
+    return runner
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("use_hsdp", [False, True])
+def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp):
+    runner = _make_compile_runner(use_hsdp=use_hsdp)
+    compile_calls = []
+
+    def _regionally_compile(model, *args, **kwargs):
+        compile_calls.append((model, args, kwargs))
+        return model
+
+    monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+
+    DiffusionModelRunner._compile_transformer(runner, "transformer")
+
+    assert compile_calls == [
+        (
+            runner.pipeline.transformer,
+            (),
+            {"dynamic": True},
+        )
+    ]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):

tests/diffusion/test_worker_wrapper_base.py

@@ -184,7 +184,7 @@ class TestWorkerWrapperBaseDelegation:
         mock_reqs = [mocker.Mock()]
         result = wrapper.execute_model(mock_reqs, mock_od_config)
 
-        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config)
+        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None)
         assert result == mock_output
 
     def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):

tests/distributed/omni_connectors/test_kv_async_prefetch.py

@@ -0,0 +1,356 @@
+"""Background KV prefetch — CPU-testable mechanism.
+
+Covers opt-in gating, the start_prefetch / consume_prefetched_kv round trip over a
+mock connector, dedup, per-call sender_info isolation, role classification, and
+miss/abort paths.  GPU H2D and D2D paths need CUDA (integration tests).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+    ReceiveRole,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# start_prefetch() skips stubs without a sender endpoint, so submissions carry it.
+_SENDER_INFO = {"host": "127.0.0.1", "zmq_port": 50051}
+
+
+class MockConnector:
+    def __init__(self):
+        self.store = {}
+        self.get_calls = []
+
+    def put(self, from_stage, to_stage, put_key, data):
+        self.store[f"{from_stage}->{to_stage}:{put_key}"] = data
+        return True, len(str(data)), None
+
+    def get(self, from_stage, to_stage, get_key, metadata=None):
+        self.get_calls.append({"get_key": get_key, "metadata": metadata})
+        key = f"{from_stage}->{to_stage}:{get_key}"
+        if key in self.store:
+            return self.store[key], len(str(self.store[key]))
+        return None
+
+
+def _make_sender_receiver(*, async_prefetch=True):
+    """Build sender+receiver managers sharing one mock connector, seed one KV."""
+    connector = MockConnector()
+
+    sender = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True
+        )
+    )
+    sender._connector = connector
+
+    receiver = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="sender",
+            stage_id="receiver",
+            need_recv_cache=True,
+            recv_timeout=1.0,
+        ),
+        async_prefetch=async_prefetch,
+    )
+    receiver._connector = connector
+    return sender, receiver, connector
+
+
+def _seed_payload(sender, req_id, *, num_layers=2, block_size=4, num_heads=4, head_dim=8):
+    kv_caches = [torch.randn(2, 5, block_size, num_heads, head_dim) for _ in range(num_layers)]
+    sender.handle_finished_requests_kv_transfer(
+        {req_id: {"block_ids": [0, 1], "seq_len": 10}}, kv_caches, block_size, "float32"
+    )
+
+
+def _req(rid):
+    """A minimal request carrying just ``request_id`` (consume_prefetched_kv only needs that)."""
+    return SimpleNamespace(request_id=rid)
+
+
+# --------------------------------------------------------------------------- #
+#  Opt-in gating
+# --------------------------------------------------------------------------- #
+
+
+def test_async_prefetch_defaults_false():
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    assert mgr._async_prefetch is False
+
+
+def test_from_od_config_enables_prefetch_flag():
+    od = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": True})
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True
+    od2 = SimpleNamespace(omni_kv_config=None)
+    assert OmniKVTransferManager.from_od_config(od2)._async_prefetch is False
+    od3 = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": False})
+    assert OmniKVTransferManager.from_od_config(od3)._async_prefetch is False
+
+
+def test_prefetch_preserved_with_connector_config():
+    # The flag must survive merging with a populated connector config.
+    od = SimpleNamespace(
+        omni_kv_config={
+            "connector_config": {"type": "mock"},
+            "need_recv_cache": True,
+            "enable_kv_async_prefetch": True,
+        },
+    )
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True
+
+
+def test_prefetch_disabled_with_cfg_companion_collector():
+    # A CFG companion collector forces the synchronous path (see _resolve_async_prefetch).
+    od = SimpleNamespace(
+        omni_kv_config={"enable_kv_async_prefetch": True},
+        cfg_kv_collect_func=lambda *a, **k: {},
+    )
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is False
+
+
+def test_from_vllm_config_ar_path_is_always_sync():
+    # The AR entry point never enables prefetch.
+    model_config = SimpleNamespace(omni_kv_config=None)
+    vllm_config = SimpleNamespace(kv_transfer_config=None)
+    assert OmniKVTransferManager.from_vllm_config(vllm_config, model_config)._async_prefetch is False
+
+
+def test_start_prefetch_noop_when_disabled():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    receiver.start_prefetch({"request_id": "r1", "kv_sender_info": None})
+    assert receiver._prefetch_futures == {}
+
+
+# --------------------------------------------------------------------------- #
+#  Round trip
+# --------------------------------------------------------------------------- #
+
+
+def test_prefetch_round_trip_returns_data():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-1")
+
+    receiver.start_prefetch({"request_id": "rid-1", "kv_sender_info": _SENDER_INFO})
+    assert "rid-1" in receiver._prefetch_futures
+
+    data, size = receiver.consume_prefetched_kv(_req("rid-1"))
+    assert data is not None
+    assert "layer_blocks" in data
+    assert data["metadata"]["seq_len"] == 10
+    assert size > 0
+    # Future is consumed (popped) on retrieval.
+    assert "rid-1" not in receiver._prefetch_futures
+
+
+def test_consume_prefetched_kv_miss_without_prefetch_returns_none():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    assert receiver.consume_prefetched_kv(_req("never-prefetched")) == (None, 0)
+
+
+def test_start_prefetch_dedups_same_request():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-dup")
+    receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO})
+    fut1 = receiver._prefetch_futures["rid-dup"]
+    receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO})
+    assert receiver._prefetch_futures["rid-dup"] is fut1
+
+
+def test_start_prefetch_skips_without_sender_info():
+    # No endpoint => no submission; the sync path handles it.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-nosi")
+    receiver.start_prefetch({"request_id": "rid-nosi", "kv_sender_info": None})
+    assert receiver._prefetch_futures == {}
+
+
+def test_consume_prefetched_kv_sweeps_stale_entries():
+    # Consuming a slot drops every other tracked prefetch.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-stale")
+    receiver.start_prefetch({"request_id": "rid-stale", "kv_sender_info": _SENDER_INFO})
+    assert "rid-stale" in receiver._prefetch_futures
+    # Miss on the current request still sweeps stale entries.
+    receiver._async_prefetch = False
+    assert receiver.consume_prefetched_kv(_req("rid-current")) == (None, 0)
+    assert receiver._prefetch_futures == {}
+
+
+# --------------------------------------------------------------------------- #
+#  consume_prefetched_kv: miss returns None, stale sweep
+# --------------------------------------------------------------------------- #
+
+
+def test_start_prefetch_sweeps_orphan():
+    # An orphan prefetch (e.g. aborted) is dropped when the next prefetch starts.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-orphan")
+    _seed_payload(sender, "rid-next")
+    receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO})
+    receiver.consume_prefetched_kv(_req("rid-orphan"))  # complete the future
+    receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO})  # re-add as leftover
+    receiver.start_prefetch({"request_id": "rid-next", "kv_sender_info": _SENDER_INFO})
+    assert "rid-orphan" not in receiver._prefetch_futures
+    assert "rid-next" in receiver._prefetch_futures
+
+
+def test_shutdown_prefetch_clears_state():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-sd")
+    receiver.start_prefetch({"request_id": "rid-sd", "kv_sender_info": _SENDER_INFO})
+    receiver.shutdown_prefetch()
+    assert receiver._prefetch_futures == {}
+    assert receiver._prefetch_executor is None
+
+
+# --------------------------------------------------------------------------- #
+#  Correctness: per-call sender_info must not touch instance state
+# --------------------------------------------------------------------------- #
+
+
+def test_receive_with_sender_info_does_not_mutate_instance_state():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-si")
+    assert receiver._sender_base_host is None
+    assert receiver._sender_base_zmq_port is None
+
+    receiver.receive_kv_cache_for_request(
+        "rid-si",
+        target_device=None,
+        sender_info={"host": "1.2.3.4", "zmq_port": 50051},
+    )
+    # The per-call endpoint must not leak into instance fields.
+    assert receiver._sender_base_host is None
+    assert receiver._sender_base_zmq_port is None
+
+
+# --------------------------------------------------------------------------- #
+#  Receive-role classification (recv_role) and the follower prefetch guard
+# --------------------------------------------------------------------------- #
+
+import vllm_omni.diffusion.distributed.parallel_state as ps  # noqa: E402
+
+
+def _patch_topo(monkeypatch, *, world_size, world_rank=0, cfg_size=1, cfg_rank=0, sp_size=1, sp_rank=0):
+    monkeypatch.setattr(ps, "get_world_group", lambda: SimpleNamespace(world_size=world_size, rank_in_group=world_rank))
+    monkeypatch.setattr(ps, "get_classifier_free_guidance_world_size", lambda: cfg_size)
+    monkeypatch.setattr(ps, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+    monkeypatch.setattr(ps, "get_cfg_group", lambda: None)
+    monkeypatch.setattr(ps, "get_sequence_parallel_world_size", lambda: sp_size)
+    monkeypatch.setattr(ps, "get_sequence_parallel_rank", lambda: sp_rank)
+    monkeypatch.setattr(ps, "get_sp_group", lambda: None)
+
+
+def _set_tp(mgr, src, tgt):
+    mgr._tp_topo = SimpleNamespace(source_tp_size=src, target_tp_size=tgt, local_rank=0)
+
+
+def _role(mgr):
+    """Recompute the receive role from the (patched) parallel state, bypassing the cache."""
+    mgr._topo_config = None
+    return mgr.topo_config.role
+
+
+def test_recv_role_world1_is_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 1, 1)
+    _patch_topo(monkeypatch, world_size=1)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_recv_role_pure_tp_is_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=2)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_recv_role_cfg_leader_and_follower(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_sp_leader_and_follower(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_tp_and_sp_leader_per_tp_rank(monkeypatch):
+    # TP=2 + SP=2: sp_rank 0 is leader (pulls), sp_rank 1 is follower (gets broadcast).
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_hsdp_world_broadcast(monkeypatch):
+    # TP inactive, world > 1 (HSDP): rank-0 leader, others follower.
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 1, 1)
+    _patch_topo(monkeypatch, world_size=2, world_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=2, world_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_uninitialized_defaults_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+
+    def _boom():
+        raise AssertionError("world group is not initialized")
+
+    monkeypatch.setattr(ps, "get_world_group", _boom)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_start_prefetch_skips_follower(monkeypatch):
+    # A follower never pulls, so start_prefetch must not submit.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-fol")
+    _set_tp(receiver, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1)
+    receiver.start_prefetch({"request_id": "rid-fol", "kv_sender_info": _SENDER_INFO})
+    assert receiver._prefetch_futures == {}
+
+
+def test_start_prefetch_submits_for_owner(monkeypatch):
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-own")
+    _set_tp(receiver, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0)
+    receiver.start_prefetch({"request_id": "rid-own", "kv_sender_info": _SENDER_INFO})
+    assert "rid-own" in receiver._prefetch_futures
+
+
+def test_consume_then_apply_attaches_payload():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-apply")
+    receiver.start_prefetch({"request_id": "rid-apply", "kv_sender_info": _SENDER_INFO})
+    data, _ = receiver.consume_prefetched_kv(_req("rid-apply"))
+    assert data is not None
+
+    req = OmniDiffusionRequest(prompts=["p"], sampling_params=OmniDiffusionSamplingParams(), request_id="rid-apply")
+    # Mirror consume_and_distribute_kv_cache's LOCAL apply path (CPU: record_stream no-op).
+    receiver._record_stream_for_prefetched(data)
+    receiver.apply_kv_cache_to_request(req, data)
+    assert req.past_key_values is not None

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,7 +1,6 @@
 import json
 import struct
 
-import numpy as np
 import pytest
 import torch
 
@@ -206,17 +205,6 @@ def test_update_sender_info_uses_configured_source_stage():
     assert manager.config.connector_config["sender_zmq_port"] == 50152
 
 
-def test_clone_received_payload_tensors_breaks_buffer_alias():
-    payload, key_tensor = _make_serialized_payload()
-    raw = np.frombuffer(bytearray(payload), dtype=np.uint8)
-    data = KVCacheTransferData.from_bytes(memoryview(raw))
-
-    OmniKVTransferManager._clone_received_payload_tensors(data)
-    raw[:] = 0
-
-    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
-
-
 def test_receive_kv_cache_uses_exponential_backoff(monkeypatch):
     config = OmniKVCacheConfig(
         connector_config={"type": "mock"},
@@ -526,6 +514,57 @@ def test_integration_flow(common_constants):
     assert req.kv_metadata["seq_len"] == 10
 
 
+def test_deferred_release_isolates_data_from_pool_buffer():
+    payload, key_tensor = _make_serialized_payload()
+
+    # ManagedBuffer-like object whose release() zeros storage (simulates reuse).
+    raw_array = bytearray(payload)
+    buf_tensor = torch.frombuffer(raw_array, dtype=torch.uint8)
+
+    class _CorruptingManagedBuffer:
+        def __init__(self, storage, tensor):
+            self._storage = storage
+            self._tensor = tensor
+
+        @property
+        def tensor(self):
+            return self._tensor
+
+        def release(self):
+            for i in range(len(self._storage)):
+                self._storage[i] = 0
+
+    managed_buf = _CorruptingManagedBuffer(raw_array, buf_tensor)
+
+    class _ManagedBufConnector:
+        def __init__(self, buf, from_stage, to_stage, req_id):
+            self._buf = buf
+            self._key = f"{from_stage}->{to_stage}:omni_{from_stage}_to_{to_stage}_kv_cache_{req_id}"
+
+        def get(self, from_stage, to_stage, get_key, metadata=None):
+            if f"{from_stage}->{to_stage}:{get_key}" == self._key:
+                return self._buf, len(self._buf._storage)
+            return None
+
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=1.0,
+    )
+    manager = OmniKVTransferManager(config)
+    req_id = "req-payload"
+    manager._connector = _ManagedBufConnector(managed_buf, "sender", "receiver", req_id)
+
+    data, size = manager.receive_kv_cache_for_request(req_id, target_device=torch.device("cpu"))
+
+    assert data is not None
+    assert size > 0
+    # Returned data must survive the pool buffer being zeroed by release().
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
 def test_manager_extraction_no_connector(kv_config, common_constants):
     """Test extraction when connector is unavailable (should still return IDs)."""
     block_size = common_constants["block_size"]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -586,8 +586,17 @@ class TestDistributedReceive:
                 return_value=0,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once()
         assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
@@ -642,6 +651,15 @@ class TestDistributedReceive:
                 return_value=1,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
         ):
             assert mgr.receive_multi_kv_cache_distributed(req) is True
 
@@ -666,7 +684,27 @@ class TestDistributedReceive:
         world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
         mgr.receive_multi_kv_cache = MagicMock(return_value=True)
 
-        with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group):
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
+        ):
             assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
@@ -722,7 +760,7 @@ class TestDistributedReceive:
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         # Verify behavior based on role
         if is_owner:
@@ -817,7 +855,7 @@ class TestDistributedReceive:
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         # Verify behavior based on role
         if role == "owner":

tests/e2e/offline_inference/test_moss_tts.py

@@ -88,6 +88,7 @@ def _get_test_config() -> str:
 # ---------------------------------------------------------------------------
 
 pytestmark = [
+    pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
     pytest.mark.full_model,
     pytest.mark.tts,
     pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),

tests/e2e/offline_inference/test_moss_tts_realtime.py

@@ -95,6 +95,7 @@ def _get_test_config() -> str:
 # ---------------------------------------------------------------------------
 
 pytestmark = [
+    pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
     pytest.mark.full_model,
     pytest.mark.tts,
     pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -63,8 +63,6 @@ test_params = [
             use_stage_cli=True,
             server_args=[
                 "--no-async-chunk",
-                "--stage-overrides",
-                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
             ],
         ),
         id="default",

tests/entrypoints/openai_api/test_qwen3_tts_websocket.py

@@ -113,9 +113,9 @@ class TestQwen3TTSWebSocket:
         session_done = result["session_done"]
 
         assert session_done is not None
-        assert session_done["total_sentences"] == 2
-        assert len(starts) == 2
-        assert len(dones) == 2
+        assert session_done["total_sentences"] == 1
+        assert len(starts) == 1
+        assert len(dones) == 1
 
         for idx, start in enumerate(starts):
             assert start["type"] == "audio.start"

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -60,6 +60,7 @@ class TestStreamingSpeechWebSocket:
             with client.websocket_connect("/v1/audio/speech/stream") as ws:
                 ws.send_json({"type": "session.config", "voice": "Vivian"})
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 start = ws.receive_json()
                 assert start["type"] == "audio.start"
@@ -73,7 +74,6 @@ class TestStreamingSpeechWebSocket:
                 done = ws.receive_json()
                 assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False}
 
-                ws.send_json({"type": "input.done"})
                 session_done = ws.receive_json()
                 assert session_done == {"type": "session.done", "total_sentences": 1}
 
@@ -113,6 +113,7 @@ class TestStreamingSpeechWebSocket:
                     }
                 )
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 start = ws.receive_json()
                 assert start["type"] == "audio.start"
@@ -126,7 +127,6 @@ class TestStreamingSpeechWebSocket:
                 done = ws.receive_json()
                 assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False}
 
-                ws.send_json({"type": "input.done"})
                 assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
 
         assert len(captured_requests) == 1
@@ -150,6 +150,7 @@ class TestStreamingSpeechWebSocket:
                     }
                 )
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 error = ws.receive_json()
                 assert error["type"] == "error"
@@ -199,6 +200,7 @@ class TestStreamingSpeechWebSocket:
                     }
                 )
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 start = ws.receive_json()
                 assert start["type"] == "audio.start"
@@ -279,6 +281,7 @@ class TestStreamingSpeechWebSocket:
                     }
                 )
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 assert ws.receive_json()["type"] == "audio.start"
                 # Real-time audio chunk (timestamps null), then the timestamp frame.
@@ -342,16 +345,19 @@ class TestStreamingSpeechWebSocket:
 
         assert speech_service._generate_audio_bytes.await_count == 0
 
-    def test_multiple_sentences_increment_indices(self, mocker: MockerFixture):
+    def test_multiple_sentences_are_buffered_into_one_request(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
 
         with TestClient(app) as client:
             with client.websocket_connect("/v1/audio/speech/stream") as ws:
                 ws.send_json({"type": "session.config", "voice": "Vivian"})
-                ws.send_json({"type": "input.text", "text": "First sentence. Second sentence. "})
+                ws.send_json({"type": "input.text", "text": "First sentence. "})
+                ws.send_json({"type": "input.text", "text": "Second sentence. "})
+                ws.send_json({"type": "input.done"})
 
-                first_start = ws.receive_json()
-                assert first_start["sentence_index"] == 0
+                start = ws.receive_json()
+                assert start["sentence_index"] == 0
+                assert start["sentence_text"] == "First sentence. Second sentence."
                 ws.receive_bytes()
                 assert ws.receive_json() == {
                     "type": "audio.done",
@@ -359,19 +365,7 @@ class TestStreamingSpeechWebSocket:
                     "total_bytes": 36,
                     "error": False,
                 }
-
-                second_start = ws.receive_json()
-                assert second_start["sentence_index"] == 1
-                ws.receive_bytes()
-                assert ws.receive_json() == {
-                    "type": "audio.done",
-                    "sentence_index": 1,
-                    "total_bytes": 36,
-                    "error": False,
-                }
-
-                ws.send_json({"type": "input.done"})
-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 2}
+                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
 
     def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
@@ -385,6 +379,7 @@ class TestStreamingSpeechWebSocket:
                 assert error == {"type": "error", "message": "Unknown message type: unknown"}
 
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
                 assert ws.receive_json()["type"] == "audio.start"
                 ws.receive_bytes()
                 assert ws.receive_json() == {
@@ -394,7 +389,6 @@ class TestStreamingSpeechWebSocket:
                     "error": False,
                 }
 
-                ws.send_json({"type": "input.done"})
                 assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
 
     def test_config_timeout_closes_session(self, mocker: MockerFixture):
@@ -419,12 +413,12 @@ class TestStreamingSpeechWebSocket:
             with client.websocket_connect("/v1/audio/speech/stream") as ws:
                 ws.send_json({"type": "session.config", "voice": "Vivian"})
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 assert ws.receive_json()["type"] == "audio.start"
                 assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"}
                 assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True}
 
-                ws.send_json({"type": "input.done"})
                 assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
 
     def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture):
@@ -453,6 +447,7 @@ class TestStreamingSpeechWebSocket:
                     }
                 )
                 ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
 
                 assert ws.receive_json()["type"] == "audio.start"
                 assert ws.receive_bytes() == b"\x01\x02"
@@ -467,7 +462,6 @@ class TestStreamingSpeechWebSocket:
                     "error": True,
                 }
 
-                ws.send_json({"type": "input.done"})
                 assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
 
     def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture):

tests/entrypoints/openai_api/test_text_splitter.py

@@ -1,253 +0,0 @@
-"""Tests for SentenceSplitter used in streaming TTS input."""
-
-import pytest
-
-from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class TestSentenceSplitterEnglish:
-    """Tests for English sentence splitting."""
-
-    def test_single_sentence_no_boundary(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Hello world")
-        assert result == []
-        assert splitter.buffer == "Hello world"
-
-    def test_single_sentence_with_boundary(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Hello world. How are you?")
-        assert len(result) == 1
-        assert result[0] == "Hello world."
-
-    def test_multiple_sentences(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Hello. How are you? I am fine! ")
-        assert len(result) == 3
-        assert result[0] == "Hello."
-        assert result[1] == "How are you?"
-        assert result[2] == "I am fine!"
-
-    def test_exclamation_mark(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Wow, that is great! Tell me more.")
-        assert len(result) == 1
-        assert result[0] == "Wow, that is great!"
-
-    def test_question_mark(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Can you hear me? I hope so.")
-        assert len(result) == 1
-        assert result[0] == "Can you hear me?"
-
-
-class TestSentenceSplitterChinese:
-    """Tests for CJK sentence splitting."""
-
-    def test_chinese_period(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("你好世界。你好吗")
-        assert len(result) == 1
-        assert result[0] == "你好世界。"
-
-    def test_chinese_exclamation(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("太好了！谢谢你")
-        assert len(result) == 1
-        assert result[0] == "太好了！"
-
-    def test_chinese_question(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("你是谁？我是小明")
-        assert len(result) == 1
-        assert result[0] == "你是谁？"
-
-    def test_chinese_comma_no_split(self):
-        """Chinese commas are clause-level and should not trigger a split."""
-        splitter = SentenceSplitter()
-        result = splitter.add_text("你好，世界")
-        assert result == []
-        assert splitter.buffer == "你好，世界"
-
-    def test_chinese_semicolon_no_split(self):
-        """Chinese semicolons are clause-level and should not trigger a split."""
-        splitter = SentenceSplitter()
-        result = splitter.add_text("第一点；第二点")
-        assert result == []
-        assert splitter.buffer == "第一点；第二点"
-
-    def test_chinese_multiple(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("你好！你好吗？我很好。")
-        assert len(result) == 3
-        assert result[0] == "你好！"
-        assert result[1] == "你好吗？"
-        assert result[2] == "我很好。"
-
-
-class TestSentenceSplitterMixed:
-    """Tests for mixed-language sentence splitting."""
-
-    def test_mixed_english_chinese(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Hello世界。How are you? ")
-        assert len(result) == 2
-        assert result[0] == "Hello世界。"
-        assert result[1] == "How are you?"
-
-
-class TestSentenceSplitterIncremental:
-    """Tests for incremental (multi-chunk) text input."""
-
-    def test_accumulation_across_chunks(self):
-        splitter = SentenceSplitter()
-        # First chunk: no boundary
-        result1 = splitter.add_text("Hello ")
-        assert result1 == []
-
-        # Second chunk: completes a sentence
-        result2 = splitter.add_text("world. How")
-        assert len(result2) == 1
-        assert result2[0] == "Hello world."
-        assert splitter.buffer == "How"
-
-    def test_word_by_word(self):
-        splitter = SentenceSplitter()
-        words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."]
-        all_sentences = []
-        for word in words:
-            all_sentences.extend(splitter.add_text(word))
-
-        assert len(all_sentences) == 1
-        assert all_sentences[0] == "Hello, how are you?"
-        # "I am fine." stays in buffer (no trailing whitespace after period)
-
-    def test_three_chunks(self):
-        splitter = SentenceSplitter()
-        splitter.add_text("The quick brown ")
-        splitter.add_text("fox jumps. ")
-        result = splitter.add_text("Over the lazy dog. ")
-        # "The quick brown fox jumps." should have been returned on second chunk
-        # "Over the lazy dog." on third chunk
-        assert len(result) == 1
-        assert result[0] == "Over the lazy dog."
-
-
-class TestSentenceSplitterFlush:
-    """Tests for flush behavior."""
-
-    def test_flush_returns_remaining(self):
-        splitter = SentenceSplitter()
-        splitter.add_text("Hello world")
-        result = splitter.flush()
-        assert result == "Hello world"
-        assert splitter.buffer == ""
-
-    def test_flush_empty_buffer(self):
-        splitter = SentenceSplitter()
-        result = splitter.flush()
-        assert result is None
-
-    def test_flush_after_sentence(self):
-        splitter = SentenceSplitter()
-        splitter.add_text("Hello world. Remaining text")
-        result = splitter.flush()
-        assert result == "Remaining text"
-
-    def test_flush_whitespace_only(self):
-        splitter = SentenceSplitter()
-        splitter.add_text("Hello. ")
-        # "Hello." extracted, buffer is " "
-        result = splitter.flush()
-        # Whitespace-only should return None
-        assert result is None
-
-    def test_flush_clears_buffer(self):
-        splitter = SentenceSplitter()
-        splitter.add_text("some text")
-        splitter.flush()
-        assert splitter.buffer == ""
-        # Second flush should return None
-        assert splitter.flush() is None
-
-
-class TestSentenceSplitterEdgeCases:
-    """Edge case tests."""
-
-    def test_empty_input(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("")
-        assert result == []
-        assert splitter.buffer == ""
-
-    def test_none_like_empty(self):
-        """Empty string should not affect buffer."""
-        splitter = SentenceSplitter()
-        splitter.add_text("Hello")
-        splitter.add_text("")
-        assert splitter.buffer == "Hello"
-
-    def test_only_punctuation(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text(". ")
-        # "." is 1 char, below default min_sentence_length of 2
-        # It will be carried forward
-        assert result == []
-
-    def test_min_sentence_length(self):
-        splitter = SentenceSplitter(min_sentence_length=10)
-        result = splitter.add_text("Hi. Hello world. ")
-        # "Hi." is 3 chars (< 10), so it gets carried to "Hello world."
-        assert len(result) == 1
-        assert "Hi." in result[0]
-        assert "Hello world." in result[0]
-
-    def test_short_segments_are_carried_until_long_enough(self):
-        splitter = SentenceSplitter(min_sentence_length=10)
-        result = splitter.add_text("Hi. Ok. Hello there. ")
-        assert result == ["Hi.Ok.Hello there."]
-        assert splitter.buffer == ""
-
-    def test_min_sentence_length_zero(self):
-        splitter = SentenceSplitter(min_sentence_length=0)
-        result = splitter.add_text("A. B. ")
-        assert len(result) == 2
-
-    def test_no_boundary_then_flush(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Hello world how are you")
-        assert result == []
-        flushed = splitter.flush()
-        assert flushed == "Hello world how are you"
-
-    def test_consecutive_punctuation(self):
-        splitter = SentenceSplitter()
-        result = splitter.add_text("Really?! Yes, really. ")
-        assert len(result) >= 1
-
-    def test_reuse_after_flush(self):
-        """Splitter can be reused after flush."""
-        splitter = SentenceSplitter()
-        splitter.add_text("First session.")
-        splitter.flush()
-
-        result = splitter.add_text("Second session. More text")
-        assert len(result) == 1
-        assert result[0] == "Second session."
-        assert splitter.buffer == "More text"
-
-
-class TestSentenceSplitterBufferLimit:
-    """Tests for buffer overflow protection."""
-
-    def test_buffer_overflow_raises(self):
-        from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE
-
-        splitter = SentenceSplitter()
-        # Fill buffer just under the limit
-        splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1))
-        # One more char should exceed the limit
-        with pytest.raises(ValueError, match="exceeded maximum size"):
-            splitter.add_text("xx")

tests/entrypoints/openai_api/test_video_server.py

@@ -11,6 +11,7 @@ import json
 import os
 import threading
 import time
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 
 import numpy as np
@@ -99,6 +100,14 @@ class BlockingVideoHandler:
             raise
 
 
+class FakeServerSocket:
+    def __init__(self):
+        self.closed = False
+
+    def close(self):
+        self.closed = True
+
+
 @pytest.fixture(autouse=True)
 def isolated_video_backends(tmp_path, monkeypatch):
     """Use isolated in-memory metadata and local storage for each test."""
@@ -111,6 +120,91 @@ def isolated_video_backends(tmp_path, monkeypatch):
     return store, tasks, storage
 
 
+@pytest.mark.asyncio
+async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch):
+    events: list[str] = []
+    serve_started = asyncio.Event()
+    http_shutdown = asyncio.Event()
+    engine_context_exited = asyncio.Event()
+    sock = FakeServerSocket()
+
+    class FakeEngine:
+        stage_configs = []
+
+        async def get_supported_tasks(self):
+            return ("generate",)
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*args, **kwargs):
+        del args, kwargs
+        events.append("engine_enter")
+        try:
+            yield FakeEngine()
+        finally:
+            events.append("engine_exit")
+            engine_context_exited.set()
+
+    async def fake_serve_http(*args, **kwargs):
+        del args, kwargs
+        events.append("serve_http")
+        serve_started.set()
+
+        async def wait_for_shutdown():
+            await http_shutdown.wait()
+            events.append("http_shutdown")
+
+        return asyncio.create_task(wait_for_shutdown())
+
+    async def fake_storage_start():
+        events.append("storage_start")
+
+    async def fake_get_vllm_config(engine_client):
+        del engine_client
+        return None
+
+    async def fake_init_app_state(engine_client, state, args):
+        del engine_client, state, args
+        events.append("init_app_state")
+
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: FastAPI())
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+    monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
+    monkeypatch.setattr(api_server, "_get_vllm_config", fake_get_vllm_config)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_init_app_state)
+    monkeypatch.setattr(api_server, "get_uvicorn_log_config", lambda args: None)
+
+    args = SimpleNamespace(
+        tool_parser_plugin="",
+        reasoning_parser_plugin="",
+        reasoning_parser=None,
+        structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+        enable_ssl_refresh=False,
+        host="127.0.0.1",
+        port=0,
+        uvicorn_log_level="info",
+        disable_uvicorn_access_log=True,
+        ssl_keyfile=None,
+        ssl_certfile=None,
+        ssl_ca_certs=None,
+        ssl_cert_reqs=None,
+        ssl_ciphers=None,
+        h11_max_incomplete_event_size=None,
+        h11_max_header_count=None,
+    )
+
+    worker_task = asyncio.create_task(api_server.omni_run_server_worker("127.0.0.1:0", sock, args))
+    await asyncio.wait_for(serve_started.wait(), timeout=2)
+
+    assert not engine_context_exited.is_set()
+
+    http_shutdown.set()
+    await asyncio.wait_for(worker_task, timeout=2)
+
+    assert sock.closed
+    assert events.index("http_shutdown") < events.index("engine_exit")
+
+
 @pytest.fixture
 def test_client():
     app = FastAPI()

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -137,8 +137,12 @@ def test_ming_compute_logits_falls_back_to_dummy_token_id():
 
 def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids():
     class FakeBackbone:
-        def __call__(self, **kwargs):
-            return kwargs["inputs_embeds"]
+        # MingLLMModel.forward calls the backbone positionally
+        # (input_ids, positions, intermediate_tensors, inputs_embeds) so the
+        # same signature works for both Qwen2Model (``positions``) and
+        # BailingMoeModel (``position_ids``).
+        def __call__(self, input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs):
+            return inputs_embeds
 
     model = _make_ming_logits_model()
     model.model = FakeBackbone()

tests/test_config_factory.py

@@ -655,6 +655,16 @@ class TestPipelineRegistration:
         assert isinstance(pipeline_cfg, PipelineConfig)
         assert pipeline_cfg.model_type == resolved_type
 
+    def test_resolve_when_autodetect_resolves_none(self):
+        """Regression test for: https://github.com/vllm-project/vllm-omni/issues/4726"""
+        deploy_path = get_deploy_config_path("ming_tts.yaml")
+        resolved_config = StageConfigFactory.create_from_model(
+            model="inclusionAI/Ming-omni-tts-0.5B",
+            deploy_config_path=deploy_path,
+        )
+        assert resolved_config is not None
+        assert len(resolved_config) > 0
+
 
 class TestResolveScheduler:
     def test_all_execution_types_handled(self):

tests/unit/mammoth_moda2/test_mammoth_moda2_config.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for MammothModa2 config classes.
+
+These tests verify the correct behavior of the Mammothmoda2Config class,
+particularly the get_text_config() method which is called during parent
+class initialization (PretrainedConfig.__init__ -> validation -> get_text_config).
+
+Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config'
+"""
+
+import pytest
+
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
+    Mammothmoda2Config,
+    Mammothmoda2Qwen2_5_VLConfig,
+    Mammothmoda2Qwen2_5_VLTextConfig,
+)
+
+
+@pytest.mark.cpu
+class TestMammothmoda2Config:
+    """Tests for Mammothmoda2Config class."""
+
+    def test_get_text_config_returns_nested_text_config(self):
+        """
+        Verify that get_text_config() returns the nested text_config from llm_config.
+
+        This is a regression test for the AttributeError that occurred when
+        get_text_config() was called during parent class initialization before
+        llm_config was set.
+        """
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+
+        assert config.llm_config is not None
+        assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig)
+
+        text_config = config.get_text_config()
+        assert text_config is not None
+        assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+        # The text_config should be the same object as llm_config.text_config
+        assert text_config is config.llm_config.text_config
+
+    def test_get_text_config_returns_none_when_llm_config_is_none(self):
+        """Verify that get_text_config() returns None when llm_config is None."""
+        config = Mammothmoda2Config(llm_config=None)
+
+        assert config.llm_config is None
+        assert config.get_text_config() is None
+
+    def test_init_sets_llm_config_before_parent_init(self):
+        """
+        Verify that llm_config is set before parent __init__ is called.
+
+        This is critical because the parent PretrainedConfig.__init__ calls
+        validation which in turn calls get_text_config(). If llm_config is not
+        set before super().__init__(), an AttributeError will be raised.
+        """
+        # This should not raise an AttributeError
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+        assert config.llm_config is not None
+
+
+@pytest.mark.cpu
+class TestMammothmoda2Qwen2_5_VLConfig:
+    """Tests for Mammothmoda2Qwen2_5_VLConfig class."""
+
+    def test_has_text_config_attribute(self):
+        """Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        assert hasattr(config, "text_config")
+        assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+    def test_text_config_has_expected_attributes(self):
+        """Verify that the nested text_config has expected attributes for token validation."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        text_config = config.text_config
+
+        assert hasattr(text_config, "vocab_size")
+        assert hasattr(text_config, "image_token_id")
+        assert hasattr(text_config, "video_token_id")

vllm_omni/benchmarks/patch/patch.py

@@ -52,6 +52,8 @@ from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
 from vllm_omni.metrics import definitions as defs
 
 _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
+_AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
+_AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
 _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"})
 _PRINT_STAGE = False
@@ -109,6 +111,32 @@ def _audio_continuity_threshold_s() -> float:
     return max(value, 0.0)
 
 
+def _audio_pcm_format() -> tuple[int, int]:
+    """Return ``(sample_rate, channels)`` of the streamed PCM response.
+
+    Defaults to 24 kHz mono (Qwen3-TTS / VoxCPM2 / most MOSS-TTS variants).
+    Override for models with a different codec output format (e.g.
+    MOSS-TTS-Local-Transformer-v1.5, which streams 48 kHz stereo PCM) via
+    ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / ``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` -
+    getting this wrong silently skews audio_duration/audio_rtf/underrun stats.
+    """
+    sample_rate = 24000
+    channels = 1
+    raw_sr = os.environ.get(_AUDIO_SAMPLE_RATE_ENV)
+    if raw_sr:
+        try:
+            sample_rate = int(raw_sr)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", _AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
+    raw_ch = os.environ.get(_AUDIO_CHANNELS_ENV)
+    if raw_ch:
+        try:
+            channels = int(raw_ch)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", _AUDIO_CHANNELS_ENV, raw_ch, channels)
+    return sample_rate, channels
+
+
 get_samples_old = datasets.get_samples
 
 _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni"
@@ -1133,10 +1161,10 @@ async def async_request_openai_audio_speech(
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
 
-    # PCM format: 16-bit signed, 24 kHz, mono
-    sample_rate = 24000
+    # PCM format: 16-bit signed; sample_rate/channels are model-dependent
+    # (see _audio_pcm_format docstring).
+    sample_rate, channels = _audio_pcm_format()
     sample_width = 2  # 16-bit = 2 bytes
-    channels = 1
 
     st = time.perf_counter()
     output.start_time = st

vllm_omni/config/__init__.py

@@ -2,9 +2,29 @@
 Configuration module for vLLM-Omni.
 """
 
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.config.omni_config import (
+    BaseVllmOmniStageConfig,
+    OmniStageCacheConfig,
+    OmniStageConnectorConfig,
+    OmniStageDiffusionParallelConfig,
+    OmniStageLoadConfig,
+    OmniStageModelConfig,
+    OmniStageParallelConfig,
+    OmniStageRuntimeConfig,
+    OmniStageSchedulerConfig,
+    StageConfigType,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+    VllmOmniDiffusionStageConfig,
+    VllmOmniGenerationStageConfig,
+    VllmOmniOrchestratorConfig,
+)
+from vllm_omni.config.pipeline_registry import register_pipeline
 from vllm_omni.config.stage_config import (
+    PIPELINE_WIDE_ENGINE_FIELDS,
     DeployConfig,
     PipelineConfig,
     StageConfig,
@@ -23,17 +43,40 @@ from vllm_omni.config.yaml_util import (
 )
 
 __all__ = [
-    "OmniModelConfig",
+    # Legacy model-level configs.
     "LoRAConfig",
+    "OmniModelConfig",
+    # Structured Omni config entry points.
+    "VllmOmniConfig",
+    "BaseVllmOmniStageConfig",
+    "VllmOmniARStageConfig",
+    "VllmOmniGenerationStageConfig",
+    "VllmOmniDiffusionStageConfig",
+    "StageConfigType",
+    # Structured Omni sub-configs.
+    "OmniStageCacheConfig",
+    "OmniStageConnectorConfig",
+    "OmniStageDiffusionParallelConfig",
+    "OmniStageLoadConfig",
+    "OmniStageModelConfig",
+    "VllmOmniOrchestratorConfig",
+    "OmniStageParallelConfig",
+    "OmniStageRuntimeConfig",
+    "OmniStageSchedulerConfig",
+    # Legacy pipeline/stage deploy config surface.
+    "PIPELINE_WIDE_ENGINE_FIELDS",
+    "DeployConfig",
+    "PipelineConfig",
     "StageConfig",
+    "StageConfigFactory",
+    "StageDeployConfig",
     "StageType",
     "StageExecutionType",
     "StagePipelineConfig",
-    "PipelineConfig",
-    "StageDeployConfig",
-    "DeployConfig",
     "load_deploy_config",
     "merge_pipeline_deploy",
+    "register_pipeline",
+    # YAML utility helpers.
     "create_config",
     "load_yaml_config",
     "merge_configs",

vllm_omni/config/config_factory.py

@@ -9,7 +9,7 @@ from dataclasses import asdict
 from pathlib import Path
 from typing import Any
 
-from transformers import PreTrainedConfig
+from transformers import PretrainedConfig
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config
 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
@@ -123,6 +123,24 @@ class StageConfigFactory:
                             cli_overrides,
                             deploy_config_path,
                         )
+
+        # --- Explicit deploy-config pipeline ---
+        # When auto-detection above resolves nothing (generic/missing HF model_type
+        # and no matching architecture), honor an explicit pipeline key in the deploy config
+        if deploy_config_path is not None:
+            deploy_path = Path(deploy_config_path)
+            if deploy_path.exists():
+                deploy_cfg = load_deploy_config(deploy_path)
+                if deploy_cfg.pipeline:
+                    pipeline_cfg = cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
+                    if pipeline_cfg is not None:
+                        return cls._create_from_registry(
+                            pipeline_cfg.model_type,
+                            pipeline_cfg,
+                            cli_overrides,
+                            deploy_config_path,
+                        )
+
         # Not in the pipeline registry — let the caller fall back to the
         # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
         return None
@@ -329,7 +347,7 @@ class StageConfigFactory:
         return build_stage_runtime_overrides(stage.stage_id, cli_overrides)
 
     @staticmethod
-    def resolve_pipeline_config(model_type: str, hf_config: PreTrainedConfig | None = None) -> PipelineConfig | None:
+    def resolve_pipeline_config(model_type: str, hf_config: PretrainedConfig | None = None) -> PipelineConfig | None:
         """Given a model type, resolve to the pipeline to be used. If the pipeline
         maps to a callable we resolve based on the HF config."""
         if model_type not in OMNI_PIPELINES:

vllm_omni/config/pipeline_registry.py

@@ -27,7 +27,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from typing import TypeAlias
 
-from transformers import PreTrainedConfig
+from transformers import PretrainedConfig
 from vllm.logger import init_logger
 
 from vllm_omni.config.stage_config import (
@@ -63,9 +63,16 @@ from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     MING_FLASH_OMNI_TTS_PIPELINE,
 )
-from vllm_omni.model_executor.models.ming_tts.pipeline import MING_TTS_PIPELINE
+from vllm_omni.model_executor.models.ming_tts.pipeline import (
+    MING_TTS_MOE_PIPELINE,
+    MING_TTS_PIPELINE,
+)
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
-from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_PIPELINE, MOSS_TTS_REALTIME_PIPELINE
+from vllm_omni.model_executor.models.moss_tts.pipeline import (
+    MOSS_TTS_LOCAL_PIPELINE,
+    MOSS_TTS_PIPELINE,
+    MOSS_TTS_REALTIME_PIPELINE,
+)
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
@@ -78,7 +85,7 @@ from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIP
 
 logger = init_logger(__name__)
 
-PipelineResolverFunc: TypeAlias = Callable[[PreTrainedConfig | None], PipelineConfig | None]
+PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineConfig | None]
 
 # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) ---
 OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
@@ -102,6 +109,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "cosyvoice3": COSYVOICE3_PIPELINE,
     "mimo_audio": MIMO_AUDIO_PIPELINE,
     "ming_tts": MING_TTS_PIPELINE,
+    "ming_tts_moe": MING_TTS_MOE_PIPELINE,
     "voxtral_tts": VOXTRAL_TTS_PIPELINE,
     "glm_tts": GLM_TTS_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,
@@ -112,6 +120,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
+    "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
     "minicpmo_4_5": MINICPMO_4_5_PIPELINE,
     "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
     "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,

vllm_omni/config/stage_config.py

@@ -690,6 +690,7 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = (
     "active_stream_window",
     "custom_voice_dir",
 )
+PIPELINE_WIDE_ENGINE_FIELDS = _PIPELINE_WIDE_ENGINE_FIELDS
 
 
 def _build_engine_args(

vllm_omni/deploy/hunyuan_image_3_moe.yaml

@@ -62,7 +62,7 @@ stages:
       max_tokens: 8192
       detokenize: true
       skip_special_tokens: false
-      include_stop_str_in_output: true,
+      include_stop_str_in_output: true
 
   - stage_id: 1
     max_num_seqs: 1
@@ -71,6 +71,8 @@ stages:
     distributed_executor_backend: "mp"
     omni_kv_config:
       need_recv_cache: true
+      enable_kv_async_prefetch: true
+      kv_prefetch_min_free_mem_ratio: 0.2
     parallel_config:
       tensor_parallel_size: 2
       enable_expert_parallel: true

vllm_omni/deploy/ming_tts_moe.yaml

@@ -0,0 +1,66 @@
+# Ming-omni-tts-16.8B-A3B (MoE) deploy: stage-0 bailing_moe LLM+flow -> stage-1 audio VAE.
+# Same two-stage pipeline as the dense ming_tts.yaml; the Stage-0 backbone is the
+# 16.8B-A3B bailing_moe MoE (~31GB bf16).
+#
+# pipeline: HF config reports model_type="bailingmm", which is not a registry key,
+# already declared by MING_TTS_MOE_PIPELINE for auto-detection.
+#
+# NOTE: ~31GB bf16 + the tiny audio VAE fit one 80GB+ GPU, so both stages colocate
+# on GPU0 (Stage-0 0.60 util ≈ 58GB + Stage-1 0.25 util ≈ 24GB ≈ 0.85, fits 97GB
+# H20). For smaller cards raise Stage-0 tensor_parallel_size or split the stages
+# onto separate devices. Multi-request batching is not yet validated (max_num_seqs: 1).
+pipeline: ming_tts_moe
+async_chunk: true
+# The MoE ships a custom tokenizer (tokenization_bailing.BailingTokenizer) that
+# must be executed to load. The model config still resolves to the registered
+# MingMoeConfig (vLLM _CONFIG_REGISTRY precedence over the repo auto_map).
+trust_remote_code: true
+dtype: bfloat16
+enable_prefix_caching: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      latent_chunk_size: 25
+      latent_left_context: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    tensor_parallel_size: 1
+    gpu_memory_utilization: 0.60
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    # Pin the unquantized MoE kernel (wer diff between vllm 0.22 vs 0.23)
+    moe_backend: triton
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 512
+      detokenize: true
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 1
+      detokenize: false

vllm_omni/deploy/moss_tts_local.yaml

@@ -0,0 +1,56 @@
+# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo)
+# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5
+# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ)
+# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every
+# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime.
+async_chunk: true
+
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.005
+      connector_get_max_wait_first_chunk: 1000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s chunks
+      codec_left_context_frames: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.60
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_1: shm
+    default_sampling_params:
+      temperature: 1.0
+      top_p: 0.95
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.15
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: shm
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/diffusion/attention/layer.py

@@ -230,6 +230,34 @@ class Attention(nn.Module):
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if torch.compiler.is_compiling() and is_forward_context_available():
+            od_config = get_forward_context().omni_diffusion_config
+            parallel_config = getattr(od_config, "parallel_config", None)
+            if getattr(parallel_config, "use_hsdp", False):
+                # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
+                # attention graph; otherwise scheduler dependency analysis can
+                # fail on the fused attention region.
+                return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
+
+        return self._forward_impl(query, key, value, attn_metadata)
+
+    @torch.compiler.disable
+    def _forward_hsdp_compile_boundary(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self._forward_impl(query, key, value, attn_metadata)
+
+    def _forward_impl(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         # Get the appropriate parallel strategy based on SP active state
         strategy = self._get_active_parallel_strategy()

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -248,20 +248,27 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
     """
     # Build DBCacheConfig with optional SCM support
     db_cache_config = _build_db_cache_config(cache_config)
+    calibrator_config = _resolve_calibrator_config(cache_config)
 
     if getattr(pipeline, "transformer_2", None) is None:
         logger.info("transformer_2 not found, enabling cache-dit for single transformer mode")
         cache_dit.enable_cache(
             BlockAdapter(
                 transformer=pipeline.transformer,
+                # For VACE, cache only the main denoising blocks. The
+                # conditioning branch (vace_blocks) has a different forward
+                # contract and produces per-step hints from the current latent
+                # plus vace_context; keeping it outside CacheDiT preserves the
+                # control signal while still accelerating the repeated backbone.
                 blocks=[pipeline.transformer.blocks],
                 forward_pattern=[ForwardPattern.Pattern_2],
                 params_modifiers=[
-                    ParamsModifier(cache_config=db_cache_config),
+                    ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
                 ],
                 has_separate_cfg=True,
             ),
             cache_config=db_cache_config,
+            calibrator_config=calibrator_config,
         )
         return build_cache_context_refresh(cache_config)
 
@@ -272,6 +279,9 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
                 pipeline.transformer_2,
             ],
             blocks=[
+                # See the single-transformer branch above: VACE conditioning
+                # blocks are intentionally recomputed each step and are not
+                # wrapped by CacheDiT's main-block Pattern_2 adapter.
                 pipeline.transformer.blocks,
                 pipeline.transformer_2.blocks,
             ],
@@ -286,17 +296,20 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
                         max_warmup_steps=cache_config.max_warmup_steps,
                         max_cached_steps=cache_config.max_cached_steps,
                     ),
+                    calibrator_config=calibrator_config,
                 ),
                 ParamsModifier(
                     cache_config=DBCacheConfig().reset(
                         max_warmup_steps=2,
                         max_cached_steps=20,
                     ),
+                    calibrator_config=calibrator_config,
                 ),
             ],
             has_separate_cfg=True,
         ),
         cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
     )
 
     refresh_trans_one = build_cache_context_refresh(cache_config)
@@ -965,6 +978,7 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
         "Wan22TI2VPipeline": enable_cache_for_wan22,
+        "Wan22VACEPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
     }

vllm_omni/diffusion/compile.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
+import torch
 import torch.nn as nn
 from vllm.logger import init_logger
 
@@ -27,7 +28,11 @@ def _matches_repeated_block(
     return len(parts) >= 2 and parts[-2] in repeated_block_attrs and parts[-1].isdigit()
 
 
-def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module:
+def regionally_compile(
+    model: nn.Module,
+    *compile_args: Any,
+    **compile_kwargs: Any,
+) -> nn.Module:
     """
     Apply regional compilation to a PyTorch model.
 
@@ -53,8 +58,9 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A
     compiled_region_count = 0
     for name, submod in model.named_modules():
         if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
-            # Compile this submodule
-            submod.compile(*compile_args, **compile_kwargs)
+            # Compile the block compute while keeping nn.Module.__call__ hooks
+            # outside the compiled graph.
+            submod.forward = torch.compile(submod.forward, *compile_args, **compile_kwargs)
             has_compiled_region = True
             compiled_region_count += 1
 

vllm_omni/diffusion/data.py

@@ -33,6 +33,62 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
+    """Normalize legacy diffusion kwargs before config construction."""
+    normalized = dict(kwargs)
+
+    # Backwards-compatibility: older callers may use a diffusion-specific
+    # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
+    if "static_lora_scale" in normalized:
+        if "lora_scale" not in normalized:
+            normalized["lora_scale"] = normalized["static_lora_scale"]
+        normalized.pop("static_lora_scale", None)
+
+    # Backwards-compatibility: map "quantization" to "quantization_config"
+    # so callers using the old field name still work.
+    if "quantization" in normalized and normalized.get("quantization_config", None) is None:
+        normalized["quantization_config"] = normalized.pop("quantization")
+    else:
+        normalized.pop("quantization", None)
+
+    # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
+    if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
+        normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
+    else:
+        normalized.pop("kv_cache_dtype", None)
+    if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
+        normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
+    else:
+        normalized.pop("kv_cache_skip_steps", None)
+    if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
+        normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
+    else:
+        normalized.pop("kv_cache_skip_layers", None)
+
+    # Handle "diffusion_attention_backend" shorthand: merge into
+    # diffusion_attention_config before field filtering.
+    diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
+    if diffusion_attn_backend is not None:
+        existing = normalized.get("diffusion_attention_config")
+        normalized["diffusion_attention_config"] = parse_attention_config(
+            existing,
+            attention_backend=diffusion_attn_backend,
+        )
+
+    # Check environment variable as fallback for cache_backend.
+    # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
+    if "cache_backend" not in normalized:
+        cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
+        normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"
+
+    # Convert optional YAML null values to empty containers.
+    for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
+        if key in normalized and normalized[key] is None:
+            normalized[key] = {}
+
+    return normalized
+
+
 def parse_kv_cache_skip_selector(
     selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -1083,55 +1139,7 @@ class OmniDiffusionConfig:
 
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
-        # Backwards-compatibility: older callers may use a diffusion-specific
-        # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale"
-        # before constructing the dataclass to avoid TypeError on unknown fields.
-        if "static_lora_scale" in kwargs:
-            if "lora_scale" not in kwargs:
-                kwargs["lora_scale"] = kwargs["static_lora_scale"]
-            kwargs.pop("static_lora_scale", None)
-
-        # Backwards-compatibility: map "quantization" to "quantization_config"
-        # so callers using the old field name still work.
-        if "quantization" in kwargs and kwargs.get("quantization_config", None) is None:
-            kwargs["quantization_config"] = kwargs.pop("quantization")
-        else:
-            kwargs.pop("quantization", None)
-
-        # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
-        if kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in kwargs:
-            kwargs["diffusion_kv_cache_dtype"] = kwargs.pop("kv_cache_dtype")
-        else:
-            kwargs.pop("kv_cache_dtype", None)
-        if kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in kwargs:
-            kwargs["diffusion_kv_cache_skip_steps"] = kwargs.pop("kv_cache_skip_steps")
-        else:
-            kwargs.pop("kv_cache_skip_steps", None)
-        if kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in kwargs:
-            kwargs["diffusion_kv_cache_skip_layers"] = kwargs.pop("kv_cache_skip_layers")
-        else:
-            kwargs.pop("kv_cache_skip_layers", None)
-
-        # Handle "diffusion_attention_backend" shorthand: merge into
-        # diffusion_attention_config before field filtering.
-        diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None)
-        if diffusion_attn_backend is not None:
-            existing = kwargs.get("diffusion_attention_config")
-            kwargs["diffusion_attention_config"] = parse_attention_config(
-                existing,
-                attention_backend=diffusion_attn_backend,
-            )
-
-        # Check environment variable as fallback for cache_backend
-        # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility
-        if "cache_backend" not in kwargs:
-            cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
-            kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-
-        # Convert optional YAML null values to empty containers.
-        for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
-            if key in kwargs and kwargs[key] is None:
-                kwargs[key] = {}
+        kwargs = normalize_omni_diffusion_kwargs(kwargs)
 
         # Filter kwargs to only include valid fields
         valid_fields = {f.name for f in fields(cls)}

vllm_omni/diffusion/distributed/utils.py

@@ -14,3 +14,26 @@ def get_local_device() -> torch.device:
         return current_omni_platform.get_torch_device(current_omni_platform.current_device())
     local_rank = int(os.environ.get("LOCAL_RANK", 0))
     return current_omni_platform.get_torch_device(local_rank)
+
+
+def build_local_sp_padding_mask(
+    batch_size: int,
+    local_seq_len: int,
+    sp_original_seq_len: int | None,
+    sp_padding_size: int,
+    sequence_parallel_rank: int,
+    device,
+):
+    """Build a per-rank SP padding mask that matches the local shard shape."""
+    if sp_original_seq_len is None or sp_padding_size <= 0:
+        return None
+
+    shard_start = sequence_parallel_rank * local_seq_len
+    shard_end = shard_start + local_seq_len
+    valid_tokens = max(0, min(sp_original_seq_len, shard_end) - shard_start)
+
+    if valid_tokens >= local_seq_len:
+        return None
+
+    token_mask = torch.arange(local_seq_len, device=device) < valid_tokens
+    return token_mask.unsqueeze(0).expand(batch_size, -1)

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -321,7 +321,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         new_req = scheduler_output.scheduled_new_reqs[0]
         result = self.collective_rpc(
             "execute_model",
-            args=(new_req.req, self.od_config),
+            args=(new_req.req, self.od_config, getattr(scheduler_output, "kv_prefetch_jobs", None)),
             unique_reply_rank=0,
             exec_all_ranks=True,
         )

vllm_omni/diffusion/forward_context.py

@@ -4,6 +4,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import torch
 import vllm.ir
 from vllm.config import VllmConfig
 
@@ -88,6 +89,40 @@ def is_forward_context_available() -> bool:
     return _forward_context is not None
 
 
+def build_local_sp_padding_mask(
+    batch_size: int,
+    local_seq_len: int,
+    device,
+):
+    """Build a per-rank SP padding mask that matches the local shard shape.
+
+    Auto-padding is applied before sequence-parallel sharding, so attention on each
+    rank must receive a mask for its local shard, not for the global padded sequence.
+    """
+    if not is_forward_context_available():
+        return None
+
+    ctx = get_forward_context()
+    if ctx.sp_original_seq_len is None or ctx.sp_padding_size <= 0:
+        return None
+
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        get_sequence_parallel_rank,
+    )
+    from vllm_omni.diffusion.distributed.utils import (
+        build_local_sp_padding_mask as build_local_sp_padding_mask_for_rank,
+    )
+
+    return build_local_sp_padding_mask_for_rank(
+        batch_size=batch_size,
+        local_seq_len=local_seq_len,
+        sp_original_seq_len=ctx.sp_original_seq_len,
+        sp_padding_size=ctx.sp_padding_size,
+        sequence_parallel_rank=get_sequence_parallel_rank(),
+        device=device,
+    )
+
+
 def get_ulysses_mode(*, default: str = "strict") -> str:
     """Resolve the Ulysses-SP mode from the current ForwardContext.
 

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -560,6 +560,8 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsCompon
                     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
                     # Scheduler step
                     latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                    # Re-cast to transformer_dtype before cfg_group.broadcast to avoid dtype mismatches across ranks
+                    latents = latents.to(transformer_dtype)
 
                 # Broadcast updated latents to all ranks
                 cfg_group.broadcast(latents, src=0)

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -59,9 +59,19 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
         self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
 
-        llm_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0)
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        llm_hidden_size = int(
+            getattr(
+                getattr(self.config.llm_config, "text_config", None),
+                "hidden_size",
+                0,
+            )
+            or 0
+        )
         if llm_hidden_size <= 0:
-            raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size")
+            raise ValueError(
+                "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
+            )
         self._reinit_caption_embedder(llm_hidden_size)
 
         # Optional: image condition refiner (Q-Former)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -198,6 +198,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         attention_kwargs: dict[str, object],
         vace_context: torch.Tensor | None,
         vace_context_scale: float,
+        boundary_timestep: float | None = None,
     ) -> torch.Tensor:
         if attention_kwargs is None:
             attention_kwargs = {}
@@ -206,6 +207,11 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                 self._current_timestep = t
                 set_forward_context_denoise_step_idx(step_idx)
 
+                if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None:
+                    current_model = self.transformer_2
+                else:
+                    current_model = self.transformer if self.transformer is not None else self.transformer_2
+
                 latent_model_input = latents.to(dtype)
                 timestep = t.expand(latents.shape[0])
 
@@ -219,6 +225,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                     "vace_context": vace_context,
                     "vace_context_scale": vace_context_scale,
                     "return_dict": False,
+                    "current_model": current_model,
                 }
                 negative_kwargs = (
                     {
@@ -229,6 +236,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                         "vace_context": vace_context,
                         "vace_context_scale": vace_context_scale,
                         "return_dict": False,
+                        "current_model": current_model,
                     }
                     if do_true_cfg
                     else None
@@ -598,7 +606,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                 )
 
         num_reference_images = 0
-        if self.transformer.vace_patch_embedding is not None:
+        # Either expert may be the one loaded (boundary_ratio 0.0/1.0 skips one),
+        # so read it from whichever transformer exists.
+        active_transformer = self.transformer if self.transformer is not None else self.transformer_2
+        if active_transformer.vace_patch_embedding is not None:
             video, mask, ref_images_processed = self.preprocess_conditions(
                 video=source_video,
                 mask=source_mask,
@@ -640,6 +651,19 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
 
+        # MoE boundary: only the two-expert checkpoints (Wan2.2 A14B) have a
+        # transformer_2 to switch to. For single-expert VACE (e.g. Wan2.1) this
+        # stays None and diffuse() uses the same transformer for every step.
+        boundary_timestep = None
+        if self.transformer_2 is not None:
+            boundary_ratio = (
+                self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio
+            )
+            if boundary_ratio is None:
+                boundary_ratio = 0.875
+                logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875")
+            boundary_timestep = boundary_ratio * self.scheduler.config.num_train_timesteps
+
         latents = self.diffuse(
             latents=latents,
             timesteps=timesteps,
@@ -650,6 +674,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             attention_kwargs=attention_kwargs,
             vace_context=vace_context,
             vace_context_scale=vace_context_scale,
+            boundary_timestep=boundary_timestep,
         )
 
         self._current_timestep = None

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -40,7 +40,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.forward_context import build_local_sp_padding_mask, get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
 from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
@@ -799,7 +799,7 @@ class WanTransformer3DModel(nn.Module):
     # The _sp_plan specifies sharding/gathering at module boundaries:
     # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True
     # - timestep_proj_prepare: Split timestep_proj for TI2V models (4D tensor)
-    # - blocks.0: Split hidden_states input at the first transformer block
+    # - _sp_shard_point: Split hidden_states before CacheDiT-wrapped transformer blocks
     # - proj_out: Gather outputs after the final projection layer
     #
     # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism)
@@ -822,10 +822,10 @@ class WanTransformer3DModel(nn.Module):
                 split_dim=1, expected_dims=4, split_output=True, auto_pad=True
             ),  # [B, seq, 6, dim]
         },
-        # Shard hidden_states at first transformer block input
-        # (after patch_embedding + flatten + transpose)
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),  # [B, seq, dim]
+        # Shard hidden_states before entering transformer blocks. This keeps
+        # CacheDiT block wrappers aligned with the local SP shard.
+        "_sp_shard_point": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
         },
         # Shard output scale/shift for TI2V (3D); T2V outputs 2D and skips sharding
         "output_scale_shift_prepare": {
@@ -856,7 +856,6 @@ class WanTransformer3DModel(nn.Module):
         quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
-
         # Store config for compatibility
         self.config = type(
             "Config",
@@ -938,6 +937,7 @@ class WanTransformer3DModel(nn.Module):
 
         # SP helper modules
         self.timestep_proj_prepare = TimestepProjPrepare()
+        self._sp_shard_point = nn.Identity()
         if is_pipeline_last_stage():
             self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
         else:
@@ -981,10 +981,11 @@ class WanTransformer3DModel(nn.Module):
             self._cached_rope_emb = rotary_emb
 
         if is_pipeline_first_stage():
-            # Patch embedding and flatten to sequence
-            # (hidden_states is sharded at blocks.0 input by _sp_plan)
+            # Patch embedding and flatten to sequence. SP sharding happens at
+            # _sp_shard_point so downstream block wrappers see local tensors.
             hidden_states = self.patch_embedding(hidden_states)
             hidden_states = hidden_states.flatten(2).transpose(1, 2)
+            hidden_states = self._sp_shard_point(hidden_states)
         else:
             if intermediate_tensors is None:
                 raise RuntimeError("intermediate_tensors must be provided for non-first PP stages")
@@ -1017,18 +1018,12 @@ class WanTransformer3DModel(nn.Module):
             parallel_config is not None
             and parallel_config.sequence_parallel_size > 1
             and parallel_config.mask_sp_padding
-            and ctx.sp_original_seq_len is not None
-            and ctx.sp_padding_size > 0
         ):
-            batch_size = hidden_states.shape[0]
-            padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
-            hidden_states_mask = torch.ones(
-                batch_size,
-                padded_seq_len,
-                dtype=torch.bool,
-                device=hidden_states.device,
+            hidden_states_mask = build_local_sp_padding_mask(
+                hidden_states.shape[0],
+                hidden_states.shape[1],
+                hidden_states.device,
             )
-            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
             if hidden_states_mask.all():
                 hidden_states_mask = None
         elif (

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -81,9 +81,7 @@ class VaceWanTransformerBlock(WanTransformerBlock):
 class WanVACETransformer3DModel(WanTransformer3DModel):
     """VACE-extended WAN Transformer with conditioning blocks for video editing."""
 
-    # TODO: `vace_blocks` are not layerwise-offloaded yet. The current offloader only
-    # supports a single block group (`blocks`); extend it to support both
-    # `vace_blocks` and `blocks`.
+    _layerwise_offload_blocks_attrs = ["vace_blocks", "blocks"]
 
     # Shard hidden_states before VACE blocks (replaces parent's blocks.0)
     _sp_plan = {

vllm_omni/diffusion/sched/base_scheduler.py

@@ -52,6 +52,7 @@ class _BaseScheduler(SchedulerInterface):
         self._running_sampling_params_key: SamplingParamsKey | None = None
         self._finished_req_ids: set[str] = set()
         self.max_num_running_reqs: int = 1
+        self._prefetch_enabled: bool = False
 
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
@@ -66,6 +67,8 @@ class _BaseScheduler(SchedulerInterface):
             self.max_num_running_reqs = max(1, int(max_num_seqs))
         except (TypeError, ValueError):
             self.max_num_running_reqs = 1
+        omni_kv = getattr(od_config, "omni_kv_config", None) or {}
+        self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False))
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -111,6 +114,22 @@ class _BaseScheduler(SchedulerInterface):
             else:
                 scheduled_cached_request_ids.append(request_id)
 
+        # Expose the next waiting request (serial mode) so the runner can
+        # prefetch its KV during this forward.  Skip a request without
+        # kv_sender_info (would target the wrong sender under multi-replica) or
+        # one already finished/aborted (would consume its sender buffer for
+        # nothing).
+        kv_prefetch_jobs: dict | None = None
+        if self._prefetch_enabled and self._waiting:
+            nxt = self._request_states.get(self._waiting[0])
+            if nxt is not None and not nxt.is_finished():
+                sender_info = getattr(nxt.req, "kv_sender_info", None)
+                if sender_info:
+                    kv_prefetch_jobs = {
+                        "request_id": nxt.request_id,
+                        "kv_sender_info": sender_info,
+                    }
+
         scheduler_output = DiffusionSchedulerOutput(
             step_id=self._step_id,
             scheduled_new_reqs=scheduled_new_reqs,
@@ -118,6 +137,7 @@ class _BaseScheduler(SchedulerInterface):
             finished_req_ids=set(self._finished_req_ids),
             num_running_reqs=len(self._running),
             num_waiting_reqs=len(self._waiting),
+            kv_prefetch_jobs=kv_prefetch_jobs,
         )
 
         # update after schedule

vllm_omni/diffusion/sched/interface.py

@@ -113,6 +113,8 @@ class DiffusionSchedulerOutput:
     finished_req_ids: set[str]
     num_running_reqs: int
     num_waiting_reqs: int
+    # next request to background-prefetch KV
+    kv_prefetch_jobs: dict | None = None
 
     @cached_property
     def scheduled_request_ids(self) -> list[str]:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -80,14 +80,29 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Cache for per-request stepwise state.
         self.state_cache: dict[str, DiffusionRequestState] = {}
 
-        # Initialize KV cache manager for connector management
+        # Initialize KV cache manager for connector management.
         self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
 
+        # Prefetch covers TP / SP / CFG-Parallel / HSDP.  Disabled when a CFG
+        # companion KV collector is set (that KV is not backgrounded).
+        has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None
+
+        self._kv_prefetch_enabled = (
+            bool(self.kv_transfer_manager.config.enable_kv_async_prefetch)
+            and not has_cfg_companion_kv
+            and self.kv_transfer_manager.config.need_recv_cache
+        )
+
+    @property
+    def target_device(self) -> torch.device | None:
+        return getattr(self.pipeline, "device", None)
+
     def _compile_transformer(self, attr_name: str) -> None:
         """Compile a transformer attribute on the pipeline with torch.compile."""
         model = getattr(self.pipeline, attr_name, None)
         if model is None:
             return
+
         try:
             setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True))
             logger.info("Model runner: %s compiled with torch.compile.", attr_name)
@@ -268,7 +283,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0,
         )
 
-    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
 
@@ -292,12 +307,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
-            # The manager handles the check for need_recv_cache internally
-            self.kv_transfer_manager.receive_multi_kv_cache_distributed(
-                req,
-                cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                target_device=getattr(self.pipeline, "device", None),
-            )
+            # Receive AR KV (fetch → distribute → apply inside the entry). prefetch on:
+            # consume prior-forward payload, sync-fallback on miss; else sync receive.
+            kv_recv_t0 = time.perf_counter()
+            if self._kv_prefetch_enabled:
+                self.kv_transfer_manager.consume_and_distribute_kv_cache(
+                    req,
+                    target_device=self.target_device,
+                )
+            else:
+                self.kv_transfer_manager.receive_multi_kv_cache_distributed(
+                    req,
+                    cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
+                    target_device=self.target_device,
+                )
+            kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000
+            logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms)
+
+            # Kick off the next request's prefetch (+ H2D) to overlap this forward.
+            if self._kv_prefetch_enabled and kv_prefetch_jobs is not None:
+                self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device)
 
             if req.sampling_params.generator is None and req.sampling_params.seed is not None:
                 if req.sampling_params.generator_device is not None:
@@ -403,7 +432,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self.kv_transfer_manager.receive_multi_kv_cache_distributed(
                     state_req,
                     cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                    target_device=getattr(self.pipeline, "device", None),
+                    target_device=self.target_device,
                 )
                 self.state_cache[request_id] = new_state
                 resolved.append(new_state)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -369,7 +369,12 @@ class DiffusionWorker:
         else:
             profiler.stop()
 
-    def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
+    def execute_model(
+        self,
+        req: OmniDiffusionRequest,
+        od_config: OmniDiffusionConfig,
+        kv_prefetch_jobs: dict | None = None,
+    ) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
         if self.lora_manager is not None:
@@ -382,7 +387,7 @@ class DiffusionWorker:
         profiler = self._get_profiler()
         ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext()
         with ctx:
-            output = self.model_runner.execute_model(req)
+            output = self.model_runner.execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
         if profiler:
             profiler.step()
         return output
@@ -646,6 +651,10 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
+        if self.model_runner is not None:
+            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+            if mgr is not None:
+                mgr.shutdown_prefetch()
         destroy_distributed_env()
 
 
@@ -1077,18 +1086,24 @@ class WorkerWrapperBase:
         """
         return self.worker.generate(requests)
 
-    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+    def execute_model(
+        self,
+        reqs: list[OmniDiffusionRequest],
+        od_config: OmniDiffusionConfig,
+        kv_prefetch_jobs: dict | None = None,
+    ) -> DiffusionOutput:
         """
         Execute a forward pass.
 
         Args:
             reqs: List of diffusion requests
             od_config: OmniDiffusionConfig configuration
+            kv_prefetch_jobs: Optional next-request KV prefetch descriptor.
 
         Returns:
             DiffusionOutput with generated results
         """
-        return self.worker.execute_model(reqs, od_config)
+        return self.worker.execute_model(reqs, od_config, kv_prefetch_jobs=kv_prefetch_jobs)
 
     def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute one diffusion step."""

vllm_omni/engine/arg_utils.py

@@ -40,7 +40,14 @@ def _register_omni_hf_configs() -> None:
         from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
             IndexTTS2Config,
         )
-        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+            MingDenseConfig,
+            MingMoeConfig,
+        )
+        from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
+            MossTTSLocalConfig,
+            MossTTSRealtimeConfig,
+        )
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
@@ -62,7 +69,10 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("dense", MingDenseConfig),
+        ("bailingmm", MingMoeConfig),
         ("indextts2", IndexTTS2Config),
+        ("moss_tts_local", MossTTSLocalConfig),
+        ("moss_tts_realtime", MossTTSRealtimeConfig),
         ("qwen3_tts", Qwen3TTSConfig),
         ("cosyvoice3", CosyVoice3Config),
         ("glm_tts", GLMTTSConfig),

vllm_omni/entrypoints/cli/serve.py

@@ -595,7 +595,6 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable layerwise (blockwise) offloading on DiT modules.",
         )
-
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/entrypoints/openai/api_server.py

@@ -575,15 +575,14 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             **uvicorn_kwargs,
         )
 
-    # NB: Await server shutdown only after the backend context is exited
-    try:
-        await shutdown_task
-    finally:
-        state = getattr(app, "state", None)
-        serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
-        if serving_speech is not None:
-            serving_speech.shutdown()
-        sock.close()
+        try:
+            await shutdown_task
+        finally:
+            state = getattr(app, "state", None)
+            serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
+            if serving_speech is not None:
+                serving_speech.shutdown()
+            sock.close()
 
 
 @asynccontextmanager

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -476,13 +476,6 @@ class StreamingSpeechSessionConfig(BaseModel):
             "Requires response_format='pcm'. Speed adjustment is not supported when streaming."
         ),
     )
-    split_granularity: Literal["sentence", "clause"] = Field(
-        default="sentence",
-        description=(
-            "Text splitting granularity: 'sentence' splits on .!?。！？, "
-            "'clause' also splits on CJK commas ， and semicolons ；."
-        ),
-    )
     word_timestamps: bool = Field(
         default=False,
         description=(

vllm_omni/entrypoints/openai/serving_speech.py

@@ -81,6 +81,7 @@ _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
 _MING_TTS_MODEL_STAGES = {"ming_tts"}
 _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
 _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
+_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"}
 _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
@@ -99,6 +100,7 @@ _TTS_MODEL_STAGES: set[str] = (
     | _MING_TTS_MODEL_STAGES
     | _MOSS_TTS_MODEL_STAGES
     | _MOSS_TTS_FULL_MODEL_STAGES
+    | _MOSS_TTS_LOCAL_MODEL_STAGES
     | _GLM_TTS_MODEL_STAGES
     | _STEP_AUDIO2_TTS_MODEL_STAGES
     | _INDEXTTS2_TTS_MODEL_STAGES
@@ -182,6 +184,18 @@ def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample:
     return header
 
 
+def _infer_audio_num_channels(audio: np.ndarray) -> int:
+    """Infer channel count before streaming PCM bytes are wrapped as WAV."""
+    if audio.ndim == 3 and audio.shape[0] == 1:
+        audio = audio[0]
+    if audio.ndim == 2:
+        if audio.shape[0] in (1, 2):
+            return int(audio.shape[0])
+        if audio.shape[1] in (1, 2):
+            return int(audio.shape[1])
+    return 1
+
+
 def _sanitize_filename(filename: str) -> str:
     """Sanitize filename to prevent path traversal attacks.
 
@@ -688,6 +702,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "moss_tts_nano"
         if model_stage in _MOSS_TTS_FULL_MODEL_STAGES:
             return "moss_tts"
+        if model_stage in _MOSS_TTS_LOCAL_MODEL_STAGES:
+            return "moss_tts"
         if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
             return "higgs_audio_v2"
         if model_stage in _HIGGS_V3_TTS_MODEL_STAGES:
@@ -1675,6 +1691,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             name = ""
         if "realtime" in name:
             return "realtime"
+        if "local" in name:
+            return "local"
         if "ttsd" in name:
             return "ttsd"
         if "soundeffect" in name:
@@ -1701,9 +1719,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 return "Input text cannot be empty"
 
         v = self._moss_variant
-        if v in (None, "tts", "realtime"):
+        if v in (None, "tts", "realtime", "local"):
             if request.ref_audio is None:
-                label = "MOSS-TTS-Nano" if v is None else ("MOSS-TTS-Realtime" if v == "realtime" else "MOSS-TTS")
+                label = (
+                    "MOSS-TTS-Nano"
+                    if v is None
+                    else (
+                        "MOSS-TTS-Realtime"
+                        if v == "realtime"
+                        else ("MOSS-TTS-Local-Transformer" if v == "local" else "MOSS-TTS")
+                    )
+                )
                 return f"{label} requires 'ref_audio' (reference audio for voice cloning)."
             return self._validate_ref_audio_format(request.ref_audio)
 
@@ -1806,11 +1832,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             params["prompt_audio_array"] = [[wav_list, sr]]
             return params
 
-        # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator):
-        # call the upstream processor server-side to produce unified codes ----
+        # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator)
+        # and MOSS-TTS-Local-Transformer-v1.5: call the upstream processor
+        # server-side to produce unified codes. Local-v1.5 ships its own
+        # AutoProcessor (processor_config.json + processing_moss_tts.py) and
+        # reuses this exact build_user_message/encode_audios_from_wav path in
+        # the offline example (examples/.../moss_tts/end2end.py:
+        # _build_unified_codes) -- it is NOT in the same boat as Realtime
+        # (no processor_config.json there), so it must not fall back to the
+        # prompt_audio_array path above (which the talker's preprocess()
+        # never reads -- info_dict["codes"]["ref"] is the only thing it
+        # consumes, so skipping this path silently drops all voice-clone
+        # conditioning and produces unconditioned/garbage audio online). ----
         proc = self._get_moss_processor()
         n_vq = int(getattr(proc.model_config, "n_vq", 32))
-        sr_target = int(getattr(proc.model_config, "sampling_rate", 24000))
+        # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate
+        # regardless of its 48 kHz stereo *output* codec -- mirrors the
+        # offline example's hardcoded encode_audios_from_wav(sampling_rate=24000)
+        # for this variant; proc.model_config.sampling_rate there is the
+        # output rate (48000), the wrong value to resample the reference into.
+        sr_target = 24000 if v == "local" else int(getattr(proc.model_config, "sampling_rate", 24000))
 
         # Reference-audio encoding + speaker caching lives in the model package
         # (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 /
@@ -1837,7 +1878,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             )
 
         user_kwargs: dict[str, Any] = {"text": request.input or ""}
-        if v in ("tts", "realtime"):
+        if v in ("tts", "local"):
             user_kwargs["reference"] = [await _encode_ref(request.ref_audio)]
         elif v == "ttsd":
             refs = [await _encode_ref(request.ref_audio)]
@@ -2555,7 +2596,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         if self._tts_tokenizer is None:
             model_name = self.engine_client.model_config.model
-            self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)
+            trust_remote_code = bool(getattr(self.engine_client.model_config, "trust_remote_code", False))
+            self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
 
         ref_text = request.ref_text
         prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None
@@ -2654,7 +2696,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         assert sr_raw is not None, (
                             "First audio chunk must include sample rate metadata for WAV streaming"
                         )
-                        wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16)
+                        num_channels = _infer_audio_num_channels(np.asarray(chunk_np))
+                        wav_header = _create_wav_header(
+                            sample_rate=sample_rate_val,
+                            num_channels=num_channels,
+                            bits_per_sample=16,
+                        )
                         yield wav_header
                         first_chunk = False
 
@@ -3539,10 +3586,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         elif self._tts_model_type == "ming_tts" and sampling_params_list:
             import copy
 
-            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import TEXT_EOS_TOKEN_ID
+            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+                MOE_TEXT_EOS_TOKEN_ID,
+                TEXT_EOS_TOKEN_ID,
+            )
+
+            hf_config = self.engine_client.model_config.hf_config
+            is_moe = getattr(hf_config, "model_type", "") == "bailingmm"
+            stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID
 
             sampling_params_list = copy.deepcopy(sampling_params_list)
-            sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)]
+            sampling_params_list[0].stop_token_ids = [int(stop_token_id)]
             if request.max_new_tokens is not None:
                 # Ming emits TEXT_EOS after the latent decode budget is exhausted, so
                 # Stage-0 needs one extra token beyond ming_max_decode_steps.

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -1,7 +1,7 @@
 """WebSocket handler for streaming text input TTS.
 
-Accepts text incrementally via WebSocket, buffers and splits at sentence
-boundaries, and generates audio per sentence using the existing TTS pipeline.
+Accepts text incrementally via WebSocket, buffers it until input.done, and
+generates audio once for the buffered input using the existing TTS pipeline.
 
 Protocol:
     Client -> Server:
@@ -41,11 +41,6 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-from vllm_omni.entrypoints.openai.text_splitter import (
-    SPLIT_CLAUSE,
-    SPLIT_SENTENCE,
-    SentenceSplitter,
-)
 from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError
 from vllm_omni.utils.forced_aligner import align as forced_align
 
@@ -63,8 +58,8 @@ class OmniStreamingSpeechHandler:
     """Handles WebSocket sessions for streaming text-input TTS.
 
     Each WebSocket connection is an independent session. Text arrives
-    incrementally, is split at sentence boundaries, and audio is generated
-    per sentence using the existing OmniOpenAIServingSpeech pipeline.
+    incrementally, is buffered until input.done, and audio is generated once
+    for the buffered input using the existing OmniOpenAIServingSpeech pipeline.
 
     Args:
         speech_service: The existing TTS serving instance (reused for
@@ -102,9 +97,7 @@ class OmniStreamingSpeechHandler:
                     await self._send_error(websocket, str(error))
                     return
 
-            boundary_re = SPLIT_CLAUSE if config.split_granularity == "clause" else SPLIT_SENTENCE
-            splitter = SentenceSplitter(boundary_re=boundary_re)
-            sentence_index = 0
+            text_parts: list[str] = []
 
             # 2. Receive text chunks until input.done
             while True:
@@ -138,23 +131,19 @@ class OmniStreamingSpeechHandler:
                     if not isinstance(text, str):
                         await self._send_error(websocket, "input.text requires a string value")
                         continue
-                    sentences = splitter.add_text(text)
-                    for sentence in sentences:
-                        await self._generate_and_send(websocket, config, sentence, sentence_index)
-                        sentence_index += 1
+                    text_parts.append(text)
 
                 elif msg_type == "input.done":
-                    # Flush remaining buffer
-                    remaining = splitter.flush()
-                    if remaining:
-                        await self._generate_and_send(websocket, config, remaining, sentence_index)
-                        sentence_index += 1
+                    full_text = "".join(text_parts).strip()
+                    total_sentences = 0
+                    if full_text:
+                        await self._generate_and_send(websocket, config, full_text, 0)
+                        total_sentences = 1
 
-                    # Send session.done
                     await websocket.send_json(
                         {
                             "type": "session.done",
-                            "total_sentences": sentence_index,
+                            "total_sentences": total_sentences,
                         }
                     )
                     return

vllm_omni/entrypoints/openai/text_splitter.py

@@ -1,120 +0,0 @@
-"""Multi-language sentence boundary detector for streaming TTS input.
-
-Buffers incoming text and splits at sentence boundaries (English and CJK),
-yielding complete sentences for audio generation.
-"""
-
-import re
-from re import Pattern
-
-# Maximum buffer size (in characters) to prevent unbounded memory growth.
-_MAX_BUFFER_SIZE = 100_000  # ~100 KB of text
-
-# Sentence-level: .!? + CJK sentence-ending 。！？
-# NOTE: English requires trailing whitespace to confirm a boundary —
-# end-of-string is NOT treated as a boundary (that is what flush() is for).
-SPLIT_SENTENCE = re.compile(
-    r"(?<=[.!?])\s+"
-    r"|(?<=[。！？])"
-)
-
-# Clause-level: adds CJK commas ， and semicolons ；
-SPLIT_CLAUSE = re.compile(
-    r"(?<=[.!?])\s+"
-    r"|(?<=[。！？，；])"
-)
-
-# Default alias
-_SENTENCE_BOUNDARY_RE = SPLIT_SENTENCE
-
-
-class SentenceSplitter:
-    """Incremental sentence splitter for streaming text input.
-
-    Buffers text and yields complete sentences when boundaries are detected.
-    Designed for TTS pipelines where text arrives incrementally (e.g., from STT).
-
-    Args:
-        min_sentence_length: Minimum character length for a sentence.
-            Sentences shorter than this are kept in the buffer to avoid
-            splitting on abbreviations like "Dr." or "U.S.".
-        boundary_re: Custom compiled regex for sentence boundaries.
-            Use ``SPLIT_SENTENCE`` (default) for sentence-level splitting,
-            ``SPLIT_CLAUSE`` for finer-grained clause-level splitting,
-            or pass your own ``re.Pattern``.
-    """
-
-    def __init__(
-        self,
-        min_sentence_length: int = 2,
-        boundary_re: Pattern[str] | None = None,
-    ) -> None:
-        self._buffer: str = ""
-        self._min_sentence_length = min_sentence_length
-        self._boundary_re = boundary_re or _SENTENCE_BOUNDARY_RE
-
-    @property
-    def buffer(self) -> str:
-        """Current buffered text."""
-        return self._buffer
-
-    def add_text(self, text: str) -> list[str]:
-        """Add text to the buffer and return any complete sentences.
-
-        Args:
-            text: Incoming text chunk.
-
-        Returns:
-            List of complete sentences extracted from the buffer.
-            May be empty if no sentence boundary was found.
-
-        Raises:
-            ValueError: If the buffer exceeds the maximum size.
-        """
-        if not text:
-            return []
-
-        self._buffer += text
-        if len(self._buffer) > _MAX_BUFFER_SIZE:
-            raise ValueError(
-                f"Text buffer exceeded maximum size ({_MAX_BUFFER_SIZE} chars). "
-                "Consider adding sentence-ending punctuation to your input."
-            )
-        return self._extract_sentences()
-
-    def flush(self) -> str | None:
-        """Flush remaining buffered text as a final sentence.
-
-        Returns:
-            The remaining buffered text (stripped), or None if buffer is empty.
-        """
-        remaining = self._buffer.strip()
-        self._buffer = ""
-        return remaining if remaining else None
-
-    def _extract_sentences(self) -> list[str]:
-        """Split buffer at sentence boundaries, keeping incomplete text buffered."""
-        parts = self._boundary_re.split(self._buffer)
-
-        if len(parts) <= 1:
-            # No boundary found — keep everything in buffer
-            return []
-
-        sentences: list[str] = []
-        carry = ""
-        # All parts except the last are complete sentences
-        for i in range(len(parts) - 1):
-            text = carry + parts[i]
-            carry = ""
-            stripped = text.strip()
-            if len(stripped) >= self._min_sentence_length:
-                sentences.append(stripped)
-            elif stripped:
-                # Too short (e.g. "Dr.") — carry forward to next part
-                carry = text
-            # else: empty, skip
-
-        # Last part stays in buffer (may be incomplete)
-        self._buffer = carry + parts[-1]
-
-        return sentences

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -54,5 +54,5 @@ class MossTTSNanoAdapter(_MossTTSAdapterBase):
 
 @register_tts_adapter
 class MossTTSAdapter(_MossTTSAdapterBase):
-    stage_keys = frozenset({"moss_tts", "moss_tts_codec"})
+    stage_keys = frozenset({"moss_tts", "moss_tts_codec", "moss_tts_local", "moss_tts_local_codec"})
     name = "moss_tts"

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -346,11 +346,8 @@ class AudioVAE(PreTrainedModel):
         h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        mean, logvar = torch.chunk(h, 2, dim=1)
-        logvar = torch.clamp(logvar, -30.0, 20.0)
-        std = torch.exp(0.5 * logvar)
-        latent = mean + std * torch.randn_like(mean)
-        latent = latent.transpose(1, 2)
+        mean = torch.chunk(h, 2, dim=1)[0]
+        latent = mean.transpose(1, 2)
         return latent, frame_num
 
     @torch.inference_mode()

vllm_omni/model_executor/models/ming_tts/__init__.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
 # NOTE: Do not import model classes in this file. Importing any
 # submodule in this package triggers __init__.py execution, and
 # both the model registry and pipeline registry import submodules

vllm_omni/model_executor/models/ming_tts/audio_prep.py

@@ -14,7 +14,6 @@ from .config_ming_tts import (
     PATCH_SIZE,
     SAMPLE_RATE,
     VAE_PATCH_SIZE,
-    VISION_START_TOKEN_ID,
     MingTTSConfig,
 )
 
@@ -232,14 +231,17 @@ def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfi
     return filtered if filtered.numel() > 0 else dummy_pos
 
 
-def _find_speaker_placeholder_positions(input_ids: torch.Tensor, hf_config: Any) -> list[int]:
-    vision_start_token_id = getattr(hf_config, "vision_start_token_id", VISION_START_TOKEN_ID)
-    vision_start_pos = (input_ids == int(vision_start_token_id)).nonzero(as_tuple=True)[0]
-    if vision_start_pos.numel() == 0:
+def _find_speaker_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> list[int]:
+    # The speaker embedding is written at the slot AFTER the placeholder marker:
+    # dense ``<|vision_start|>``+1 (the ``<|vision_pad|>`` slot), moe ``<spk>``+1
+    # (the ``<audioPatch>`` slot). See upstream prepare_input_embed.
+    placeholder_token_id = int(cfg.speaker_placeholder_token_id)
+    placeholder_pos = (input_ids == placeholder_token_id).nonzero(as_tuple=True)[0]
+    if placeholder_pos.numel() == 0:
         return []
 
     slots = []
-    for pos in vision_start_pos:
+    for pos in placeholder_pos:
         slot = int(pos.item()) + 1
         if slot < int(input_ids.shape[0]):
             slots.append(slot)

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -42,6 +42,12 @@ from .constants import (
     LLM_HIDDEN_SIZE,
     LLM_VOCAB_SIZE,
     MAX_DECODE_STEPS,
+    MOE_AUDIO_DUMMY_TOKEN_ID,
+    MOE_AUDIO_END_TOKEN_ID,
+    MOE_AUDIO_EOS_TOKEN_ID,
+    MOE_AUDIO_START_TOKEN_ID,
+    MOE_SPK_TOKEN_ID,
+    MOE_TEXT_EOS_TOKEN_ID,
     PATCH_SIZE,
     SAMPLE_RATE,
     STOP_HEAD_MIN_STEPS,
@@ -89,10 +95,126 @@ class MingDenseConfig(PretrainedConfig):
         return self.llm_config
 
 
+class BailingMoeConfig(PretrainedConfig):
+    """Stage-0 LLM config for the Ming MoE variant.
+
+    Ported from upstream inclusionAI ``configuration_bailing_moe.py``; the
+    community vLLM ``BailingMoeModel`` is duck-typed over these attributes.
+    Only the fields the backbone reads are kept; unknown keys flow through
+    ``**kwargs`` to ``PretrainedConfig``.
+    """
+
+    model_type = "bailing_moe"
+
+    def __init__(
+        self,
+        vocab_size: int = 126464,
+        hidden_size: int = 2048,
+        intermediate_size: int | None = None,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 0,
+        hidden_act: str = "silu",
+        use_qkv_bias: bool = False,
+        use_bias: bool = True,
+        rms_norm_eps: float = 1e-05,
+        norm_head: bool = False,
+        tie_word_embeddings: bool = False,
+        max_position_embeddings: int = 32768,
+        rope_theta: float = 10000.0,
+        use_cache: bool = True,
+        rope_scaling: Any = None,
+        pad_token_id: int = 126081,
+        num_experts: int = 64,
+        num_shared_experts: int = 2,
+        num_experts_per_tok: int = 6,
+        norm_topk_prob: bool = True,
+        moe_intermediate_size: int | None = None,
+        first_k_dense_replace: int = 0,
+        head_dim: int | None = None,
+        output_router_logits: bool = False,
+        multi_gate: bool = False,
+        image_patch_token: int = 126346,
+        use_grouped_gemm: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_bias = use_bias
+        self.norm_head = norm_head
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.rope_scaling = rope_scaling
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.norm_topk_prob = norm_topk_prob
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.multi_gate = multi_gate
+        self.image_patch_token = image_patch_token
+        self.use_grouped_gemm = use_grouped_gemm
+        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+def _coerce_bailing_moe_config(value: Any) -> BailingMoeConfig:
+    if isinstance(value, BailingMoeConfig):
+        cfg = value
+    elif isinstance(value, PretrainedConfig):
+        cfg = BailingMoeConfig.from_dict(value.to_dict())
+    elif isinstance(value, dict):
+        cfg = BailingMoeConfig.from_dict(dict(value))
+    else:
+        raise TypeError(f"Unsupported llm_config type for Ming MoE config: {type(value)!r}")
+    # The multimodal bailing_moe declares 3D mrope (rope_scaling type "3D").
+    # Ming-TTS Stage-0 carries no vision/video tokens, so mrope degenerates to
+    # standard 1D RoPE (all three position dims are equal). Strip rope_scaling
+    # so vLLM's get_rope uses plain RoPE — it has no "3D" scaling type.
+    cfg.rope_scaling = None
+    cfg.rope_parameters = {"rope_type": "default", "rope_theta": cfg.rope_theta}
+    return cfg
+
+
+class MingMoeConfig(PretrainedConfig):
+    # Upstream inclusionAI Ming-omni-tts-16.8B-A3B HF config reports model_type="bailingmm".
+    model_type = "bailingmm"
+
+    def __init__(
+        self,
+        llm_config: BailingMoeConfig | dict[str, Any] | None = None,
+        ditar_config: dict[str, Any] | None = None,
+        aggregator_config: dict[str, Any] | None = None,
+        audio_tokenizer_config: AudioVAEConfig | dict[str, Any] | None = None,
+        architectures: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.llm_config = _coerce_bailing_moe_config(llm_config or {})
+        self.ditar_config = dict(ditar_config or {})
+        self.aggregator_config = dict(aggregator_config or {})
+        self.audio_tokenizer_config = _coerce_audio_vae_config(audio_tokenizer_config)
+        super().__init__(architectures=architectures, **kwargs)
+
+    def get_text_config(self, decoder: bool = False, **kwargs: Any) -> BailingMoeConfig:
+        del decoder, kwargs
+        return self.llm_config
+
+
 @dataclass
 class MingTTSConfig:
     """Flat config object shared by Stage-1 and Stage-2. Build via from_hf_config()."""
 
+    model_variant: str = "dense"  # "dense" (qwen2 backbone) | "moe" (bailing_moe backbone)
+
     llm_hidden_size: int = LLM_HIDDEN_SIZE
     llm_vocab_size: int = LLM_VOCAB_SIZE
     llm_config: dict[str, Any] = field(default_factory=dict)
@@ -124,6 +246,7 @@ class MingTTSConfig:
     audio_start_token_id: int = AUDIO_START_TOKEN_ID
     audio_end_token_id: int = AUDIO_END_TOKEN_ID
     audio_eos_token_id: int = AUDIO_EOS_TOKEN_ID
+    speaker_placeholder_token_id: int = VISION_START_TOKEN_ID  # dense <|vision_start|>; moe <spk>
 
     @classmethod
     def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig:
@@ -144,7 +267,12 @@ class MingTTSConfig:
 
         enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP)
 
+        # Variant is decided by the Stage-0 LLM backbone family, not the wrapper
+        # model_type: dense -> qwen2, MoE -> bailing_moe.
+        model_variant = "moe" if str(llm_dict.get("model_type", "")) == "bailing_moe" else "dense"
+
         cfg = cls(
+            model_variant=model_variant,
             llm_hidden_size=llm_dict.get("hidden_size", LLM_HIDDEN_SIZE),
             llm_vocab_size=llm_dict.get("vocab_size", LLM_VOCAB_SIZE),
             llm_config=llm_dict,
@@ -158,6 +286,15 @@ class MingTTSConfig:
             sample_rate=atc_sample_rate,
             audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP,
         )
+        if model_variant == "moe":
+            # The bailing tokenizer uses a different vocab; override the dense
+            # (Qwen2) special-token defaults with the bailing token IDs.
+            cfg.text_eos_token_id = MOE_TEXT_EOS_TOKEN_ID
+            cfg.audio_dummy_token_id = MOE_AUDIO_DUMMY_TOKEN_ID
+            cfg.audio_start_token_id = MOE_AUDIO_START_TOKEN_ID
+            cfg.audio_end_token_id = MOE_AUDIO_END_TOKEN_ID
+            cfg.audio_eos_token_id = MOE_AUDIO_EOS_TOKEN_ID
+            cfg.speaker_placeholder_token_id = MOE_SPK_TOKEN_ID
         return cfg
 
     def validate(self) -> None:
@@ -195,7 +332,9 @@ __all__ = [
     "LLM_HIDDEN_SIZE",
     "LLM_VOCAB_SIZE",
     "MAX_DECODE_STEPS",
+    "BailingMoeConfig",
     "MingDenseConfig",
+    "MingMoeConfig",
     "MingTTSConfig",
     "PATCH_SIZE",
     "SAMPLE_RATE",

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -15,6 +15,22 @@ VISION_START_TOKEN_ID = 151652  # <|vision_start|>
 TEXT_EOS_TOKEN_ID = 151669  # <text_eos>
 
 
+# ---------------------------------------------------------------------------
+# MoE (bailing_moe / Ming-omni-tts-16.8B-A3B) token IDs
+# (confirmed from the bailing tokenizer_config.json — different vocab than the
+# dense Qwen2 tokenizer above). The bailing tokenizer has NO <text_eos>; the
+# audio AR loop terminates purely on the stop-head (see upstream
+# modeling_bailingmm.sample), so <end_of_audio> doubles as the AR stop token.
+# ---------------------------------------------------------------------------
+
+MOE_AUDIO_DUMMY_TOKEN_ID = 126357  # <audioPatch>
+MOE_AUDIO_START_TOKEN_ID = 126358  # <audio>
+MOE_AUDIO_END_TOKEN_ID = 126359  # </audio>
+MOE_AUDIO_EOS_TOKEN_ID = 126356  # <end_of_audio>
+MOE_TEXT_EOS_TOKEN_ID = 126356  # no <text_eos> in bailing; reuse <end_of_audio> as AR stop
+MOE_SPK_TOKEN_ID = 126368  # <spk> — speaker-embedding placeholder (dense uses <|vision_start|>)
+
+
 # ---------------------------------------------------------------------------
 # Architectural constants (confirmed from original config.json)
 # ---------------------------------------------------------------------------

vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUDAGraph-accelerated CFM (flow-matching) diffusion head for Ming-TTS.
+
+The per-AR-step flow head runs ``steps`` ODE iterations, each a tiny DiT forward
+(history+patch tokens, CFG cond+null batched). In eager mode these are dozens of
+micro kernels per step and the wall time is dominated by kernel-launch / Python
+overhead, not GPU math. Capturing the whole sampling loop + Aggregator + stop
+head into one CUDA graph removes that overhead.
+
+Adapted from adaspeech/ming_flash ``CFMGraphExecutor``. Captured for batch=1
+(Ming Stage-0 runs ``max_num_seqs: 1``); other batch sizes fall back to eager.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.common.ming.fm import build_timesteps
+
+
+class CFMSampler(nn.Module):
+    """Unrolled CFM ODE solver, graph-capturable (no ``torch.randn``/``.item()``).
+
+    All randomness (``y0``, ``sde_rnd``) and the timestep schedule ``t`` are
+    passed in as float32 tensors so the loop is deterministic during replay and
+    the integration matches the eager float32 path. Equivalent to eager
+    ``Solver.integrate`` for ``temperature==0`` (the TTS default).
+    """
+
+    def __init__(self, dit_model: nn.Module, steps: int = 10):
+        super().__init__()
+        self.model = dit_model
+        self.steps = steps
+
+    @torch.no_grad()
+    def forward(self, llm_cond, lat_cond, y0, t, sde_args, sde_rnd):
+        # y0 / t / sde_rnd / sde_args are float32; integrate in float32 and cast
+        # only the DiT inputs to the model dtype, mirroring eager FlowLoss.sample.
+        model_dtype = llm_cond.dtype
+        for step in range(self.steps):
+            dt = t[step + 1] - t[step]
+            pred_cfg = self.model.forward_with_cfg(
+                x=y0.to(model_dtype),
+                t=t[step].to(model_dtype),
+                c=llm_cond,
+                latent_history=lat_cond,
+                cfg_scale=sde_args[0],
+                patch_size=y0.shape[1],
+            )
+            pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+            flow = (pred + (pred - null_pred) * sde_args[0]).float()
+            y0 = y0 + flow * dt
+            if step < self.steps - 1:
+                y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
+        return y0
+
+
+class CFMGraphExecutor:
+    """Capture CFM sampling + Aggregator + stop head into one CUDA graph.
+
+    First ``execute`` captures the graph; later calls copy inputs into the
+    static placeholders and ``replay()``. Runtime controls (cfg/sigma/temp) are
+    fed via the ``sde_args`` placeholder, so they are tunable without recapture.
+    """
+
+    def __init__(
+        self,
+        cfm_sampler: CFMSampler,
+        aggregator: nn.Module,
+        stop_head: nn.Module,
+        *,
+        patch_size: int,
+        latent_dim: int,
+        steps: int = 10,
+    ):
+        self.cfm_sampler = cfm_sampler
+        self.aggregator = aggregator
+        self.stop_head = stop_head
+        self.patch_size = int(patch_size)
+        self.latent_dim = int(latent_dim)
+        self.steps = int(steps)
+        self.initialized = False
+        self.graph: torch.cuda.CUDAGraph | None = None
+
+    def _timesteps(self, device):
+        # Raw EPSS schedule + sway in float32, identical to eager build_timesteps.
+        return build_timesteps(self.steps, device=device, dtype=torch.float32, use_epss=True, sway_sampling_coef=-1.0)
+
+    @torch.no_grad()
+    def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature):
+        bsz = llm_cond.shape[0]
+        device = llm_cond.device
+        # ODE state / noise / timesteps are float32 to match the eager path.
+        y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+        t = self._timesteps(device)
+        sde_rnd = torch.randn((self.steps, bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+
+        if not self.initialized:
+            self._initialize_graph(llm_cond, his_lat, y0, t, sde_rnd)
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0] = float(cfg_strength)
+        self._sde_args_ph[1] = float(sigma)
+        self._sde_args_ph[2] = float(temperature)
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        self.graph.replay()
+
+        return self._gen_lat_ph.clone(), self._embeds_ph.clone(), self._stop_ph.clone()
+
+    def _initialize_graph(self, llm_cond, his_lat, y0, t, sde_rnd):
+        model_dtype = llm_cond.dtype
+        self._llm_cond_ph = torch.empty_like(llm_cond)
+        self._his_lat_ph = torch.empty_like(his_lat)
+        self._y0_ph = torch.empty_like(y0)  # float32
+        self._t_ph = torch.empty_like(t)  # float32
+        self._sde_args_ph = torch.empty(3, device=llm_cond.device, dtype=torch.float32)
+        self._sde_rnd_ph = torch.empty_like(sde_rnd)  # float32
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0], self._sde_args_ph[1], self._sde_args_ph[2] = 2.0, 0.25, 0.0
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        def _run():
+            # Solver runs in float32; cast the final latent to the model dtype
+            # for the Aggregator / stop head (matches eager ``.to(dtype)``).
+            gen_lat = self.cfm_sampler(
+                self._llm_cond_ph, self._his_lat_ph, self._y0_ph, self._t_ph, self._sde_args_ph, self._sde_rnd_ph
+            ).to(model_dtype)
+            embeds = self.aggregator(gen_lat)
+            stop = self.stop_head(self._llm_cond_ph[:, -1, :]).softmax(dim=-1)
+            return gen_lat, embeds, stop
+
+        # Warm up on a side stream to stabilize allocations before capture.
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            for _ in range(3):
+                _run()
+        torch.cuda.current_stream().wait_stream(warmup)
+
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self._gen_lat_ph, self._embeds_ph, self._stop_ph = _run()
+
+        self.initialized = True

vllm_omni/model_executor/models/ming_tts/fm/dit.py

@@ -87,12 +87,17 @@ class DiT(nn.Module):
         return x
 
     def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
-        if not cfg_scale == 1:
-            x = torch.cat([x, x], dim=0)
-            latent_history = torch.cat([latent_history, latent_history], dim=0)
-            fake_latent = torch.zeros_like(c)
-            c = torch.cat([c, fake_latent], dim=0)
+        # Always run the classifier-free-guidance path (cond + uncond). TTS
+        # inference always uses cfg>1, and branching on ``cfg_scale`` is avoided
+        # because under CUDA-graph capture it is a CUDA tensor — ``bool(cfg_scale)``
+        # would force a device sync and break capture. ``expand`` (not ``repeat``)
+        # avoids a fresh allocation each step.
+        x = torch.cat([x, x], dim=0)
+        latent_history = torch.cat([latent_history, latent_history], dim=0)
+        fake_latent = torch.zeros_like(c)
+        c = torch.cat([c, fake_latent], dim=0)
         if t.ndim == 0:
-            t = t.repeat(x.shape[0])
+            t = t.reshape(1)
+        t = t.expand(x.shape[0])
         model_out = self.forward(x, t, c, latent_history)
         return model_out[:, -patch_size:, :]

vllm_omni/model_executor/models/ming_tts/ming_tts.py

@@ -182,7 +182,7 @@ class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin)
                 use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)),
             )
         if speaker_embeddings is not None and len(speaker_embeddings) > 0:
-            speaker_slots = _find_speaker_placeholder_positions(input_ids, self.vllm_config.model_config.hf_config)
+            speaker_slots = _find_speaker_placeholder_positions(input_ids, self.ming_config)
             if len(speaker_slots) < len(speaker_embeddings):
                 raise RuntimeError(
                     "Could not locate enough speaker placeholder slots: "

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -3,6 +3,7 @@
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
 from __future__ import annotations
 
+import os
 from collections.abc import Iterable
 from typing import Any
 
@@ -51,6 +52,9 @@ from .patch_emission import (
 
 logger = init_logger(__name__)
 
+# CFM ODE integration steps; must match CFM.sample(steps=...) default (10).
+_CFM_STEPS = 10
+
 
 class MingLLMModel(nn.Module):
     hf_to_vllm_mapper = WeightsMapper(
@@ -66,7 +70,16 @@ class MingLLMModel(nn.Module):
         self.vllm_config = vllm_config
         self.prefix = prefix
         self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config)
-        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        if self.ming_config.model_variant == "moe":
+            from vllm.model_executor.models.bailing_moe import BailingMoeModel
+
+            # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly
+            # (no get_text_config()), so re-wrap with the nested bailing_moe config.
+            llm_config = vllm_config.model_config.hf_config.get_text_config()
+            llm_vllm_config = vllm_config.with_hf_config(llm_config)
+            self.model = BailingMoeModel(vllm_config=llm_vllm_config, prefix=maybe_prefix(prefix, "model"))
+        else:
+            self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
         self.linear_proj_audio = Aggregator(
             in_channels=self.ming_config.latent_dim,
             llm_input_dim=self.ming_config.llm_hidden_size,
@@ -86,6 +99,11 @@ class MingLLMModel(nn.Module):
         self._pending_postprocess_updates: dict[str, dict[str, Any]] = {}
         self._last_ming_next_token_ids = None
         self._last_text_mode = False
+        # CUDAGraph for the flow-matching (CFM) diffusion head — the TTFP/latency
+        # hot spot. Built lazily on the first batch=1 CUDA decode; set
+        # MING_CFM_CUDAGRAPH=0 to force the eager flow head.
+        self._cfm_graph = None
+        self._cfm_graph_enabled = os.environ.get("MING_CFM_CUDAGRAPH", "1") == "1"
 
     def embed_input_ids(
         self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any
@@ -118,11 +136,14 @@ class MingLLMModel(nn.Module):
             min_stop_step=int(self.ming_config.stop_head_min_steps),
             default_max_decode_steps=self.ming_config.max_decode_steps,
         )
+        # Positional call: Qwen2Model names the 2nd arg ``positions`` while
+        # BailingMoeModel names it ``position_ids`` — both share the signature
+        # (input_ids, positions, intermediate_tensors, inputs_embeds).
         backbone_out = self.model(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
+            input_ids,
+            positions,
+            intermediate_tensors,
+            inputs_embeds,
         )
         if isinstance(backbone_out, IntermediateTensors):
             self._last_ming_next_token_ids = None
@@ -294,6 +315,33 @@ class MingLLMModel(nn.Module):
             logprobs_tensors=None,
         )
 
+    def _maybe_build_cfm_graph(self, z_diff_cond: torch.Tensor):
+        """Lazily build the CFM CUDAGraph executor (batch=1, CUDA only)."""
+        if not self._cfm_graph_enabled:
+            return None
+        if z_diff_cond.shape[0] != 1 or not z_diff_cond.is_cuda:
+            return None
+        if self._cfm_graph is not None:
+            return self._cfm_graph
+        try:
+            from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
+
+            sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS)
+            self._cfm_graph = CFMGraphExecutor(
+                sampler,
+                self.linear_proj_audio,
+                self.stop_head,
+                patch_size=self.ming_config.patch_size,
+                latent_dim=self.ming_config.latent_dim,
+                steps=_CFM_STEPS,
+            )
+            logger.info("Ming CFM CUDAGraph enabled (steps=%d, patch=%d).", _CFM_STEPS, self.ming_config.patch_size)
+        except Exception as exc:
+            logger.warning("Ming CFM CUDAGraph unavailable (%s); using eager flow head.", exc)
+            self._cfm_graph_enabled = False
+            return None
+        return self._cfm_graph
+
     def _decode_one_step(
         self,
         *,
@@ -317,23 +365,45 @@ class MingLLMModel(nn.Module):
         z_diff_cond = hidden_states.to(dtype=self.fm_dtype).unsqueeze(1)
         if not torch.isfinite(z_diff_cond).all():
             raise RuntimeError("Non-finite z_diff_cond before FlowLoss.sample().")
-        sampled_token_latent = self.flowloss.sample(
-            z=z_diff_cond,
-            latent_history=latent_history,
-            cfg=cfg_scale,
-            patch_size=self.ming_config.patch_size,
-            sigma=sigma,
-            temperature=temperature,
-        )
+
+        # Fast path: CUDAGraph-captured flow head (CFM sampling + Aggregator +
+        # stop head). Falls back permanently to the eager path on any failure.
+        sampled_token_latent = next_embeds = stop_probs = None
+        # cfg < 1e-5 disables CFG; the eager flow head handles this with a
+        # dedicated unconditional branch (c=zeros) that the captured graph does
+        # not replicate, so fall back to eager there to match FlowLoss.sample.
+        graph_exec = self._maybe_build_cfm_graph(z_diff_cond) if cfg_scale >= 1e-5 else None
+        if graph_exec is not None:
+            try:
+                sampled_token_latent, next_embeds, stop_full = graph_exec.execute(
+                    z_diff_cond, latent_history, cfg_scale, sigma, temperature
+                )
+                stop_probs = stop_full[:, 1]
+            except Exception as exc:
+                logger.warning("Ming CFM CUDAGraph failed (%s); falling back to eager flow head.", exc)
+                self._cfm_graph = None
+                self._cfm_graph_enabled = False
+                sampled_token_latent = next_embeds = stop_probs = None
+
+        if sampled_token_latent is None:
+            sampled_token_latent = self.flowloss.sample(
+                z=z_diff_cond,
+                latent_history=latent_history,
+                cfg=cfg_scale,
+                patch_size=self.ming_config.patch_size,
+                sigma=sigma,
+                temperature=temperature,
+            )
+            # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H].
+            next_embeds = self.linear_proj_audio(sampled_token_latent)
+            stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1]
+
         expected_shape = (hidden_states.shape[0], self.ming_config.patch_size, self.ming_config.latent_dim)
         if tuple(sampled_token_latent.shape) != expected_shape:
             raise RuntimeError(
                 f"FlowLoss output shape mismatch: got {tuple(sampled_token_latent.shape)}, expected {expected_shape}"
             )
         new_history = torch.cat([latent_history[:, self.ming_config.patch_size :, :], sampled_token_latent], dim=1)
-        # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H].
-        next_embeds = self.linear_proj_audio(sampled_token_latent)
-        stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1]
         if not torch.isfinite(sampled_token_latent).all():
             raise RuntimeError("Non-finite sampled_token_latent in Ming decode step.")
         if not torch.isfinite(next_embeds).all():

vllm_omni/model_executor/models/ming_tts/pipeline.py

@@ -13,7 +13,10 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.ming_tts"
 MING_TTS_PIPELINE = PipelineConfig(
     model_type="ming_tts",
     model_arch="MingTTSForConditionalGeneration",
-    hf_architectures=("MingTTSForConditionalGeneration",),
+    hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"),
+    # The dense (0.5B) and MoE (16.8B) share architectures=["BailingMMNativeForConditionalGeneration"]
+    # Here we disambiguate by the model_type the upstream HF config reports
+    hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "dense",
     stages=(
         StagePipelineConfig(
             stage_id=0,
@@ -52,3 +55,17 @@ MING_TTS_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# MoE variant (inclusionAI/Ming-omni-tts-16.8B-A3B)
+# Same two-stage topology and model class as the above,
+# but with a distinct model_type reported by the upstream HF config.
+# Keep this for auto-detection of deploy config yaml.
+MING_TTS_MOE_PIPELINE = PipelineConfig(
+    model_type="ming_tts_moe",
+    model_arch="MingTTSForConditionalGeneration",
+    hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"),
+    # Disambiguate from the dense-0.5B ming_tts pipeline
+    hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "bailingmm",
+    stages=MING_TTS_PIPELINE.stages,
+)

vllm_omni/model_executor/models/ming_tts/prompt_assembly.py

@@ -176,6 +176,84 @@ def build_dense_prompt_token_ids(
     )
 
 
+def build_moe_prompt_token_ids(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    text: str,
+    instruction: str | None = None,
+    prompt_text: str | None = None,
+    speaker_count: int = 0,
+    prompt_patch_count: int = 0,
+) -> list[int]:
+    """Stage-0 prompt for the Ming MoE (bailing_moe) variant.
+
+    Mirrors the ``model_type != 'dense'`` branch of upstream
+    ``BailingMMNativeForConditionalGeneration.prepare_input_embed``: no system
+    message, ``<role>HUMAN/ASSISTANT</role>`` turns, and a ``<spk><audioPatch></spk>``
+    speaker placeholder (vs. the dense ``<|vision_*|>`` placeholder).
+    """
+    speaker_prompt = []
+    for idx in range(int(speaker_count)):
+        speaker_prompt.extend(
+            tokenizer.encode(f"  speaker_{idx + 1}:")
+            + tokenizer.encode("<spk>")
+            + tokenizer.encode("<audioPatch>")
+            + tokenizer.encode("</spk>\n")
+        )
+    instruction_prompt = (
+        tokenizer.encode(instruction) + tokenizer.encode("<|endoftext|>") if instruction is not None else []
+    )
+    prompt_text_tokens = (
+        tokenizer.encode(prompt_text) if int(prompt_patch_count) > 0 and prompt_text is not None else []
+    )
+    audio_patch_token_id = tokenizer.convert_tokens_to_ids("<audioPatch>")
+    if audio_patch_token_id == tokenizer.unk_token_id:
+        raise ValueError("Ming tokenizer is missing required <audioPatch> token.")
+    prompt_latent_tokens = [audio_patch_token_id] * int(prompt_patch_count)
+    text_input_prefix = (
+        []
+        if all(token in text for token in ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: "))
+        else tokenizer.encode(" Text input:\n")
+    )
+    return (
+        tokenizer.encode("<role>HUMAN</role>")
+        + tokenizer.encode(prompt)
+        + speaker_prompt
+        + text_input_prefix
+        + prompt_text_tokens
+        + tokenizer.encode(text)
+        + tokenizer.encode("<role>ASSISTANT</role>")
+        + instruction_prompt
+        + tokenizer.encode("<audio>")
+        + prompt_latent_tokens
+    )
+
+
+def build_prompt_token_ids(model_variant: str, tokenizer: Any, **kwargs: Any) -> list[int]:
+    """Dispatch to the dense or MoE Stage-0 prompt builder by model variant."""
+    if model_variant == "moe":
+        return build_moe_prompt_token_ids(tokenizer, **kwargs)
+    return build_dense_prompt_token_ids(tokenizer, **kwargs)
+
+
+def detect_model_variant(tokenizer: Any) -> str:
+    """Infer the Ming variant from the tokenizer's special tokens.
+
+    The bailing (MoE) tokenizer carries ``<role>`` chat markers; the dense
+    Qwen2 tokenizer does not. Used to pick the prompt format when the caller
+    does not pass an explicit variant.
+    """
+    try:
+        role_id = tokenizer.convert_tokens_to_ids("<role>")
+    except Exception:
+        return "dense"
+    unk = getattr(tokenizer, "unk_token_id", None)
+    if role_id is None or (unk is not None and role_id == unk):
+        return "dense"
+    return "moe"
+
+
 def build_ming_dense_prompt(
     tokenizer: Any,
     *,
@@ -189,7 +267,10 @@ def build_ming_dense_prompt(
     speaker_embedding: Any = None,
     use_zero_spk_emb: bool = False,
     request_id: str | None = None,
+    model_variant: str | None = None,
 ) -> dict[str, Any]:
+    if model_variant is None:
+        model_variant = detect_model_variant(tokenizer)
     instruction_text = create_instruction(instruction)
     speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb)
     effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls)
@@ -217,7 +298,8 @@ def build_ming_dense_prompt(
             prompt_latent_value, patch_size=PATCH_SIZE, latent_dim=LATENT_DIM
         )
 
-    prompt_token_ids = build_dense_prompt_token_ids(
+    prompt_token_ids = build_prompt_token_ids(
+        model_variant,
         tokenizer,
         prompt=prompt,
         text=text,
@@ -266,7 +348,10 @@ __all__ = [
     "DEFAULT_PROMPT",
     "build_dense_prompt_token_ids",
     "build_ming_dense_prompt",
+    "build_moe_prompt_token_ids",
+    "build_prompt_token_ids",
     "create_instruction",
+    "detect_model_variant",
     "estimate_decode_step_window_for_duration",
     "estimate_decode_steps_for_duration",
     "parse_duration_seconds",

vllm_omni/model_executor/models/ming_tts/validation.py

@@ -14,6 +14,9 @@ from .constants import (
     LATENT_DIM,
     LLM_HIDDEN_SIZE,
     LLM_VOCAB_SIZE,
+    MOE_AUDIO_DUMMY_TOKEN_ID,
+    MOE_AUDIO_EOS_TOKEN_ID,
+    MOE_TEXT_EOS_TOKEN_ID,
     PATCH_SIZE,
     SAMPLE_RATE,
 )
@@ -70,19 +73,23 @@ def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any:
 
 def validate_ming_tts_config(cfg: Any) -> None:
     """Run before GPU allocation/weight loading. Raises ValueError on mismatches."""
-    if cfg.audio_dummy_token_id != 151705:
+    is_moe = getattr(cfg, "model_variant", "dense") == "moe"
+    exp_audio_dummy = MOE_AUDIO_DUMMY_TOKEN_ID if is_moe else 151705
+    exp_audio_eos = MOE_AUDIO_EOS_TOKEN_ID if is_moe else 151704
+    exp_text_eos = MOE_TEXT_EOS_TOKEN_ID if is_moe else 151669
+    if cfg.audio_dummy_token_id != exp_audio_dummy:
         raise ValueError(
-            f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected 151705 (<audioPatch>). "
+            f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected {exp_audio_dummy} (<audioPatch>). "
             "Wrong tokenizer/checkpoint?"
         )
-    if cfg.audio_eos_token_id != 151704:
+    if cfg.audio_eos_token_id != exp_audio_eos:
         raise ValueError(
-            f"audio_eos_token_id={cfg.audio_eos_token_id}, expected 151704 (<end_of_audio>). "
+            f"audio_eos_token_id={cfg.audio_eos_token_id}, expected {exp_audio_eos} (<end_of_audio>). "
             "Wrong tokenizer/checkpoint?"
         )
-    if cfg.text_eos_token_id != 151669:
+    if cfg.text_eos_token_id != exp_text_eos:
         raise ValueError(
-            f"text_eos_token_id={cfg.text_eos_token_id}, expected 151669 (<text_eos>). Wrong tokenizer/checkpoint?"
+            f"text_eos_token_id={cfg.text_eos_token_id}, expected {exp_text_eos}. Wrong tokenizer/checkpoint?"
         )
 
     if cfg.audio_tokenizer_config is None:
@@ -102,13 +109,16 @@ def validate_ming_tts_config(cfg: Any) -> None:
             f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. "
             "Check ditar_config.history_patch_size."
         )
-    if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
-        raise ValueError(
-            f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
-            "Check llm_config.hidden_size."
-        )
-    if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
-        raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
+    # Absolute hidden/vocab pins are dense-only; the MoE backbone carries its
+    # own (larger) dims, cross-checked against llm_config.hidden_size below.
+    if not is_moe:
+        if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
+            raise ValueError(
+                f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
+                "Check llm_config.hidden_size."
+            )
+        if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
+            raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
     if cfg.sample_rate != SAMPLE_RATE:
         raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.")
 

vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py

@@ -45,6 +45,8 @@ class MossAudioTokenizerConfig(PretrainedConfig):
         causal_transformer_context_duration: float = 10.0,
         encoder_kwargs: list[dict[str, Any]] | None = None,
         decoder_kwargs: list[dict[str, Any]] | None = None,
+        number_channels: int = 1,
+        enable_channel_interleave: bool = True,
         quantizer_type: str = "rlfq",
         quantizer_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
@@ -54,6 +56,8 @@ class MossAudioTokenizerConfig(PretrainedConfig):
         self.sampling_rate = sampling_rate
         self.downsample_rate = downsample_rate
         self.causal_transformer_context_duration = causal_transformer_context_duration
+        self.number_channels = number_channels
+        self.enable_channel_interleave = enable_channel_interleave
         self.encoder_kwargs = encoder_kwargs or _default_encoder_kwargs()
         self.decoder_kwargs = decoder_kwargs or _default_decoder_kwargs()
         if quantizer_kwargs is None:
@@ -201,13 +205,21 @@ class _Attention(nn.Module):
     """Causal multi-head self-attention with RoPE, no streaming KV cache."""
 
     def __init__(
-        self, embed_dim: int, num_heads: int, causal: bool, max_period: float, device=None, dtype=None
+        self,
+        embed_dim: int,
+        num_heads: int,
+        causal: bool,
+        max_period: float,
+        context: int | None = None,
+        device=None,
+        dtype=None,
     ) -> None:
         super().__init__()
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         self.causal = causal
         self.max_period = max_period
+        self.context = context
         kw = {"device": device, "dtype": dtype}
         self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw)
         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw)
@@ -234,7 +246,38 @@ class _Attention(nn.Module):
         qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]  # each (B, H, T, D)
         q, k = _apply_rope(q, k, self.max_period)
-        out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
+        if self.context is not None and self.context < T:
+            # Local-windowed causal attention: query i may only see keys in
+            # [i - context + 1, i]. Matches upstream's per-stage receptive
+            # field — wider context windows are not numerically equivalent.
+            #
+            # A dense (T, T) mask is O(T^2) memory (e.g. 4 GiB at T=65536) —
+            # for long utterances decoded in a single non-streaming call this
+            # OOMs. Process queries in blocks against only their reachable
+            # key range so peak mask memory is O(block * (block + context))
+            # regardless of T.
+            block = 4096
+            if T <= block:
+                positions = torch.arange(T, device=x.device)
+                delta = positions.view(-1, 1) - positions.view(1, -1)
+                mask = (delta >= 0) & (delta < self.context)
+                out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+            else:
+                outs = []
+                for start in range(0, T, block):
+                    end = min(start + block, T)
+                    k_start = max(0, start - self.context + 1)
+                    q_pos = torch.arange(start, end, device=x.device)
+                    k_pos = torch.arange(k_start, end, device=x.device)
+                    delta = q_pos.view(-1, 1) - k_pos.view(1, -1)
+                    blk_mask = (delta >= 0) & (delta < self.context)
+                    out_blk = F.scaled_dot_product_attention(
+                        q[:, :, start:end], k[:, :, k_start:end], v[:, :, k_start:end], attn_mask=blk_mask
+                    )
+                    outs.append(out_blk)
+                out = torch.cat(outs, dim=2)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
         out = out.transpose(1, 2).reshape(B, T, self.embed_dim)
         return self.out_proj(out)
 
@@ -248,6 +291,7 @@ class _TransformerLayer(nn.Module):
         causal: bool,
         max_period: float,
         layer_scale: float | None,
+        context: int | None = None,
         device=None,
         dtype=None,
     ) -> None:
@@ -255,7 +299,7 @@ class _TransformerLayer(nn.Module):
         kw = {"device": device, "dtype": dtype}
         self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **kw)
         self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **kw)
-        self.attn = _Attention(d_model, num_heads, causal, max_period, device, dtype)
+        self.attn = _Attention(d_model, num_heads, causal, max_period, context, device, dtype)
         self.ff1 = nn.Linear(d_model, dim_feedforward, bias=False, **kw)
         self.ff2 = nn.Linear(dim_feedforward, d_model, bias=False, **kw)
         ls_kw = cast(dict[str, object], kw)
@@ -282,13 +326,16 @@ class _Transformer(nn.Module):
         causal: bool,
         max_period: float,
         layer_scale: float | None,
+        context: int | None = None,
         device=None,
         dtype=None,
     ) -> None:
         super().__init__()
         self.layers = nn.ModuleList(
             [
-                _TransformerLayer(d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, device, dtype)
+                _TransformerLayer(
+                    d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, context, device, dtype
+                )
                 for _ in range(num_layers)
             ]
         )
@@ -307,11 +354,12 @@ class _ProjectedTransformer(nn.Module):
     ) -> None:
         super().__init__()
         self.downsample_ratio: int = 1
-        self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if d_model != input_dimension else nn.Identity()
+        # Upstream always materializes a learned Linear here regardless of
+        # whether input_dimension == d_model — substituting Identity when
+        # the dims happen to match silently drops a trained projection.
+        self.in_proj = nn.Linear(input_dimension, d_model, bias=False)
         self.transformer = _Transformer(d_model=d_model, **kwargs)
-        self.out_proj = (
-            nn.Linear(d_model, output_dimension, bias=False) if d_model != output_dimension else nn.Identity()
-        )
+        self.out_proj = nn.Linear(d_model, output_dimension, bias=False)
 
     def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         x = self.in_proj(x.transpose(1, 2))  # (B, D, T) → (B, T, d_model)
@@ -456,10 +504,21 @@ class _ResidualQ(nn.Module):
 
 
 def _build_modules(
-    specs: list[dict[str, Any]], is_downsample: bool, context_duration: float, sampling_rate: float
-) -> nn.ModuleList:
+    specs: list[dict[str, Any]], is_downsample: bool, default_context_duration: float, frame_rate: float
+) -> tuple[nn.ModuleList, float]:
+    """Build encoder/decoder stage modules.
+
+    ``frame_rate`` tracks the running frames-per-second at the current stage
+    (updated after each module by its downsample/upsample ratio, mirroring
+    upstream's ``current_frame_rate``). Each ``Transformer`` module's local
+    attention window (``context``, in frames) is derived from its own
+    ``context_duration`` (seconds) at the rate *at that point in the stack* —
+    not a global constant — so it must be computed here, not by the caller.
+    Returns the final rate too, since the decoder's starting rate is the
+    encoder's bottleneck rate, not the raw sampling rate.
+    """
     modules: list[nn.Module] = []
-    rate = float(sampling_rate)
+    rate = float(frame_rate)
     for spec in specs:
         spec = dict(spec)
         if spec["module_type"] == "PatchedPretransform":
@@ -471,16 +530,18 @@ def _build_modules(
             spec.pop("positional_embedding", None)
             spec.pop("norm", None)
             spec.pop("causal", None)
+            context_duration = float(spec.pop("context_duration", default_context_duration))
             m = _ProjectedTransformer(
                 module_type="Transformer",
                 causal=True,
                 max_period=spec.pop("max_period", 10000),
                 layer_scale=spec.pop("layer_scale", None),
+                context=int(round(rate * context_duration)),
                 **spec,
             )
         modules.append(m)
         rate = rate / m.downsample_ratio if is_downsample else rate * m.downsample_ratio
-    return nn.ModuleList(modules)
+    return nn.ModuleList(modules), rate
 
 
 class MossAudioTokenizerModel(PreTrainedModel):
@@ -494,12 +555,27 @@ class MossAudioTokenizerModel(PreTrainedModel):
         super().__init__(config)
         self.sampling_rate = config.sampling_rate
         self.downsample_rate = config.downsample_rate
+        # Real v1 checkpoints store these as an explicit JSON `null` rather
+        # than omitting the key, so `getattr(..., default)` doesn't apply —
+        # normalize None to the same defaults here.
+        self.number_channels = getattr(config, "number_channels", 1) or 1
+        self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True)
+        if self.enable_channel_interleave is None:
+            self.enable_channel_interleave = True
 
         ctx = config.causal_transformer_context_duration
-        sr = float(config.sampling_rate)
-        self.encoder = _build_modules(config.encoder_kwargs, is_downsample=True, context_duration=ctx, sampling_rate=sr)
-        self.decoder = _build_modules(
-            copy.deepcopy(config.decoder_kwargs), is_downsample=False, context_duration=ctx, sampling_rate=sr
+        channel_interleave_factor = (
+            self.number_channels if (self.enable_channel_interleave and self.number_channels > 1) else 1
+        )
+        frame_rate = float(config.sampling_rate) * channel_interleave_factor
+        self.encoder, bottleneck_rate = _build_modules(
+            config.encoder_kwargs, is_downsample=True, default_context_duration=ctx, frame_rate=frame_rate
+        )
+        self.decoder, _ = _build_modules(
+            copy.deepcopy(config.decoder_kwargs),
+            is_downsample=False,
+            default_context_duration=ctx,
+            frame_rate=bottleneck_rate,
         )
 
         kw = dict(config.quantizer_kwargs)
@@ -512,15 +588,29 @@ class MossAudioTokenizerModel(PreTrainedModel):
         wav_list: list[torch.Tensor],
         num_quantizers: int | None = None,
     ) -> MossAudioTokenizerEncoderOutput:
-        """Encode a list of 1-D waveform tensors → RVQ codes (NQ, B, T)."""
+        """Encode a list of waveform tensors → RVQ codes (NQ, B, T).
+
+        Each element of ``wav_list`` is ``(T,)``/``(1, T)`` for a mono codec
+        (``number_channels == 1``) or ``(number_channels, T)`` otherwise.
+        """
         device = wav_list[0].device
         B = len(wav_list)
-        max_len = max(w.shape[-1] for w in wav_list)
-        x = torch.zeros(B, 1, max_len, device=device)
+        C = self.number_channels
+        normalized: list[torch.Tensor] = []
         lengths = torch.zeros(B, device=device, dtype=torch.long)
         for i, w in enumerate(wav_list):
-            x[i, 0, : w.shape[-1]] = w
-            lengths[i] = w.shape[-1]
+            if C == 1:
+                w_i = w.unsqueeze(0) if w.dim() == 1 else w
+            else:
+                if w.dim() != 2 or w.shape[0] != C:
+                    raise ValueError(f"Expected wav_list[{i}] to have shape ({C}, T), got {tuple(w.shape)}.")
+                w_i = w
+            normalized.append(w_i)
+            lengths[i] = w_i.shape[-1]
+        max_len = int(lengths.max().item())
+        x = torch.zeros(B, C, max_len, device=device, dtype=wav_list[0].dtype)
+        for i, w_i in enumerate(normalized):
+            x[i, :, : w_i.shape[-1]] = w_i
         return self._encode(x, lengths, num_quantizers)
 
     @torch.no_grad()
@@ -541,6 +631,45 @@ class MossAudioTokenizerModel(PreTrainedModel):
             lengths[i] = c.shape[-1]
         return self._decode(codes, lengths)
 
+    def _flatten_channels_for_codec(
+        self,
+        input_values: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Pad to a ``downsample_rate`` multiple, then interleave channels into one stream.
+
+        Mirrors upstream's ``_flatten_channels_for_codec``: for a stereo
+        codec, ``(B, C, T)`` becomes ``(B, 1, T*C)`` with samples interleaved
+        frame-by-frame across channels, so the mono encoder/decoder stack can
+        process a multi-channel waveform unchanged.
+        """
+        if input_values.shape[-1] % self.downsample_rate != 0:
+            pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate)
+            input_values = F.pad(input_values, (0, pad_length))
+        if self.number_channels > 1 and self.enable_channel_interleave:
+            input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1)
+            input_lengths = input_lengths * self.number_channels
+        return input_values, input_lengths
+
+    def _restore_channels_from_codec(
+        self,
+        output_values: torch.Tensor,
+        output_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Inverse of ``_flatten_channels_for_codec``: de-interleave back to ``(B, C, T)``."""
+        if self.number_channels == 1 or not self.enable_channel_interleave:
+            return output_values.float(), output_lengths
+        output_values = (
+            output_values.squeeze(1)
+            .contiguous()
+            .view(output_values.shape[0], -1, self.number_channels)
+            .transpose(1, 2)
+            .contiguous()
+            .float()
+        )
+        output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor")
+        return output_values, output_lengths
+
     def _encode(
         self,
         x: torch.Tensor,
@@ -549,9 +678,7 @@ class MossAudioTokenizerModel(PreTrainedModel):
     ) -> MossAudioTokenizerEncoderOutput:
         if x.dim() == 2:
             x = x.unsqueeze(1)
-        T = x.shape[-1]
-        if T % self.downsample_rate:
-            x = F.pad(x, (0, self.downsample_rate - T % self.downsample_rate))
+        x, lengths = self._flatten_channels_for_codec(x, lengths)
         e, e_len = x, lengths
         for m in self.encoder:
             e, e_len = m(e, e_len)
@@ -567,6 +694,7 @@ class MossAudioTokenizerModel(PreTrainedModel):
         d, d_len = z, lengths
         for m in self.decoder:
             d, d_len = m(d, d_len)
+        d, d_len = self._restore_channels_from_codec(d, d_len)
         return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len)
 
 

vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.qwen3 import Qwen3Config
 
 
@@ -235,7 +236,80 @@ class MossTTSRealtimeConfig(PretrainedConfig):
         return self.language_config
 
 
+class MossTTSLocalConfig(PretrainedConfig):
+    """Config for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+    Like ``MossTTSDelayConfig``, the HF checkpoint stores a nested Qwen3 config
+    (``qwen3_config``) which we unwrap via ``get_text_config()``. It also stores
+    a nested GPT2-style config (``gpt2_config``) describing the single-layer
+    local depth transformer that decodes the ``n_vq`` audio codebooks per frame
+    (see ``modeling_moss_tts_local_depth.py``).
+    """
+
+    model_type = "moss_tts_local"
+
+    def __init__(
+        self,
+        qwen3_config: dict | None = None,
+        gpt2_config: dict | None = None,
+        n_vq: int = 12,
+        audio_vocab_size: int = 1024,
+        audio_pad_token_id: int = 1024,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_start_token_id: int = 151669,
+        audio_end_token_id: int = 151670,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_slot_token_id: int = 151656,
+        sampling_rate: int = 48000,
+        audio_tokenizer_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2",
+        local_text_head_mode: str = "binary",
+        **kwargs: object,
+    ) -> None:
+        if qwen3_config is None:
+            qwen3_config = {}
+        if isinstance(qwen3_config, dict):
+            qwen3_config = dict(qwen3_config)
+            qwen3_config.pop("model_type", None)
+            self.qwen3_config = Qwen3Config(**qwen3_config)
+        else:
+            self.qwen3_config = qwen3_config
+
+        if isinstance(gpt2_config, dict):
+            gpt2_config = dict(gpt2_config)
+            gpt2_config.pop("model_type", None)
+            self.gpt2_config = GPT2Config(**gpt2_config)
+        else:
+            self.gpt2_config = gpt2_config
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.n_vq = n_vq
+        self.audio_vocab_size = audio_vocab_size
+        self.audio_pad_token_id = audio_pad_token_id
+        self.audio_pad_code = audio_pad_token_id
+        self.im_start_token_id = im_start_token_id
+        self.im_end_token_id = im_end_token_id
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.audio_user_slot_token_id = audio_user_slot_token_id
+        self.audio_assistant_slot_token_id = audio_assistant_slot_token_id
+        self.sampling_rate = sampling_rate
+        self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path
+        self.local_text_head_mode = local_text_head_mode
+
+    def get_text_config(self, **_: object) -> Qwen3Config:
+        return self.qwen3_config
+
+
 AutoConfig.register("moss_tts_delay", MossTTSDelayConfig)
 AutoConfig.register("moss_tts_realtime", MossTTSRealtimeConfig)
+AutoConfig.register("moss_tts_local", MossTTSLocalConfig)
 
-__all__ = ["MossTTSDelayConfig", "MossTTSRealtimeConfig", "MossTTSLocalTransformerConfig"]
+__all__ = [
+    "MossTTSDelayConfig",
+    "MossTTSRealtimeConfig",
+    "MossTTSLocalTransformerConfig",
+    "MossTTSLocalConfig",
+]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-import os
 from collections.abc import Iterable
 from typing import Any
 
@@ -42,10 +41,8 @@ class MossTTSCodecDecoder(nn.Module):
     ``n_vq=32`` (MOSS-TTS) or ``n_vq=16`` (all other variants) without
     swapping weights.
 
-    The codec checkpoint path defaults to the value stored in
-    ``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can
-    be overridden by setting the environment variable
-    ``MOSS_TTS_CODEC_PATH``.
+    The codec checkpoint path comes from
+    ``vllm_config.model_config.hf_config.codec_model_name_or_path``.
     """
 
     input_modalities = "audio"
@@ -64,11 +61,19 @@ class MossTTSCodecDecoder(nn.Module):
 
         cfg = vllm_config.model_config.hf_config
         self._n_vq: int = int(getattr(cfg, "n_vq", getattr(cfg, "rvq", 16)))
-        self._codec_path: str = str(getattr(cfg, "codec_model_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"))
+        self._codec_path: str = str(
+            getattr(
+                cfg,
+                "codec_model_name_or_path",
+                getattr(cfg, "audio_tokenizer_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"),
+            )
+        )
 
-        # Resolved at load_weights() time
+        # Resolved at load_weights() time, once the codec checkpoint's own
+        # config (sampling rate, channel count) is known.
         self._codec: MossAudioTokenizerModel | None = None
         self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None
+        self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
 
     # ------------------------------------------------------------------
@@ -119,7 +124,7 @@ class MossTTSCodecDecoder(nn.Module):
         """
         sr_tensor = self._sr_tensor
         empty = torch.zeros((0,), dtype=torch.float32)
-        info_list: list[dict[str, Any]] = runtime_additional_information or []
+        info_list: list[dict[str, Any]] = runtime_additional_information or [{}]
         num_req = max(len(info_list), 1)
 
         if self._codec is None:
@@ -147,7 +152,7 @@ class MossTTSCodecDecoder(nn.Module):
         # kwargs if available, else assume one request.
         ids_flat = input_ids.reshape(-1).to(dtype=torch.long)
         num_scheduled_tokens = kwargs.get("num_scheduled_tokens")
-        if isinstance(num_scheduled_tokens, list) and len(num_scheduled_tokens) == num_req:
+        if isinstance(num_scheduled_tokens, (list, tuple)) and len(num_scheduled_tokens) == num_req:
             offsets = [0]
             for n in num_scheduled_tokens:
                 offsets.append(offsets[-1] + int(n))
@@ -176,12 +181,8 @@ class MossTTSCodecDecoder(nn.Module):
             codebook_size = self._codec.config.codebook_size
             codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1)
 
-            # left_context_size: number of left-context frames prepended to
-            # this chunk by the stage input processor (overlap-add pattern,
-            # same as Qwen3-TTS code2wav).  We decode them together with the
-            # real chunk for smooth boundaries, then trim them from the output.
             meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
-            left_ctx = meta.get("left_context_size", 0)
+            left_ctx = self._runtime_value(info, meta, "left_context_size", 0)
             if isinstance(left_ctx, (list, tuple)):
                 left_ctx = int(left_ctx[0]) if left_ctx else 0
             elif isinstance(left_ctx, torch.Tensor):
@@ -196,28 +197,40 @@ class MossTTSCodecDecoder(nn.Module):
             if out.audio is None:
                 continue
 
-            wav = out.audio.reshape(-1).to(dtype=torch.float32).cpu()
+            # ``out.audio`` is ``(1, C, T)``; keep the channel axis for
+            # stereo codecs (Local-v1.5) and flatten to ``(T,)`` for mono
+            # ones (Delay/Realtime) to preserve their existing output shape.
+            wav = out.audio[0].to(dtype=torch.float32).cpu()
             if out.audio_lengths is not None:
-                wav = wav[: int(out.audio_lengths[0].item())]
+                wav = wav[..., : int(out.audio_lengths[0].item())]
 
-            # Trim left-context samples.
+            # Trim left-context samples (per-channel sample axis, so the
+            # trim amount is identical for mono and interleaved-stereo).
             if left_ctx > 0:
-                trim = min(left_ctx * self._codec.downsample_rate, wav.shape[0])
+                trim = min(left_ctx * self._codec.downsample_rate, wav.shape[-1])
                 if trim < left_ctx * self._codec.downsample_rate:
                     logger.warning(
                         "left_ctx trim (%d samples) exceeds wav length (%d); returning empty audio.",
                         left_ctx * self._codec.downsample_rate,
-                        wav.shape[0],
+                        wav.shape[-1],
                     )
-                wav = wav[trim:]
+                wav = wav[..., trim:]
 
-            audios[i] = wav
+            audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav
 
         return OmniOutput(
             text_hidden_states=None,
             multimodal_outputs={"model_outputs": audios, "sr": srs},
         )
 
+    @staticmethod
+    def _runtime_value(info: Any, meta: dict[str, Any], name: str, default: Any = None) -> Any:
+        if name in meta:
+            return meta[name]
+        if isinstance(info, dict) and name in info:
+            return info[name]
+        return default
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -234,11 +247,10 @@ class MossTTSCodecDecoder(nn.Module):
         for _ in weights:
             pass
 
-        codec_path = os.environ.get("MOSS_TTS_CODEC_PATH", self._codec_path)
+        codec_path = self._codec_path
         logger.info("Loading MOSS Audio Tokenizer from %s", codec_path)
 
-        codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
-        codec = MossAudioTokenizerModel(codec_cfg)
+        codec_cfg, codec = self._build_codec(codec_path)
 
         model_loader = DefaultModelLoader(self.vllm_config.load_config)
         source = DefaultModelLoader.Source(
@@ -255,10 +267,19 @@ class MossTTSCodecDecoder(nn.Module):
         # and the rest stay at their random init, which produces noise that
         # sounds correct in duration but is structurally garbage.
         _SUFFIX_REMAP: list[tuple[str, str]] = [
+            # v1 (MOSS-Audio-Tokenizer) naming.
             (".self_attn.in_projs.0.", ".attn.in_proj."),
             (".self_attn.out_projs.0.", ".attn.out_proj."),
             (".linear1.", ".ff1."),
             (".linear2.", ".ff2."),
+            # v2 (MOSS-Audio-Tokenizer-v2) uses different submodule names for
+            # the same attention/FFN sublayers — no trailing "s"/index on
+            # in_proj/out_proj, and an `ffn.{0,2}` Sequential instead of
+            # separate linear1/linear2 attributes.
+            (".self_attn.in_proj.", ".attn.in_proj."),
+            (".self_attn.out_proj.", ".attn.out_proj."),
+            (".ffn.0.", ".ff1."),
+            (".ffn.2.", ".ff2."),
             (".layer_scale_1.", ".ls1."),
             (".layer_scale_2.", ".ls2."),
             (".input_proj.", ".in_proj."),
@@ -295,11 +316,14 @@ class MossTTSCodecDecoder(nn.Module):
         codec.to(device=device, dtype=torch.float32)
         codec.eval()
         self._codec = codec
+        self._n_channels = int(getattr(codec_cfg, "number_channels", 1) or 1)
+        self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32)
 
         logger.info(
-            "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d",
+            "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d, n_channels=%d",
             codec_cfg.sampling_rate,
             codec_cfg.num_quantizers,
+            self._n_channels,
         )
 
         self._maybe_enable_decoder_cudagraph(device)
@@ -310,6 +334,24 @@ class MossTTSCodecDecoder(nn.Module):
         # mirror that here.
         return {f"_codec.{name}" for name, _ in codec.named_parameters()}
 
+    def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]:
+        try:
+            from transformers import AutoConfig, AutoModel
+
+            codec_cfg = AutoConfig.from_pretrained(codec_path, trust_remote_code=True)
+            codec = AutoModel.from_config(codec_cfg, trust_remote_code=True)
+            logger.info("Using MOSS Audio Tokenizer remote-code classes from %s", codec_path)
+            return codec_cfg, codec
+        except Exception:
+            logger.exception(
+                "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; "
+                "falling back to vendored codec."
+            )
+
+        codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
+        codec = MossAudioTokenizerModel(codec_cfg)
+        return codec_cfg, codec
+
     def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None:
         """Capture CUDA Graphs for the codec decoder if enforce_eager is False."""
         if getattr(self.vllm_config.model_config, "enforce_eager", True):

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -0,0 +1,223 @@
+# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Per-frame depth transformer for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+A 1-layer GPT2-style block that decodes the ``n_vq`` audio codebook codes for
+one audio frame, run inside the talker's ``make_omni_output`` independent of
+vLLM's main scheduler -- mirrors ``MossTTSRealtimeLocalTransformer``
+(``modeling_moss_tts_local.py``) in role, but the algorithm and numerics are
+faithful to the official ``gpt2_decoder.py`` / ``modeling_moss_tts.py``
+(GPT2-style LayerNorm + bias, SiLU MLP, **interleaved/GPT-J-style RoPE** --
+not vLLM's neox-style concat-half rotation) rather than Realtime's
+Qwen3-style ``CodePredictorBaseModel``.
+
+Its KV cache and positions reset to 0 every audio frame: there is no
+cross-frame state. Submodule names (``h.0.*`` / ``ln_f``) match the
+checkpoint 1:1 so ``load_weights()`` needs no remapping.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import (
+    _sample_token,
+)
+
+
+class _MossTTSLocalAttention(nn.Module):
+    """GPT2-style fused-QKV self-attention with interleaved RoPE.
+
+    Faithful to the official ``MossTTSNanoGPT2Attention``: ``rotate_half``
+    operates on even/odd index pairs (GPT-J style), and ``cos``/``sin`` are
+    built via ``repeat_interleave(2, dim=-1)`` rather than the neox-style
+    concat-half construction vLLM uses elsewhere.
+    """
+
+    def __init__(self, hidden_size: int, n_head: int, rope_base: float) -> None:
+        super().__init__()
+        if hidden_size % n_head != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by n_head={n_head}")
+        self.n_head = n_head
+        self.head_dim = hidden_size // n_head
+        self.embed_dim = hidden_size
+        self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+        inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _rope_cos_sin(
+        self, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device))
+        cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)
+        sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)
+        return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim)
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        even = x[..., ::2]
+        odd = x[..., 1::2]
+        return torch.stack((-odd, even), dim=-1).reshape_as(x)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call."""
+        batch_size, seq_len, _ = hidden_states.shape
+        qkv = self.c_attn(hidden_states)
+        query, key, value = qkv.split(self.embed_dim, dim=-1)
+        query = query.view(batch_size, seq_len, self.n_head, self.head_dim)
+        key = key.view(batch_size, seq_len, self.n_head, self.head_dim)
+        value = value.view(batch_size, seq_len, self.n_head, self.head_dim)
+
+        cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype)
+        query = query * cos + self._rotate_half(query) * sin
+        key = key * cos + self._rotate_half(key) * sin
+
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True)
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
+        return self.c_proj(attn_output)
+
+
+class _MossTTSLocalMLP(nn.Module):
+    def __init__(self, hidden_size: int, inner_size: int) -> None:
+        super().__init__()
+        self.fc_in = nn.Linear(hidden_size, inner_size, bias=True)
+        self.fc_out = nn.Linear(inner_size, hidden_size, bias=True)
+        self.act = nn.SiLU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.fc_out(self.act(self.fc_in(hidden_states)))
+
+
+class _MossTTSLocalBlock(nn.Module):
+    def __init__(self, hidden_size: int, n_head: int, inner_size: int, rope_base: float, eps: float) -> None:
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=eps)
+        self.attn = _MossTTSLocalAttention(hidden_size, n_head, rope_base)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=eps)
+        self.mlp = _MossTTSLocalMLP(hidden_size, inner_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))
+        hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
+        return hidden_states
+
+
+class MossTTSLocalDepthTransformer(nn.Module):
+    """Per-frame depth transformer for MOSS-TTS-Local-Transformer-v1.5.
+
+    Per frame:
+      - position 0's input is the backbone's last hidden state; its output
+        feeds BOTH the binary continue/stop head (``local_text_lm_head``)
+        and codebook-0's head (``audio_lm_heads[0]``) simultaneously.
+      - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is
+        re-embedded (``audio_embeddings[c]``) and appended as the next
+        position, re-prefilling the block over the growing (<=n_vq) sequence
+        with a fresh causal mask each call -- mathematically identical to
+        incremental KV-cache decoding since attention is strictly causal and
+        only the last position is ever read, but avoids any cache plumbing
+        given the trivially short sequence length.
+    """
+
+    def __init__(self, gpt2_config, hidden_size: int | None = None) -> None:
+        super().__init__()
+        self.hidden_size = int(hidden_size if hidden_size is not None else gpt2_config.n_embd)
+        n_head = int(gpt2_config.n_head)
+        inner_size = int(gpt2_config.n_inner)
+        eps = float(getattr(gpt2_config, "layer_norm_epsilon", 1e-5))
+        rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0))
+        self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)])
+        self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
+
+    def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
+        hidden_states = seq_embeds
+        for block in self.h:
+            hidden_states = block(hidden_states)
+        return self.ln_f(hidden_states)
+
+    @torch.no_grad()
+    def generate_frame(
+        self,
+        backbone_last_hidden: torch.Tensor,  # (B, H)
+        audio_lm_heads: nn.ModuleList,  # n_vq x Linear(H -> audio_vocab_size)
+        audio_embeddings: nn.ModuleList,  # n_vq x Embedding(audio_vocab_size, H)
+        local_text_lm_head: nn.Module,  # Linear(H -> 2): [continue, stop]
+        *,
+        n_vq: int,
+        do_sample: bool = True,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        text_temperature: float = 1.0,
+        text_top_k: int = 50,
+        text_top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+        history_per_codebook: list[list[int]] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Generate one audio frame for batch B.
+
+        Returns ``(should_continue, codes)``: ``should_continue`` is a
+        ``(B,)`` bool tensor (``True`` iff the binary head picked the
+        "continue" candidate, i.e. logits index 0); ``codes`` is a
+        ``(B, n_vq)`` LongTensor of sampled codebook indices.
+
+        ``history_per_codebook[c]`` is a list of recently-emitted token ids
+        for codebook ``c``; when ``repetition_penalty != 1.0`` those tokens'
+        logits get scaled down before sampling (mirrors upstream's
+        ``_apply_repetition_penalty``).
+        """
+        batch_size = backbone_last_hidden.shape[0]
+        dtype = self.ln_f.weight.dtype
+
+        embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
+        embeds[:, 0, :] = backbone_last_hidden.to(dtype)
+
+        hidden = self._forward_prefix(embeds[:, :1, :])
+        local_hidden = hidden[:, 0, :]
+
+        binary_logits = local_text_lm_head(local_hidden).float()
+        binary_choice = _sample_token(binary_logits, text_temperature, text_top_k, text_top_p, do_sample)
+        should_continue = binary_choice.eq(0)
+        import os as _os
+
+        if _os.environ.get("MOSS_TTS_DEBUG_STOP"):
+            import logging as _logging
+
+            _logging.getLogger("moss_tts_debug").warning(
+                "binary_logits=%s choice=%s", binary_logits.tolist(), binary_choice.tolist()
+            )
+
+        codes = backbone_last_hidden.new_zeros((batch_size, n_vq), dtype=torch.long)
+        for channel_index in range(n_vq):
+            channel_logits = audio_lm_heads[channel_index](local_hidden).float()
+            if (
+                repetition_penalty != 1.0
+                and history_per_codebook is not None
+                and channel_index < len(history_per_codebook)
+            ):
+                hist = history_per_codebook[channel_index]
+                if hist:
+                    hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device)
+                    sel = channel_logits.index_select(-1, hist_t)
+                    pos = sel > 0
+                    sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
+                    channel_logits.index_copy_(-1, hist_t, sel)
+            channel_token = _sample_token(channel_logits, temperature, top_k, top_p, do_sample)
+            codes[:, channel_index] = channel_token
+
+            if channel_index + 1 < n_vq:
+                embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype)
+                hidden = self._forward_prefix(embeds[:, : channel_index + 2, :])
+                local_hidden = hidden[:, channel_index + 1, :]
+
+        return should_continue, codes
+
+
+__all__ = ["MossTTSLocalDepthTransformer"]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from typing import Any
 
 import torch
@@ -22,11 +22,15 @@ from vllm.v1.sample.metadata import SamplingMetadata
 
 from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
     MossTTSDelayConfig,
+    MossTTSLocalConfig,
     MossTTSRealtimeConfig,
 )
 from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import (
     MossTTSRealtimeLocalTransformer,
 )
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+    MossTTSLocalDepthTransformer,
+)
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 logger = init_logger(__name__)
@@ -36,6 +40,40 @@ def _maybe_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
 
 
+def _iter_state_row_spans(
+    states: Sequence[dict[str, Any]],
+    spans: Sequence[tuple[int, int]] | None,
+    num_rows: int,
+) -> Iterable[tuple[dict[str, Any], int, int]]:
+    """Yield state/logit-row ranges without assuming equal request row counts."""
+    if num_rows <= 0 or not states:
+        return
+
+    if spans is not None:
+        if len(spans) != len(states):
+            raise RuntimeError(f"request_token_spans has {len(spans)} entries for {len(states)} request states")
+        if all(row_end <= num_rows for _, row_end in spans):
+            for state, (row_start, row_end) in zip(states, spans, strict=False):
+                if row_start < row_end:
+                    yield state, int(row_start), int(row_end)
+            return
+        if num_rows != len(states):
+            raise RuntimeError(
+                "request_token_spans describe full hidden rows, but current logits "
+                f"have {num_rows} rows for {len(states)} states and are not one-row-per-request"
+            )
+
+    if num_rows == len(states):
+        for i, state in enumerate(states):
+            yield state, i, i + 1
+        return
+
+    raise RuntimeError(
+        "MOSS-TTS compute_logits requires request_token_spans or one-row-per-request logits; "
+        f"got {num_rows} rows for {len(states)} request states"
+    )
+
+
 # ---------------------------------------------------------------------------
 # MossTTSDelayTalkerForGeneration
 # ---------------------------------------------------------------------------
@@ -209,14 +247,13 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         n_vq = self.n_vq
         device = logits.device
         vocab_size = logits.shape[-1]
-        rows_per_state = max(1, logits.shape[0] // max(1, len(states)))
 
-        for i, state in enumerate(states):
-            if not isinstance(state, dict):
-                continue
-            row_start = i * rows_per_state
-            row_end = min(row_start + rows_per_state, logits.shape[0])
-            if row_start >= row_end:
+        for state, row_start, row_end in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            logits.shape[0],
+        ):
+            if state is None:
                 continue
             row = logits[row_start:row_end]
 
@@ -569,6 +606,7 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         """
         if isinstance(model_outputs, OmniOutput):
             self._batch_state = None
+            self._batch_state_spans = None
             return model_outputs
 
         hidden = model_outputs  # (S, H)
@@ -590,11 +628,11 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         # Stash the per-row state for compute_logits to apply masks. logits
         # rows align with hidden rows, which align with info_dicts in order.
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+        self._batch_state_spans = kwargs.get("request_token_spans")
 
         # Per-request (start, end) hidden-row spans from the runner. In mixed
         # prefill+decode steps the per-request token counts differ, so an equal
-        # rows-per-request split would sample codes from the wrong rows; fall
-        # back to that split only when spans are unavailable.
+        # rows-per-request split would sample codes from the wrong rows.
         spans = kwargs.get("request_token_spans")
 
         # One accumulated-codes tensor per request, in batch order. The generic
@@ -607,18 +645,18 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
 
         if hidden.numel() > 0:
             num_rows = hidden.shape[0]
-            rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
+            if spans is None or len(spans) != len(info_dicts):
+                raise RuntimeError(
+                    "MOSS-TTS make_omni_output requires request_token_spans "
+                    f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}"
+                )
 
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start, row_end = spans[i]
+                row_end = min(int(row_end), num_rows)
                 if row_start >= row_end:
                     per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
@@ -864,13 +902,12 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         logits = hidden_states.new_full((B, V), float("-inf"))
 
         states = self._batch_state or []
-        rows_per_state = max(1, B // max(1, len(states) or 1))
-        for i, state in enumerate(states):
-            r0 = i * rows_per_state
-            r1 = min(r0 + rows_per_state, B)
-            if r0 >= r1:
-                continue
-            if not isinstance(state, dict):
+        for state, r0, r1 in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            B,
+        ):
+            if state is None:
                 logits[r0:r1, self.text_pad_id] = 0.0
                 continue
             if state.get("is_stopping"):
@@ -991,6 +1028,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
     ) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             self._batch_state = None
+            self._batch_state_spans = None
             return model_outputs
 
         hidden = model_outputs  # (S, H)
@@ -1004,9 +1042,10 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                 info["audio_state"] = {"is_stopping": False, "step": 0}
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+        self._batch_state_spans = kwargs.get("request_token_spans")
 
-        # See delay talker: real per-request row spans from the runner, with the
-        # equal-split as fallback when unavailable.
+        # See delay talker: real per-request row spans from the runner are
+        # required because mixed prefill+decode steps have unequal row counts.
         spans = kwargs.get("request_token_spans")
         # Per-request accumulated codes in batch order, pre-filled with empty
         # placeholders so every skip path (stopped / eos / bos / pad) keeps
@@ -1016,16 +1055,16 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         have_codes = False
         if hidden.numel() > 0 and info_dicts:
             num_rows = hidden.shape[0]
-            rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
+            if spans is None or len(spans) != len(info_dicts):
+                raise RuntimeError(
+                    "MOSS-TTS make_omni_output requires request_token_spans "
+                    f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}"
+                )
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     continue
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start, row_end = spans[i]
+                row_end = min(int(row_end), num_rows)
                 if row_start >= row_end:
                     continue
 
@@ -1189,7 +1228,430 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         return loaded
 
 
+# ---------------------------------------------------------------------------
+# MossTTSLocalTalkerForGeneration
+# ---------------------------------------------------------------------------
+
+
+class MossTTSLocalTalkerForGeneration(nn.Module):
+    """Stage-0 talker for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+    Architecture differs from both existing variants:
+
+    * Backbone (Qwen3; checkpoint prefix ``transformer.*``, unlike Delay/
+      Realtime's ``language_model.*``) consumes additive text+audio fusion
+      the same way Delay does (real text embeddings + masked per-codebook
+      audio embeddings from ``info["codes"]["ref"]`` / ``info["audio_codes"]
+      ["current"]``), but there is no delay-pattern countdown: the prompt
+      already ends right at the audio_start boundary, and every step from
+      the first decode onward drives exactly one audio frame.
+    * The checkpoint's own ``text_lm_head`` is loaded (trained for a
+      teacher-forced text loss) but never consulted at inference: the
+      official reference's own ``generate()`` hardcodes the next
+      text-channel input to ``audio_assistant_slot_token_id`` regardless of
+      what ``text_lm_head`` would predict, because
+      ``local_text_head_mode == "binary"`` for this checkpoint.
+      ``compute_logits`` therefore synthesises a one-hot logit row
+      (Realtime-style), forced to ``audio_assistant_slot_token_id`` while
+      continuing or ``im_end_token_id`` once the local transformer's binary
+      head decides to stop.
+    * Per-step audio generation runs ``MossTTSLocalDepthTransformer``
+      (1-layer GPT2-style local transformer, interleaved RoPE, ``n_vq=12``
+      codebooks -- structurally different from Realtime's Qwen3-style
+      ``CodePredictorBaseModel``) inside ``make_omni_output``. Stop
+      condition: the local transformer's binary head selects the
+      "audio_end" candidate (logits index 1).
+    """
+
+    have_multimodal_outputs: bool = True
+    has_preprocess: bool = True
+    has_postprocess: bool = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config: MossTTSLocalConfig = vllm_config.model_config.hf_config
+
+        qwen3_cfg = self.config.qwen3_config
+        hidden_size = int(qwen3_cfg.hidden_size)
+        self.hidden_size = hidden_size
+        self.n_vq: int = int(self.config.n_vq)
+        self.audio_vocab_size: int = int(self.config.audio_vocab_size)
+        self.audio_pad_token_id: int = int(self.config.audio_pad_token_id)
+        self.text_vocab_size: int = int(qwen3_cfg.vocab_size)
+
+        self.audio_assistant_slot_token_id: int = int(self.config.audio_assistant_slot_token_id)
+        self.im_end_token_id: int = int(self.config.im_end_token_id)
+
+        # Qwen3 backbone. get_text_config() is sufficient for vLLM to size
+        # KV cache / heads correctly, as long as AutoConfig resolves to the
+        # local MossTTSLocalConfig rather than the HF remote-code class.
+        self.model = Qwen3Model(
+            vllm_config=vllm_config,
+            prefix=_maybe_prefix(prefix, "model"),
+        )
+
+        # Real text LM head -- loaded from the checkpoint but never consulted
+        # at inference (see class docstring); kept purely as a destination
+        # for ``text_lm_head.weight`` so load_weights() loads every tensor.
+        self.text_lm_head = ParallelLMHead(
+            self.text_vocab_size,
+            hidden_size,
+            bias=False,
+            prefix=_maybe_prefix(prefix, "text_lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(self.text_vocab_size)
+
+        # Per-codebook audio embeddings/heads. audio_pad_token_id (==
+        # audio_vocab_size) is outside the embedding table -- masked to zero
+        # contribution rather than stored as a 1025th row (mirrors upstream
+        # ``_build_inputs_embeds``'s ``valid_mask`` masking).
+        self.audio_embeddings = nn.ModuleList(
+            [nn.Embedding(self.audio_vocab_size, hidden_size) for _ in range(self.n_vq)]
+        )
+        self.audio_lm_heads = nn.ModuleList(
+            [nn.Linear(hidden_size, self.audio_vocab_size, bias=False) for _ in range(self.n_vq)]
+        )
+        self.local_text_lm_head = nn.Linear(hidden_size, 2, bias=False)
+        self.local_transformer = MossTTSLocalDepthTransformer(self.config.gpt2_config, hidden_size=hidden_size)
+
+        self._batch_state: list[dict[str, Any]] | None = None
+        # Stacked embedding cache built after load_weights() for vectorised
+        # additive-fusion gathers (avoids an n_vq-iteration Python loop).
+        self._stacked_audio_emb_w: torch.Tensor | None = None
+
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
+            ("audio_codes", "current"),
+            ("audio_codes", "accumulated"),
+            ("hidden_states", "last"),
+        }
+
+    # ------------------------------------------------------------------
+    # vLLM hooks
+    # ------------------------------------------------------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        return self.model.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **_: Any,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: SamplingMetadata | None = None,
+    ) -> torch.Tensor | None:
+        """Synthesise a one-hot text logit row per request.
+
+        The text channel is never freely generated here: the upstream
+        reference hardcodes it to ``audio_assistant_slot_token_id`` while
+        audio frames keep coming, and only emits ``im_end_token_id`` once the
+        local transformer's binary head decides to stop (see class
+        docstring for why the real ``text_lm_head`` is bypassed).
+        """
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hidden_states is None or hidden_states.numel() == 0:
+            return None
+
+        num_rows = hidden_states.shape[0]
+        logits = hidden_states.new_full((num_rows, self.text_vocab_size), float("-inf"))
+
+        states = self._batch_state or []
+        for state, row_start, row_end in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            num_rows,
+        ):
+            if state is not None and state.get("is_stopping"):
+                logits[row_start:row_end, self.im_end_token_id] = 0.0
+            else:
+                logits[row_start:row_end, self.audio_assistant_slot_token_id] = 0.0
+        if not states:
+            logits[:, self.audio_assistant_slot_token_id] = 0.0
+        return logits
+
+    # ------------------------------------------------------------------
+    # Embedding (text + audio codebooks, additive)
+    # ------------------------------------------------------------------
+
+    def _audio_embed(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes``: ``(T, n_vq)`` long, possibly containing ``audio_pad_token_id``
+        (outside the valid embedding range). Returns the ``(T, H)`` additive
+        audio embedding, masking out pad positions per codebook.
+        """
+        device = codes.device
+        valid_mask = codes.ne(self.audio_pad_token_id)  # (T, n_vq)
+        safe_codes = codes.masked_fill(~valid_mask, 0).clamp(0, self.audio_vocab_size - 1)
+        if self._stacked_audio_emb_w is not None:
+            n_vq_idx = torch.arange(self.n_vq, device=device)
+            codes_t = safe_codes.t()  # (n_vq, T)
+            gathered = self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t]  # (n_vq, T, H)
+            gathered = gathered * valid_mask.t().unsqueeze(-1)
+            return gathered.sum(0)
+        embeds: torch.Tensor | None = None
+        for i, emb_layer in enumerate(self.audio_embeddings):
+            col = emb_layer(safe_codes[:, i]) * valid_mask[:, i].unsqueeze(-1)
+            embeds = col if embeds is None else embeds + col
+        return embeds
+
+    def preprocess(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """Build per-step input embeddings (text + audio additive fusion).
+
+        Prefill: initialise the per-request ``{is_stopping, step}`` state.
+        Decode: combine the forced text token (from ``compute_logits``) with
+        the audio codes the local transformer produced last step.
+        """
+        device = input_ids.device
+        span_len = int(input_ids.shape[0])
+        audio_state = info_dict.get("audio_state")
+        is_first_call = not isinstance(audio_state, dict)
+
+        if span_len > 1 or is_first_call:
+            embeds = self.model.embed_tokens(input_ids)
+
+            ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
+            ref_offset = int(info_dict.get("ref_offset", 0))
+            last_ref_row = None
+            if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
+                if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
+                    ref_codes = ref_codes.view(-1, self.n_vq)
+                if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2:
+                    end_off = ref_offset + span_len
+                    chunk = ref_codes[ref_offset:end_off]
+                    if chunk.numel() > 0 and chunk.shape[0] == span_len:
+                        codes = chunk.to(device=device, dtype=torch.long)
+                        embeds = embeds + self._audio_embed(codes)
+                        last_ref_row = codes[-1].detach()
+
+            current_codes = (
+                last_ref_row
+                if last_ref_row is not None
+                else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device)
+            )
+
+            max_new_frames_req = info_dict.get("max_new_frames")
+            if isinstance(max_new_frames_req, (list, tuple)) and max_new_frames_req:
+                max_new_frames_req = max_new_frames_req[0]
+            try:
+                max_new_frames = int(max_new_frames_req) if max_new_frames_req is not None else -1
+            except (TypeError, ValueError):
+                max_new_frames = -1
+
+            info_update: dict[str, Any] = {
+                "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames},
+                "audio_codes": {"current": current_codes},
+                "ref_offset": ref_offset + span_len,
+            }
+            return input_ids, embeds, info_update
+
+        # Decode step: input_ids is the forced text token from compute_logits.
+        text_embed = self.model.embed_tokens(input_ids.reshape(1))  # (1, H)
+        prev_codes = (info_dict.get("audio_codes", {}) or {}).get("current")
+        if isinstance(prev_codes, torch.Tensor) and prev_codes.numel() == self.n_vq:
+            codes = prev_codes.to(device=device, dtype=torch.long).unsqueeze(0)  # (1, n_vq)
+            combined = text_embed + self._audio_embed(codes)
+        else:
+            combined = text_embed
+        return input_ids, combined, {}
+
+    def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
+        if hidden_states.numel() == 0:
+            return {}
+        return {"hidden_states": {"last": hidden_states[-1].detach()}}
+
+    # ------------------------------------------------------------------
+    # Per-step audio generation via local transformer
+    # ------------------------------------------------------------------
+
+    def make_omni_output(
+        self,
+        model_outputs: torch.Tensor | OmniOutput,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            self._batch_state = None
+            self._batch_state_spans = None
+            return model_outputs
+
+        hidden = model_outputs  # (S, H)
+        info_dicts: list[dict[str, Any]] = (
+            kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
+        )
+
+        for info in info_dicts:
+            if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict):
+                info["audio_state"] = {"is_stopping": False, "step": 0, "max_new_frames": -1}
+
+        self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+        self._batch_state_spans = kwargs.get("request_token_spans")
+
+        spans = kwargs.get("request_token_spans")
+        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
+        have_codes = False
+
+        if hidden.numel() > 0 and info_dicts:
+            num_rows = hidden.shape[0]
+            if spans is None or len(spans) != len(info_dicts):
+                raise RuntimeError(
+                    "MOSS-TTS make_omni_output requires request_token_spans "
+                    f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}"
+                )
+            for i, info in enumerate(info_dicts):
+                if not isinstance(info, dict):
+                    continue
+                row_start, row_end = spans[i]
+                row_end = min(int(row_end), num_rows)
+                if row_start >= row_end:
+                    continue
+
+                state = info.get("audio_state", {}) or {}
+                if state.get("is_stopping"):
+                    continue  # already stopped -- no more audio frames
+
+                last_h = hidden[row_end - 1].unsqueeze(0)  # (1, H)
+
+                # Sampling parameters mirror Realtime's defaults; repetition
+                # penalty applies over a 50-frame rolling window per codebook.
+                rep_window = 50
+                acc_for_hist = (info.get("audio_codes", {}) or {}).get("accumulated")
+                if isinstance(acc_for_hist, torch.Tensor) and acc_for_hist.numel() > 0:
+                    tail = acc_for_hist[-rep_window:].long().cpu().tolist()
+                    hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)]
+                else:
+                    hist_per_cb = [[] for _ in range(self.n_vq)]
+
+                should_continue_t, new_codes = self.local_transformer.generate_frame(
+                    last_h,
+                    self.audio_lm_heads,
+                    self.audio_embeddings,
+                    self.local_text_lm_head,
+                    n_vq=self.n_vq,
+                    do_sample=True,
+                    temperature=0.8,
+                    top_k=30,
+                    top_p=0.6,
+                    repetition_penalty=1.1,
+                    history_per_codebook=hist_per_cb,
+                )
+                should_continue = bool(should_continue_t[0].item())
+                new_codes = new_codes.squeeze(0)  # (n_vq,)
+
+                step = int(state.get("step", 0))
+                max_new_frames = int(state.get("max_new_frames", -1))
+                if not should_continue or (0 <= max_new_frames <= step):
+                    state["is_stopping"] = True
+                    state["step"] = step + 1
+                    info["audio_codes"] = {
+                        "current": new_codes,
+                        "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
+                    }
+                    continue  # don't append the stop frame to accumulated
+
+                acc = (info.get("audio_codes", {}) or {}).get("accumulated")
+                if isinstance(acc, torch.Tensor) and acc.numel() > 0:
+                    updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0)
+                else:
+                    updated_acc = new_codes.unsqueeze(0)
+
+                info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
+                state["step"] = step + 1
+                per_req_codes[i] = updated_acc
+                have_codes = True
+
+        if not have_codes:
+            return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+
+        # Local-v1.5 emits raw codes (no delay pattern), so signal that to the
+        # chunk processor the same way Realtime does, via ``meta.finished``
+        # (read as "skip de-delay" by ``talker2codec_async_chunk``).
+        device = hidden.device
+        return OmniOutput(
+            text_hidden_states=hidden,
+            multimodal_outputs={
+                "codes": {"audio": per_req_codes},
+                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
+            },
+        )
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Map checkpoint weight names to vLLM-Omni module names.
+
+        Checkpoint layout (prefix ``transformer.*`` for the backbone --
+        unlike Delay/Realtime's ``language_model.*``):
+          transformer.*               → model.*            (Qwen3 backbone)
+          text_lm_head.weight         → text_lm_head.weight
+          audio_embeddings.{i}.weight → audio_embeddings.{i}.weight
+          audio_lm_heads.{i}.weight   → audio_lm_heads.{i}.weight
+          local_text_lm_head.weight   → local_text_lm_head.weight
+          local_transformer.h.0.*     → local_transformer.h.0.*
+          local_transformer.ln_f.*    → local_transformer.ln_f.*
+
+        Every "tied" pair in the upstream config is stored as a full,
+        separate tensor in this checkpoint (confirmed via the raw
+        safetensors header: 438 keys, no dedup), so aside from the backbone
+        prefix this is pure 1:1 name-based loading.
+        """
+        loaded: set[str] = set()
+        params_dict = dict(self.named_parameters())
+
+        backbone_weights: list[tuple[str, torch.Tensor]] = []
+        skipped: list[str] = []
+
+        for name, tensor in weights:
+            if name.startswith("transformer."):
+                backbone_weights.append((name[len("transformer.") :], tensor))
+                continue
+            if name in params_dict:
+                default_weight_loader(params_dict[name], tensor)
+                loaded.add(name)
+            else:
+                skipped.append(name)
+
+        backbone_loaded = self.model.load_weights(iter(backbone_weights))
+        for n in backbone_loaded:
+            loaded.add(f"model.{n}")
+
+        self._stacked_audio_emb_w = torch.stack(
+            [e.weight.detach() for e in self.audio_embeddings], dim=0
+        )  # (n_vq, audio_vocab_size, hidden_size)
+
+        logger.info(
+            "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)",
+            len(loaded),
+            len(params_dict),
+            len(skipped),
+            skipped[:5],
+        )
+        not_loaded = [n for n in params_dict if n not in loaded]
+        if not_loaded:
+            logger.warning("[MossTTSLocal] %d params NOT loaded (first 5: %s)", len(not_loaded), not_loaded[:5])
+        return loaded
+
+
 __all__ = [
     "MossTTSDelayTalkerForGeneration",
     "MossTTSRealtimeTalkerForGeneration",
+    "MossTTSLocalTalkerForGeneration",
 ]

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -80,9 +80,39 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig(
     ),
 )
 
-# The pipeline config is otherwise the same for all 5 variants; the per-variant
-# differences (n_vq, backbone size, generation strategy) are encoded in the
-# HF config.json and the deploy YAML. Realtime is split out because it has a
-# different talker architecture (MossTTSRealtime vs MossTTSDelayModel).
+MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
+    model_type="moss_tts_local",
+    model_arch="MossTTSLocalModel",  # different talker class: GPT2-style local depth transformer
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="moss_tts_local",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
+            sampling_constraints={"detokenize": False},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="moss_tts_local_codec",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="MossTTSCodecDecoder",
+            sync_process_input_func=f"{_PROC}.talker2codec",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
 
-__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE"]
+# The pipeline config is otherwise the same for all variants; the per-variant
+# differences (n_vq, backbone size, generation strategy) are encoded in the
+# HF config.json and the deploy YAML. Realtime and Local are split out because
+# they have different talker architectures from the delay variant
+# (MossTTSRealtime / MossTTSLocalModel vs MossTTSDelayModel).
+
+__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE", "MOSS_TTS_LOCAL_PIPELINE"]

vllm_omni/model_executor/models/registry.py

@@ -294,6 +294,12 @@ _OMNI_MODELS = {
         "modeling_moss_tts_talker",
         "MossTTSRealtimeTalkerForGeneration",
     ),
+    # MossTTSLocalModel: MOSS-TTS-Local-Transformer-v1.5
+    "MossTTSLocalModel": (
+        "moss_tts",
+        "modeling_moss_tts_talker",
+        "MossTTSLocalTalkerForGeneration",
+    ),
     # Stage-1 codec decoder (shared by all 5 variants)
     "MossTTSCodecDecoder": (
         "moss_tts",

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -11,12 +11,16 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 
 def ar2dit(
     source_outputs: list[Any],
-    prompts: OmniTokensPrompt | TextPrompt | None = None,
+    prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None,
     _requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
     """Convert AR stage outputs to DiT stage inputs."""
     ar_outputs = source_outputs
 
+    # Normalize prompts to list
+    if not isinstance(prompts, list):
+        prompts = [prompts] if prompts is not None else [{}]
+
     dit_inputs: list[OmniTokensPrompt] = []
     for ar_output, prompt in zip(ar_outputs, prompts):
         addi_info = prompt["additional_information"]

vllm_omni/platforms/cuda/platform.py

@@ -213,6 +213,11 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         free, _ = torch.cuda.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.cuda.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)

vllm_omni/platforms/interface.py

@@ -161,6 +161,10 @@ class OmniPlatform(Platform):
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        raise NotImplementedError
+
     @classmethod
     def create_autocast_context(
         cls,

vllm_omni/platforms/musa/platform.py

@@ -146,17 +146,14 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
-        """Get the free memory on the MUSA device.
-
-        Args:
-            device: Optional device to query. If None, uses current device.
-
-        Returns:
-            Free memory in bytes.
-        """
         free, _ = torch.musa.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.musa.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.musa.get_device_name(device_id)