[Bugfix] Preserve session IDs when upgrading Omni engine requests

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
[CI/Build] Fix MOSS CPU tests on ROCm workers (#8447 )
2026-10-04 03:19:39 +08:00 · 2026-10-03 09:13:51 +08:00 · 2026-10-02 18:03:14 +00:00 · 2026-10-02 09:14:39 -07:00 · 2026-10-02 08:49:50 -07:00 · 2026-10-02 15:43:55 +00:00
239 changed files with 24273 additions and 2186 deletions
@@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts
  - vllm_omni/worker/gpu_model_runner.py
  - vllm_omni/model_executor/models/qwen3_tts/
  - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
  - vllm_omni/model_executor/models/common/snake_activation.py
  - vllm_omni/model_executor/models/common/alias_free_activation.py
  - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni
  - vllm_omni/entrypoints/openai/serving_chat.py
  - vllm_omni/model_executor/models/qwen3_omni/
  - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
  - vllm_omni/model_executor/models/common/snake_activation.py
  - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
  - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
  - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
  - vllm_omni/deploy/qwen3_omni_moe.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml
  - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml

 x-wan22: &wan22
@@ -124,6 +124,30 @@ steps:
          - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
        mirror_hardwares: h100_2

+      - label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
+          - tests/helpers/runtime.py
+          - tests/helpers/stage_config.py
+          - vllm_omni/core/sched/omni_ar_scheduler.py
+          - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
+          - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+          - vllm_omni/engine/orchestrator.py
+          - vllm_omni/engine/stage_engine_core_proc.py
+          - vllm_omni/engine/stage_pool.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        timeout_in_minutes: 40
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_2
+
      - label: "Omni · MiniCPM-o 4.5 Test"
        source_file_dependencies: omni_minicpmo_4_5_function
        timeout_in_minutes: 50
@@ -25,6 +25,16 @@ steps:
          - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model"
        mirror_hardwares: h100_1

+      - label: "Simple · MRV2 FA3 Graph Bounds Test"
+        timeout_in_minutes: 15
+        source_file_dependencies:
+          - tests/worker_v2/test_capture_query_bound_gpu.py
+          - tests/worker_v2/test_decode_query_bound_gpu.py
+          - vllm_omni/worker_v2/model_states/omni_model_state.py
+        commands:
+          - "pytest -sv tests/worker_v2/test_capture_query_bound_gpu.py tests/worker_v2/test_decode_query_bound_gpu.py -m 'core_model and cuda and H100' --run-level=core_model"
+        mirror_hardwares: h100_1
+
      - label: "Simple · Diffusion Test"
        timeout_in_minutes: 30
        commands:
@@ -9,15 +9,6 @@ Review like a maintainer: direct, selective, and focused on issues that CI does
 not prove. Prefer a few high-confidence findings over exhaustive commentary.
 Zero findings is a valid result.

-## Target-version evidence
-
-Resolve Python/dependency compatibility, pytest selectors and hardware needs
-from the frozen head's metadata, requirements, CI, and fixtures. Run levels such
-as `core_model` do not imply CPU-only execution. Read the target's accepted
-module contracts; draft or deferred architecture is not a merge requirement.
-Recheck only affected assumptions when reviewing another release. Historical
-coverage does not establish compatibility with a new snapshot.
-
 ## Quality contract

 Make every finding:
@@ -112,12 +103,6 @@ If docs and live code disagree, verify the code/tests and report the drift.
 | [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. |
 | [find-simplifications](../find-simplifications/SKILL.md) | Every review; run a diff-scoped subtraction and simplification pass after correctness blockers. |

-### Batch sessions and review helpers
-
-For a requested batch session, unanswered review replies, or inline-coordinate
-validation, read [batch-review.md](references/process/batch-review.md). The
-repository owns the helper scripts and their offline regression tests.
-
 ### Delivery and reviewer coordination

 | Reference | Read when |
@@ -125,15 +110,6 @@ repository owns the helper scripts and their offline regression tests.
 | [maintainer-style-study.md](references/delivery/maintainer-style-study.md) | Findings are ready for concise maintainer-style delivery. |
 | [review-requests.md](references/delivery/review-requests.md) | The user asks to identify, suggest, request, or ping code-owner reviewers. |

-## Optional review workers
-
-For independent defect discovery on a large diff, or when the user selects
-`auto`, `codex`, `claude`, or `none`, read
-[reviewer-backends.md](references/process/reviewer-backends.md). Pass immutable
-base/head SHAs and a bounded read-only scope. The parent owns domain routing,
-verification, deduplication, severity, and any separately authorized posting.
-Missing optional tooling falls back to direct review without blocking it.
-
 ## Workflow

 ### 1. Freeze and report the snapshot
@@ -20,15 +20,11 @@ were reverted or broken. Check:
  tolerances are deterministic or explicitly controlled;
 - normal, invalid, boundary, feature-off, failure/cancellation, and regression
  paths are covered where the diff changes them;
- read the target's pytest configuration, collection hooks, hardware helpers,
-  and CI together: run-level/domain/backend/SKU/card-count selectors must place
-  tests in the intended lane. `core_model` alone is not CPU-only; use explicit
-  hardware selection such as `core_model and cpu` only when supported by the
-  target. Zero selected or skipped accelerator tests are not validation.
+- run-level and domain markers place hardware/model tests in the intended CI
+  lane without silently skipping the contract.

-Map source symbols to tests with bounded `rg` searches and inspect the target's
-test-writing guide. Components may mirror source while model E2E and feature
-suites live elsewhere; imports, fixtures, and assertions establish coverage.
+Map source symbols to tests with bounded `rg` searches; do not assume the test
+tree mirrors production paths.

 ## Runtime check

@@ -1,50 +0,0 @@
-# Requested batch review and helper tools
-
-Use only for a requested batch session, review follow-up, or proposed inline
-comments. Commands below run from this skill directory, require `gh`, `jq`,
-Python, and Bash, and only read GitHub. Resolve the connected reviewer's login;
-do not infer it from examples or repository ownership.
-
-This is generic candidate discovery, not a daily eligibility policy. If the
-user's session specifies eligibility labels, a changed head since prior review,
-a daily cap, or another selection rule, build that eligible set from current
-metadata first. Do not let the legacy selector's own-PR/already-reviewed filters
-silently override it: select directly when those filters would discard an
-explicitly eligible re-review. User-selected PRs do not need to pass the helper.
-
-1. Check outstanding replies first when the session includes follow-up:
-   `bash scripts/check_replies.sh --reviewer <login> --days 14`.
-   Reopen each returned thread and inspect its current resolution before
-   deciding whether a response is needed. A reply can already be addressed by
-   code or another maintainer; do not send an automatic acknowledgement.
-2. Select candidates with
-   `bash scripts/select_prs.sh --reviewer <login> --days 7 --limit 5`.
-   This preserves the existing selector: open recent PRs, excluding own,
-   draft/WIP/"don't merge", and already-reviewed PRs, with low-review-count
-   candidates first. It does not filter ordinary docs/CI PRs or prove
-   mergeability, CI success, or complete search coverage.
-3. Freeze and review each selected PR using the normal execution contract.
-   Keep CI conclusions separate from conflicts/mergeability. Match depth to
-   changed behavior, not candidate rank; no quota requires manufacturing
-   findings. Search/review APIs can be limited or fail; inspect raw status when
-   coverage matters rather than treating a zero count as a verified clean state.
-4. Before separately authorized inline posting, validate the proposed payload:
-   `printf '%s' "$REVIEW_JSON" | bash scripts/verify_line_numbers.sh <PR>`.
-   Input is an object with a `comments` array containing `path`, `line`,
-   optional `side` (default `RIGHT`), and optional same-side `start_line`.
-   It checks actual diff hunks, including deletion/rename paths. It does not
-   verify the finding or submit it. Re-read base/head before and after its
-   live diff fetch; require equality with the reviewed snapshot and use that
-   head as `commit_id` for any authorized post. Otherwise refresh the review
-   and line mapping first.
-5. Preserve a short session record of pinned heads, findings, validation gaps,
-   and already-posted comment links so a repeat pass reviews only the delta.
-
-Run helper regression tests without project dependencies or GitHub access:
-
-```bash
-python tests/test_helpers.py
-```
-
-The helpers are bounded triage aids, not a scheduler. They do not authorize
-follow-up comments, review events, code changes, or recurring monitoring.
@@ -1,80 +0,0 @@
-# Reviewer Backends
-
-Use this reference when delegating defect discovery to Codex, Claude, or another reviewer. The parent `review-pr` workflow remains responsible for gates, domain routing, verification, prioritization, and any GitHub writes.
-
-## Backend Selection
-
-| Choice | Behavior |
-|--------|----------|
-| `auto` | Use one available reviewer backend; fall back to direct review |
-| `codex` | Use an available Codex review agent or skill |
-| `claude` | Use an available Claude review agent or CLI |
-| `none` | Review directly without delegation |
-
-Honor an explicit user choice. Do not block when the requested backend is unavailable; report the fallback locally and continue directly unless the user required that backend specifically.
-
-Use one backend for ordinary code changes. Use at most two reviewers for large or high-risk changes, either with non-overlapping path scopes or as independent passes. Do not run multiple reviewers merely to vote on the same conclusion.
-
-## Common Read-Only Handoff
-
-Give every backend the same concrete target:
-
-```yaml
-repository: /absolute/path/to/vllm-omni
-base_ref: origin/main
-base_sha: <resolved-base-sha>
-head_sha: <captured-pr-head-sha>
-merge_base_sha: <resolved-merge-base-sha>
-scope: complete-diff | path-partition
-paths: []
-focus: []
-constraints:
-  read_only: true
-  no_github_writes: true
-  no_delegation: true
-output:
-  finding: "[P0-P3] Imperative title — path:line"
-  exhaustive: true
-  include_assessment: true
-  include_test_gaps: true
-  include_residual_risks: true
-```
-
-Capture immutable SHAs before delegation. Ask the reviewer to:
-
-1. Read the applicable `AGENTS.md`.
-2. Inspect the complete assigned diff and enough surrounding code to understand it.
-3. Check relevant tests and call sites.
-4. Report only concrete regressions introduced by the change.
-5. Cite the smallest changed-line range that demonstrates each finding.
-6. Continue through the full assigned scope after finding the first issue.
-
-The worker must not modify files, create commits, push branches, post comments, submit a review event, or delegate again.
-
-## Backend Adapters
-
-### Codex
-
-Use an available Codex review agent or review skill; do not assume that a named system skill is installed. Give it the common handoff plus selected vLLM-Omni focus areas. Use it for a complete change or bounded path partition, not for a one-function context lookup.
-
-### Claude
-
-Pass the same handoff and output format to the available Claude review surface. Do not assume a particular command, model, or installation path. Explicitly include the read-only and no-GitHub-write constraints in the prompt.
-
-### Direct fallback
-
-If no backend is available, the parent performs the same checks. Keep the same evidence threshold and comment budget. Do not mention missing tooling in the GitHub review unless the user asks.
-
-## Aggregate Results
-
-The parent reviewer must:
-
-1. Re-read the PR head SHA. If it changed, discard stale line mappings and rerun affected scopes.
-2. Verify each finding against code, tests, and the captured diff.
-3. Drop speculative, pre-existing, intentional, or style-only findings.
-4. Deduplicate by root cause and affected call path rather than by line alone.
-5. Classify `blocking` independently from `P0`-`P3`; priority alone does not determine the verdict.
-6. Apply the vLLM-Omni comment budget after collecting the complete result.
-7. Preserve material omitted findings, test gaps, and residual risks in the local report.
-
-Only the parent may turn selected findings into GitHub comments, and only with user authorization.
@@ -1,87 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-# check_replies.sh — Find unanswered contributor replies to your review comments.
-#
-# Usage: ./check_replies.sh --reviewer lishunyang12 [--days 14]
-#
-# Scans PRs you've reviewed for threads where a contributor replied
-# but you haven't responded yet. Outputs actionable list.
-
-set -euo pipefail
-
-REPO="vllm-project/vllm-omni"
-REVIEWER=""
-DAYS=14
-
-while [[ $# -gt 0 ]]; do
-  case $1 in
-    --reviewer) REVIEWER="$2"; shift 2 ;;
-    --days)     DAYS="$2"; shift 2 ;;
-    *) echo "Unknown arg: $1" >&2; exit 1 ;;
-  esac
-done
-
-if [ -z "$REVIEWER" ]; then
-  echo "Usage: $0 --reviewer <github_login>" >&2
-  exit 1
-fi
-
-SINCE=$(date -d "-${DAYS} days" +%Y-%m-%d 2>/dev/null || date -v-${DAYS}d +%Y-%m-%d)
-
-echo "Scanning for unanswered replies to ${REVIEWER} (last ${DAYS} days)..." >&2
-
-# Find PRs the reviewer has commented on recently
-REVIEWED_PRS=$(gh api "search/issues?q=repo:${REPO}+reviewed-by:${REVIEWER}+is:pr+is:open+updated:>=${SINCE}&per_page=50" \
-  --jq '.items[].number' 2>/dev/null)
-
-UNANSWERED=0
-
-for PR_NUM in $REVIEWED_PRS; do
-  # Get all review comments on this PR
-  COMMENTS=$(gh api "repos/${REPO}/pulls/${PR_NUM}/comments" \
-    --jq '[.[] | {id, user: .user.login, body: .body[:100], in_reply_to_id, created_at, path, url: .html_url}]' 2>/dev/null)
-
-  # Find threads: reviewer's comments that got a reply, but reviewer didn't reply back
-  # 1. Get reviewer's comment IDs
-  REVIEWER_IDS=$(echo "$COMMENTS" | jq -r --arg r "$REVIEWER" '[.[] | select(.user == $r) | .id]')
-
-  # 2. Find replies to reviewer's comments (from others)
-  REPLIES=$(echo "$COMMENTS" | jq -c --arg r "$REVIEWER" --argjson rids "$REVIEWER_IDS" '
-    [.[] |
-      select(.user != $r) |
-      select(.in_reply_to_id != null) |
-      select([.in_reply_to_id] | inside($rids))
-    ]
-  ')
-
-  REPLY_COUNT=$(echo "$REPLIES" | jq 'length')
-  if [ "$REPLY_COUNT" -eq 0 ]; then
-    continue
-  fi
-
-  # 3. For each reply, check if reviewer replied after it
-  while IFS= read -r reply; do
-    THREAD_ID=$(echo "$reply" | jq -r '.in_reply_to_id')
-    REPLY_DATE=$(echo "$reply" | jq -r '.created_at')
-    REPLY_USER=$(echo "$reply" | jq -r '.user')
-    REPLY_URL=$(echo "$reply" | jq -r '.url')
-
-    # Check if reviewer replied after this reply
-    FOLLOWUP=$(echo "$COMMENTS" | jq --arg r "$REVIEWER" --arg date "$REPLY_DATE" --argjson tid "$THREAD_ID" '
-      [.[] | select(.user == $r) | select(.in_reply_to_id == $tid) | select(.created_at > $date)] | length
-    ')
-
-    if [ "$FOLLOWUP" -eq 0 ]; then
-      REPLY_BODY=$(echo "$reply" | jq -r '.body')
-      echo "PR #${PR_NUM} | @${REPLY_USER} replied | ${REPLY_URL}"
-      echo "  > ${REPLY_BODY}"
-      echo ""
-      UNANSWERED=$((UNANSWERED + 1))
-    fi
-  done < <(echo "$REPLIES" | jq -c '.[]')
-done
-
-if [ "$UNANSWERED" -eq 0 ]; then
-  echo "No unanswered replies found." >&2
-fi
@@ -1,73 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-# select_prs.sh — Pick reviewable PRs from vllm-omni for a review session.
-#
-# Usage: ./select_prs.sh [--days 7] [--limit 5] [--reviewer lishunyang12]
-#
-# Filters:
-#   - Excludes your own PRs (--reviewer)
-#   - Excludes WIP / Draft / "Don't merge" PRs
-#   - Excludes pure docs/CI-only merged PRs
-#   - Excludes PRs you already reviewed
-#   - Prioritizes zero-review PRs
-#
-# Output: JSON array of {number, title, author, createdAt, reviewCount}
-
-set -euo pipefail
-
-REPO="vllm-project/vllm-omni"
-DAYS=7
-LIMIT=5
-REVIEWER=""
-
-while [[ $# -gt 0 ]]; do
-  case $1 in
-    --days)    DAYS="$2"; shift 2 ;;
-    --limit)   LIMIT="$2"; shift 2 ;;
-    --reviewer) REVIEWER="$2"; shift 2 ;;
-    *) echo "Unknown arg: $1" >&2; exit 1 ;;
-  esac
-done
-
-SINCE=$(date -d "-${DAYS} days" +%Y-%m-%d 2>/dev/null || date -v-${DAYS}d +%Y-%m-%d)
-
-echo "Fetching open PRs from last ${DAYS} days..." >&2
-
-# Fetch all open PRs from the time range
-ALL_PRS=$(gh pr list --repo "$REPO" --state open --limit 200 \
-  --json number,title,author,createdAt,isDraft \
-  --search "created:>=${SINCE}")
-
-# Filter in a single jq pass
-FILTERED=$(echo "$ALL_PRS" | jq --arg reviewer "$REVIEWER" '
-  [.[] |
-    # Skip own PRs
-    select(.author.login != $reviewer) |
-    # Skip drafts
-    select(.isDraft == false) |
-    # Skip WIP / Dont merge
-    select(.title | test("\\[WIP\\]|\\[Draft\\]|Don.t[ ]+merge|\\[Don.t\\]"; "i") | not)
-  ]
-')
-
-# For each remaining PR, check review count and whether reviewer already reviewed
-echo "Checking review status for $(echo "$FILTERED" | jq length) PRs..." >&2
-
-echo "$FILTERED" | jq -c '.[]' | while read -r pr; do
-  NUM=$(echo "$pr" | jq -r '.number')
-
-  # Get review count
-  REVIEW_COUNT=$(gh api "repos/${REPO}/pulls/${NUM}/reviews" --jq 'length' 2>/dev/null || echo "0")
-
-  # Check if reviewer already reviewed
-  if [ -n "$REVIEWER" ]; then
-    ALREADY=$(gh api "repos/${REPO}/pulls/${NUM}/reviews" \
-      --jq "[.[] | select(.user.login==\"${REVIEWER}\")] | length" 2>/dev/null || echo "0")
-    if [ "$ALREADY" -gt 0 ]; then
-      continue
-    fi
-  fi
-
-  echo "$pr" | jq --argjson rc "$REVIEW_COUNT" '. + {reviewCount: $rc}'
-done | jq -s "sort_by(.reviewCount) | .[:${LIMIT}]"
@@ -1,129 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Check review comment coordinates against the actual lines of a PR diff."""
-
-import argparse
-import ast
-import json
-import subprocess
-import sys
-
-
-def parse_hunk(text):
-    """Parse a unified-diff hunk header without a runtime dependency."""
-    parts = text.split(" ", 4)
-    if len(parts) < 4 or parts[0] != "@@" or not parts[3].startswith("@@"):
-        return None
-    bounds = []
-    for value, sign in zip(parts[1:3], ("-", "+")):
-        if not value.startswith(sign):
-            return None
-        numbers = value[1:].split(",")
-        if len(numbers) not in (1, 2) or not all(number.isdecimal() for number in numbers):
-            return None
-        bounds.extend((numbers[0], numbers[1] if len(numbers) == 2 else None))
-    return tuple(bounds)
-
-
-def diff_path(value):
-    """Decode a git path, including C-quoted UTF-8 bytes."""
-    if value.startswith('"'):
-        value = ast.literal_eval(value).encode("latin1").decode("utf-8")
-    return None if value == "/dev/null" else value[2:]
-
-
-def diff_lines(diff):
-    """Map each file/side to commentable lines, preserving rename/delete paths."""
-    lines = {}
-    old_path = new_path = None
-    old_remaining = new_remaining = 0
-    for text in diff.splitlines():
-        if text.startswith("diff --git "):
-            old_path = new_path = None
-            old_remaining = new_remaining = 0
-        elif hunk := parse_hunk(text):
-            old_line, old_count, new_line, new_count = hunk
-            old_line, new_line = int(old_line), int(new_line)
-            old_remaining = int(old_count) if old_count is not None else 1
-            new_remaining = int(new_count) if new_count is not None else 1
-        elif old_remaining or new_remaining:
-            # A hunk's content can itself start with --- or +++.
-            path = new_path or old_path
-            if text.startswith((" ", "-")):
-                lines.setdefault((path, "LEFT"), set()).add(old_line)
-                old_line += 1
-                old_remaining -= 1
-            if text.startswith((" ", "+")):
-                lines.setdefault((path, "RIGHT"), set()).add(new_line)
-                new_line += 1
-                new_remaining -= 1
-        elif text.startswith("--- "):
-            old_path = diff_path(text[4:])
-        elif text.startswith("+++ "):
-            new_path = diff_path(text[4:])
-    return lines
-
-
-def validate_comments(comments, lines):
-    """Validate single-line and same-side multiline review comments."""
-    valid = True
-    for comment in comments:
-        path = comment["path"]
-        side = comment.get("side", "RIGHT")
-        line = comment["line"]
-        start = comment.get("start_line", line)
-        start_side = comment.get("start_side", side)
-        available = lines.get((path, side), set())
-        if (
-            type(line) is not int
-            or type(start) is not int
-            or start < 1
-            or line < start
-            or start_side != side
-            or line - start + 1 > len(available)
-            or not all(number in available for number in range(start, line + 1))
-        ):
-            print(
-                f"ERROR: {path}:{start}-{line} ({side}) is outside diff lines",
-                file=sys.stderr,
-            )
-            valid = False
-        else:
-            print(f"OK: {path}:{line} ({side})", file=sys.stderr)
-    return valid
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("pr_number", type=int)
-    args = parser.parse_args()
-    try:
-        review = json.load(sys.stdin)
-        comments = review["comments"]
-        if not isinstance(comments, list):
-            raise TypeError("comments must be an array")
-        result = subprocess.run(
-            [
-                "gh",
-                "pr",
-                "diff",
-                str(args.pr_number),
-                "--repo",
-                "vllm-project/vllm-omni",
-            ],
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-        return 0 if validate_comments(comments, diff_lines(result.stdout)) else 1
-    except subprocess.CalledProcessError as error:
-        print(error.stderr, file=sys.stderr, end="")
-        return 1
-    except (OSError, ValueError, KeyError, TypeError, SyntaxError) as error:
-        print(f"ERROR: {error}", file=sys.stderr)
-        return 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())
@@ -1,8 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-# Usage: echo '<review_json>' | ./verify_line_numbers.sh <pr_number>
-# Requires Python 3 and gh. Exits nonzero for invalid comments or fetch failures.
-set -euo pipefail
-
-exec python3 "$(dirname "${BASH_SOURCE[0]}")/verify_line_numbers.py" "$@"
@@ -1,227 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Exercise repository review helpers without network access."""
-
-import json
-import os
-import subprocess
-import tempfile
-import unittest
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parents[1]
-REVIEW = ROOT / "scripts"
-
-
-class HelperTests(unittest.TestCase):
-    def setUp(self):
-        self.temporary = tempfile.TemporaryDirectory()
-        self.addCleanup(self.temporary.cleanup)
-        self.root = Path(self.temporary.name)
-        self.bin = self.root / "bin"
-        self.bin.mkdir()
-        gh = self.bin / "gh"
-        gh.write_text("""#!/usr/bin/env python3
-import json, os, sys
-from pathlib import Path
-args = sys.argv[1:]
-root = Path(os.environ['FIXTURES'])
-if os.environ.get('FAIL_GH'):
-    print('GitHub unavailable', file=sys.stderr)
-    raise SystemExit(1)
-if args[:2] == ['pr', 'list']:
-    print((root / 'prs.json').read_text())
-elif args[:2] == ['pr', 'diff']:
-    print((root / 'diff').read_text(), end='')
-elif args[0] == 'api' and args[1].startswith('search/issues?'):
-    print('1')
-elif args[0] == 'api' and args[1].endswith('/comments'):
-    print((root / 'comments.json').read_text())
-elif args[0] == 'api' and args[1].endswith('/reviews'):
-    print('0')
-else:
-    raise SystemExit('Unexpected gh arguments: ' + repr(args))
-""")
-        gh.chmod(0o755)
-        self.env = dict(
-            os.environ,
-            PATH=f"{self.bin}{os.pathsep}{os.environ['PATH']}",
-            FIXTURES=str(self.root),
-            PYTHONDONTWRITEBYTECODE="1",
-        )
-
-    def run_script(self, script, *args, data=None):
-        return subprocess.run(
-            ["bash", str(script), *args],
-            input=data,
-            env=self.env,
-            text=True,
-            capture_output=True,
-            check=False,
-        )
-
-    def test_select_filters_drafts_wip_and_own_prs(self):
-        prs = [
-            {"number": i, "title": title, "author": {"login": author}, "isDraft": draft}
-            for i, (title, author, draft) in enumerate(
-                [
-                    ("Fix bug", "author", False),
-                    ("[WIP] Feature", "author", False),
-                    ("[Draft] Feature", "author", False),
-                    ("Don't merge yet", "author", False),
-                    ("Feature", "reviewer", False),
-                    ("Feature", "author", True),
-                ],
-                1,
-            )
-        ]
-        (self.root / "prs.json").write_text(json.dumps(prs))
-        result = self.run_script(REVIEW / "select_prs.sh", "--reviewer", "reviewer")
-        self.assertEqual(result.returncode, 0, result.stderr)
-        self.assertEqual([pr["number"] for pr in json.loads(result.stdout)], [1])
-
-    def verify(self, comments, diff=None):
-        (self.root / "diff").write_text(
-            diff
-            or (
-                "diff --git a/src/example.py b/src/example.py\n"
-                "--- a/src/example.py\n+++ b/src/example.py\n"
-                "@@ -10,2 +20,2 @@\n-old\n+new\n context\n"
-            )
-        )
-        return self.run_script(
-            REVIEW / "verify_line_numbers.sh",
-            "1",
-            data=json.dumps({"comments": comments}),
-        )
-
-    def test_verify_valid_nested_path_and_both_sides(self):
-        result = self.verify(
-            [
-                {"path": "src/example.py", "line": line, "side": side, "body": "Check"}
-                for side, line in [("RIGHT", 20), ("RIGHT", 21), ("LEFT", 10)]
-            ]
-        )
-        self.assertEqual(result.returncode, 0, result.stderr)
-        self.assertEqual(result.stderr.count("OK:"), 3)
-
-    def test_verify_rejects_outside_hunk_or_missing_file(self):
-        for path, line in [("src/example.py", 999), ("missing.py", 20)]:
-            with self.subTest(path=path):
-                result = self.verify([{"path": path, "line": line, "body": "Check"}])
-                self.assertNotEqual(result.returncode, 0)
-
-    def test_verify_deleted_file_on_left(self):
-        diff = "diff --git a/gone.py b/gone.py\n--- a/gone.py\n+++ /dev/null\n@@ -1 +0,0 @@\n-old\n"
-        result = self.verify([{"path": "gone.py", "line": 1, "side": "LEFT", "body": "Check"}], diff)
-        self.assertEqual(result.returncode, 0, result.stderr)
-        result = self.verify([{"path": "gone.py", "line": 1, "side": "RIGHT", "body": "Check"}], diff)
-        self.assertNotEqual(result.returncode, 0)
-
-    def test_verify_new_file_with_explicit_zero_old_count(self):
-        diff = "diff --git a/new.py b/new.py\n--- /dev/null\n+++ b/new.py\n@@ -0,0 +1,2 @@\n+one\n+two\n"
-        result = self.verify([{"path": "new.py", "line": 2, "side": "RIGHT", "body": "Check"}], diff)
-        self.assertEqual(result.returncode, 0, result.stderr)
-        result = self.verify([{"path": "new.py", "line": 1, "side": "LEFT", "body": "Check"}], diff)
-        self.assertNotEqual(result.returncode, 0)
-
-    def test_verify_checks_multiline_comment_start(self):
-        result = self.verify(
-            [
-                {
-                    "path": "src/example.py",
-                    "start_line": 1,
-                    "start_side": "RIGHT",
-                    "line": 20,
-                    "side": "RIGHT",
-                    "body": "Check",
-                }
-            ]
-        )
-        self.assertNotEqual(result.returncode, 0)
-
-    def test_verify_multiple_files_hunks_and_renamed_path(self):
-        diff = (
-            "diff --git a/old.py b/new.py\n--- a/old.py\n+++ b/new.py\n"
-            "@@ -1 +1 @@\n-old\n+new\n"
-            "@@ -10 +10 @@\n-old\n+new\n"
-            "diff --git a/second.py b/second.py\n--- a/second.py\n+++ b/second.py\n"
-            "@@ -1 +1 @@\n-old\n+new\n"
-        )
-        result = self.verify(
-            [
-                {"path": path, "line": line, "body": "Check"}
-                for path, line in [("new.py", 1), ("new.py", 10), ("second.py", 1)]
-            ],
-            diff,
-        )
-        self.assertEqual(result.returncode, 0, result.stderr)
-        result = self.verify([{"path": "second.py", "line": 10, "body": "Check"}], diff)
-        self.assertNotEqual(result.returncode, 0)
-
-    def test_verify_quoted_path_and_header_like_content(self):
-        diff = (
-            'diff --git "a/a\\tb.py" "b/a\\tb.py"\n'
-            '--- "a/a\\tb.py"\n+++ "b/a\\tb.py"\n'
-            "@@ -1 +1 @@\n--- content\n+++ content\n"
-        )
-        result = self.verify([{"path": "a\tb.py", "line": 1, "body": "Check"}], diff)
-        self.assertEqual(result.returncode, 0, result.stderr)
-
-    def test_verify_rejects_malformed_input_and_fetch_failure(self):
-        for data in ("not json", '{"comments": null}'):
-            result = self.run_script(REVIEW / "verify_line_numbers.sh", "1", data=data)
-            self.assertNotEqual(result.returncode, 0)
-        self.env["FAIL_GH"] = "1"
-        result = self.verify([{"path": "src/example.py", "line": 20, "body": "Check"}])
-        self.assertNotEqual(result.returncode, 0)
-        self.assertIn("GitHub unavailable", result.stderr)
-
-    def replies(self, answered):
-        comments = [
-            {
-                "id": 100,
-                "user": "reviewer",
-                "body": "Please fix",
-                "in_reply_to_id": None,
-                "created_at": "2026-09-01T00:00:00Z",
-                "url": "root",
-            },
-            {
-                "id": 101,
-                "user": "author",
-                "body": "Fixed",
-                "in_reply_to_id": 100,
-                "created_at": "2026-09-02T00:00:00Z",
-                "url": "reply",
-            },
-        ]
-        if answered:
-            comments.append(
-                {
-                    "id": 102,
-                    "user": "reviewer",
-                    "body": "Thanks",
-                    "in_reply_to_id": 100,
-                    "created_at": "2026-09-03T00:00:00Z",
-                    "url": "answer",
-                }
-            )
-        (self.root / "comments.json").write_text(json.dumps(comments))
-        return self.run_script(REVIEW / "check_replies.sh", "--reviewer", "reviewer")
-
-    def test_replies_ignores_answered_thread(self):
-        result = self.replies(answered=True)
-        self.assertEqual(result.returncode, 0, result.stderr)
-        self.assertEqual(result.stdout, "")
-        self.assertIn("No unanswered replies found", result.stderr)
-
-    def test_replies_reports_unanswered_without_false_clean_summary(self):
-        result = self.replies(answered=False)
-        self.assertEqual(result.returncode, 0, result.stderr)
-        self.assertIn("@author replied", result.stdout)
-        self.assertNotIn("No unanswered replies found", result.stderr)
-
-
-if __name__ == "__main__":
-    unittest.main()
@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.

 ## Model-specific variables

-The inventory includes 66 variables read by a single model or pipeline family. They are
+The inventory includes 93 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.

@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the

 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 38 | Move a stable setting into typed stage or model configuration. |
+| Promote | 64 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
-| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
 | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |

 The disposition is a migration target, not a statement that the existing
@@ -0,0 +1,22 @@
+# Pre-sharded HSDP loading
+
+Use `--use-hsdp --hsdp-shard-size 8 --hsdp-weight-load-strategy pre_sharded`
+to shard transformer parameters before loading checkpoint values. Set the shard
+size to the available GPU count (or configure replication explicitly).
+The default strategy, `full`, retains the existing load-then-shard behavior.
+The same `hsdp_weight_load_strategy` option is available through Python and stage
+engine configuration.
+
+The pre-sharded strategy reads rank-local slices from ordinary Hugging Face
+safetensors checkpoints into FSDP-owned storage. No converted checkpoint is
+required. Cosmos 3 and Cosmos 3 Edge transformer blocks release their initialized
+parameter storage to meta as they are constructed; nonpersistent buffers retain
+their values.
+Encoders and VAEs follow their existing loading paths.
+
+This strategy requires the default diffusion pipeline loader, complete dedicated
+transformer checkpoint sources, and supported checkpoint key mappings. Tensor
+layouts must match the runtime layout; tensor transforms are unsupported.
+Checkpoint tensor dtypes must match runtime tensor dtypes at preflight; dtype conversion is unsupported.
+Quantization and LoRA are unsupported. Every worker needs checkpoint access.
+`num_weight_load_threads` controls the safetensors reader's threads per rank.
@@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles:
 | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS |
 | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
+| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU |

 ```bash
 # MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
@@ -350,12 +352,73 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local
 (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
 its existing `talker_mtp` entry point for V1.

+### Selecting one or two Qwen3-TTS stages
+
+Choose the deployment profile explicitly with `--deploy-config`. The default
+`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to
+Code2Wav, which produces audio. Its stage devices can be placed on the same
+GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for
+one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly,
+without a Code2Wav engine or stage connector.
+
+```bash
+# Two stages on one GPU; the two-stage pipeline remains independently selectable.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+
+# One stage on one GPU; no inter-stage transport or Code2Wav scheduling.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml
+```
+
+The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1,
+an in-process worker and disabled prefix caching. Select the full profile;
+adding `talker_stream_decode: true` to a two-stage Talker is rejected because
+its latent output belongs to the Code2Wav input contract. Single-stage options
+live in `stages[].additional_config`; two-stage transport and first-frame
+options live in `connectors.*.extra`.
+
+The single-stage profile defaults `talker_stream_first_audio` to `false` for
+throughput. To opt into a separate first-frame delivery queue, save this
+overlay beside the bundled deployment files and pass its path to
+`--deploy-config`:
+
+```yaml
+base_config: qwen3_tts_fused_single_gpu.yaml
+stages:
+  - stage_id: 0
+    additional_config:
+      talker_stream_decode: true
+      talker_stream_first_audio: true
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+```
+
+This queue uses PCM from the single-stage decoder. The two-stage
+`talker_first_audio` option below uses its own first-frame decoder and does
+not require single-stage execution. The two options cannot be enabled
+together. Single-stage prefix-cache reset with
+`reset_running_requests=True` returns `False` while requests are active;
+wait for completion or abort them before retrying.
+
+Reference-cloning context is selected independently by each profile. The
+single-stage profile explicitly uses 72 reference frames, matching the
+`qwen3_tts.yaml` context boundary. A shorter
+`additional_config.ref_code_context_frames` (for example 25) can change
+timbre continuity even when WER is similar;
+the throughput measurements for CustomVoice do not qualify Base voice
+cloning quality. Reference layout, grouping and priming are model hooks;
+the shared worker does not select Qwen's context length or codebook layout.
+
 ### Qwen3-TTS first-frame delivery and rollback

 The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`,
 `qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles
 enable the `talker_first_audio` connector option, as does the experimental
-single-GPU profile. This changes the default CUDA MRV2 streaming path:
+two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile.
+This changes the default CUDA MRV2 streaming path:
 residual prediction runs eagerly after the Talker sample, the Talker loads
 an additional first-frame decoder with its weights and CUDA graphs, and the
 orchestrator orders its audio before subsequent Code2Wav chunks.
@@ -451,3 +514,50 @@ for that comparison; this switch only controls automatic MPS management.

 MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
 the intended GPU, and warm the complete pipeline before measuring performance.
+
+## MOSS-TTS Local 1.5 with Model Runner V2
+
+`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
+the shared runtime introduced for Qwen3-TTS:
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This profile inherits the batching, codec graph buckets and 1-frame/15-frame
+chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
+Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
+capabilities while retaining its V1 `talker_mtp` implementation, sampling
+defaults and explicit request-seed handling. Codec chunks use the native data
+plane; the internal Talker retains its final-only orchestrator output policy.
+
+On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
+The codec retains its own CUDA Graph capture/replay, including batch buckets
+through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
+of the stateful decoder to avoid lengthy compilation at startup. This is
+independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
+It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
+the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
+the codec's capture buckets when compilation is disabled.
+
+For sustained C128 serving on a large-memory CUDA GPU, select the
+separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
+stream slots per stage, a 32 GiB Talker KV budget, and the codec's
+`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
+The bounded KV budget leaves room for codec state and graphs; it does not
+guarantee that 128 maximum-length prompts fit simultaneously. See the
+[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
+for activation, backend comparisons, memory requirements and benchmark commands.
+
+Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
+NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
+MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
+set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
+`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
+
+Event-driven orchestration remains independently selectable with
+`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
+when comparing these modes. Model-runner selection does not change the
+orchestration default or enable experimental reference encoding, chunk ramps,
+generation-output draining or MPS.
@@ -4,6 +4,9 @@ This document defines the selection and extension contract for diffusion
 attention backends. User-facing backend choices, installation, and tuning are
 in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).

+The [execution contract PoC](attention_execution_contract_poc.md) describes the
+experimental post-selection capability and lifecycle contract from RFC #7226.
+
 ## Scope

 The contract applies to DiT and other diffusion attention layers. It is
@@ -0,0 +1,133 @@
+# Diffusion attention execution contract PoC
+
+This prototype accompanies [RFC #7226](https://github.com/vllm-project/vllm-omni/issues/7226)
+and extends the [attention selection design](attention_backend_selection.md).
+It demonstrates path-specific capabilities with dense BF16 FA4, NPU/ROCm
+routing examples, and tensor-state lifetime with a test-only attention module.
+
+## Execution contract
+
+`ExecutionContext` describes the requested execution path. `ExecutionPathResult`
+reports its identity, support status and reason, and compilation mode.
+`AttentionBackend.resolve_capabilities()` provides conservative pre-construction
+results. After initialization, `Attention.resolve_execution_path()` supplies the
+active parallel, paged-KV, and HSDP context; the backend combines it with the
+selected kernel and normalized metadata.
+
+Resolution runs outside compiled execution. FA4 reports `SUPPORTED` and
+`CUSTOM_OP` for dense, noncausal BF16 without parallel or HSDP boundaries when
+its kernel accepts the head dimensions. Dimension validation delegates to FA4's
+architecture-specific rules through `backends/utils/fa.py`. Kernel rejections
+report `UNSUPPORTED` with an actionable reason. Missing private validators and
+other unmigrated paths report `UNMIGRATED` with advisory `EAGER_ONLY` defaults.
+`requested_support()` checks a fullgraph request; it does not enforce selection
+or change existing execution.
+
+FA4 execution uses an opaque custom op. Its fake output preserves Q's batch,
+sequence, and head count and uses V's head dimension. Output is contiguous even
+when Q is noncontiguous. Metadata normalization is
+shared with dispatch. Producers must update published mask semantics when masks
+change; unpublished masks remain runtime-dependent to avoid synchronization.
+Callers resolve again when execution metadata changes.
+
+## NPU and ROCm worked examples
+
+The same `FLASH_ATTN` backend resolves MindIE paths on NPU and AITER paths on
+ROCm. Both examples report `UNMIGRATED`: their route is inspected, while hardware
+correctness and compilation remain unvalidated. `EAGER_ONLY` is an advisory
+default, not a measured compiler limitation. Existing execution is unchanged.
+
+| Input | NPU / MindIE | ROCm / AITER |
+| --- | --- | --- |
+| Dense | `npu_dense` | `rocm_dense` |
+| Padding mask | `npu_masked`: expand the key mask to `[B, 1, Q, K]` | `rocm_masked_varlen`: unpad, call varlen attention, then restore rows |
+| Packed inputs | Opt-in `[real, pad]` contract; choose varlen or prefix-KV slicing from the environment | Forward complete packed-document boundaries to varlen attention |
+| Incomplete packing | Rebuild a mask when possible; otherwise reject before the kernel | Shared CUDA-like normalization rejects incomplete metadata |
+| Unpublished mask semantics | Route to the masked wrapper | Keep resolution runtime-dependent without reading mask values |
+
+NPU resolution reuses `_resolve_packed_seq_npu`; ROCm reuses the metadata
+normalization used by dispatch. AITER identity comes from initialization rather
+than caller-provided context. Quantized, paged, piecewise, and parallel paths
+remain outside these worked examples.
+
+The new L1 tests use CPU tensors and replace only vendor kernel calls. They
+check path identity, forwarded masks/boundaries, fallback rejection, and
+conservative fullgraph-request handling. They do not test vendor numerics or
+compiler support. Existing L1 CI already collects this test file.
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn.py \
+  -k 'npu_contract or rocm_contract' -m 'core_model and cpu' -q
+```
+
+### Hardware validation required before migration
+
+On an Ascend/MindIE or ROCm/AITER runner, record device, driver, PyTorch, vendor
+library, and compiler versions, then validate the corresponding rows above:
+
+1. Compare eager outputs with FP32 SDPA for dense and padding-mask inputs,
+   including unequal Q/K lengths. Compare valid query rows; do not assume both
+   platforms define padded-query outputs identically.
+2. Compare packed outputs with separate per-document references. On NPU, cover
+   both varlen and laser prefix-slicing modes, plus mask reconstruction and the
+   missing-mask error. A padding fallback does not establish multi-document
+   isolation. On ROCm, cover multiple real documents.
+3. Check output shape, dtype, device, finite values, and unchanged inputs.
+   Start BF16 comparisons at `atol=rtol=1e-2` and have the platform maintainer
+   confirm the tolerance against its reference tests.
+4. Run `torch.compile(fullgraph=True, dynamic=True)` with the platform's supported
+   compiler over multiple sequence lengths. Record graph breaks and recompiles;
+   choose `TRACEABLE`, `CUSTOM_OP`, or verified `EAGER_ONLY` from that evidence.
+   A custom-op path additionally needs schema/fake checks and repeated replay.
+
+Promote only the paths validated on that platform. These hardware checks have
+not been run locally; ROCm/NPU fullgraph support is not claimed.
+
+## State lifetime example
+
+`test_attention_state_lifetime.py` prepares an owned query-scale tensor eagerly
+and passes it explicitly to an opaque attention op. The compiled module retains
+its state after the caller drops its reference and releases ownership when the
+compiled callable is deleted. Equivalent instances reuse a graph with their own
+state values.
+
+This example tests ownership and release, dynamic replay across instances,
+custom-op schema/fake behavior, and Inductor execution. It adds no production
+backend or state registry. Planning keys, caching, fallback-policy declarations,
+opaque handles, CUDA graph capture, and TRTLLM integration are deferred.
+
+## Validation
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn_compile.py \
+  tests/diffusion/attention/test_flash_attn.py \
+  tests/diffusion/attention/test_attention_capabilities.py \
+  tests/diffusion/attention/test_attention_state_lifetime.py -q -rs
+```
+
+The focused suite covers capability decisions and metadata changes, existing
+padding/mask regressions, tensor-state lifetime, and real FA4 fullgraph execution
+against FP32 SDPA and eager FA4. Representative Q/K and V dimensions are
+(32, 32), (64, 64), (80, 48), (192, 128), and (256, 256), with batches 1 and 2
+and sequence lengths up to 1,024. Invalid dimensions are checked against actual
+kernel errors. These samples exercise the contract rather than define support.
+BF16 comparisons use `atol=rtol=1e-2`; `torch.library.opcheck` checks schema and
+fake-output correctness.
+
+Single-graph reuse is asserted in the tensor-state test. The FA4 numerical tests
+retain compiler caches across shapes within each case but do not require one
+graph across different batch sizes, Q/K length equality, or head dimensions.
+
+Tested environment: GB300, PyTorch `2.13.0+cu130`, FA4 `4.0.0b18`, CUTLASS DSL
+`4.6.2` with CUDA 13 libraries, Quack `0.6.4`, and TVM FFI `0.1.11`.
+Real-kernel tests require CUDA and CuTe FA4; their availability is checked at runtime.
+
+Before the rebase, the focused suite passed all 97 tests without skips, including
+11 CPU NPU/ROCm routing tests. After rebasing onto upstream `e3be42e05`, the
+expanded suite reports 99 passed and 11 failed in the local vLLM 0.28 environment.
+All 11 failures are import errors: upstream now requires `compute_layout_strides`
+from vLLM, which this environment lacks. Both noncontiguous FA4 regression cases
+and real-kernel schema/fake checks pass. The full suite must be rerun with the
+vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
+unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
+warnings remain.
@@ -68,6 +68,7 @@ guides list selectable backends, platform defaults, installation, and tuning.
 The design contracts separate selection mechanics from backend algorithms:

 - [Attention Backend Selection](feature/attention_backend_selection.md)
+- [Attention Execution Contract PoC](feature/attention_execution_contract_poc.md)
 - [Skip-Softmax](feature/skip_softmax.md)

 #### CPU offloading
@@ -116,4 +117,5 @@ The pre-#5137 pages are preserved in the
 [legacy module archive](module/archive/README.md) for historical reference and
 are not active design contracts.

+- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md)
 - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md)
@@ -0,0 +1,22 @@
+# Qwen3-Omni MRv2 performance
+
+This model adaptation builds on #8184's existing stage runtime and output
+contracts. It adds Talker CPU metadata parsing, sampled-token embedding
+handoff, an optional first-frame decoder, length-grouped Code2Wav graph
+buckets, and opt-in incremental short-KV code prediction. The H200 profile
+also enables conservative fused sampling and bounded cuDNN algorithm search.
+The shared predictor keeps the upstream re-prefill helper and its default
+sampling behavior; short KV is explicitly selected by the deployment profile.
+
+Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200
+profile. Preserve a control with `code_predictor_fused_sampling=false` and
+`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement.
+cuDNN search restores process settings on exit and has startup cost. Kernel
+arithmetic is not bitwise-equivalent; the first-frame decoder uses additional
+model memory. Full-duplex performance is not established by turn-mode tests.
+
+First-frame delivery requires `talker_first_audio: true` in connector extra,
+a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled.
+Other deployments retain the regular Code2Wav path.
+`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks;
+the base V1 profile retains its original blocks and indexed text cache.
@@ -16,6 +16,10 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
    --trust-remote-code
 ```

+To enable [incremental prefill](#incremental-prefill), set
+`enable_prefix_caching: true` on stage 0 (the thinker) in the deploy config. Otherwise
+the server falls back to rebuilding the prompt per query.
+
 ### Run the Example Client

 ```bash
@@ -54,7 +58,7 @@ WebSocket /v1/video/chat/stream

 When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.

-Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
+Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. With incremental prefill active, every buffered frame is included instead of stride-sampling, with the same decode-failure exclusion. Excluded frames are not replaced with other buffered frames. The incremental path waits for the whole window to finish decoding. The legacy path waits only for its sampled positions, so an unselected frame still decoding does not delay the query.

 **Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.

@@ -64,10 +68,10 @@ Queries stride-sample the buffered frames, keeping the last frame, then exclude
 | ------- | ------ | --------- | ------------- |
 | `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
 | `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
-| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
-| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
+| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. Legacy path only: with incremental prefill active, every buffered frame is submitted and frame density is the client's responsibility (clients push discrete frames at their own rate). |
+| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Below the cap both paths keep every frame. Legacy evicts the oldest frame once full. Incremental prefill drops the oldest half at once, then appends until full again. |
 | `system_prompt` | string or null | null | Optional custom system prompt. |
-| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
+| `use_audio_in_video` | bool | `true` | Incremental prefill: forwarded on warmup and every query when this is true, so multimodal hashes stay aligned. Legacy: forwarded only when this query has input audio. |
 | `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. |
 | `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
 | `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
@@ -107,9 +111,42 @@ The server accepts these legacy field names and rewrites them before validation.

 EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.

+## Incremental Prefill
+
+Server-side only (stage-0 `enable_prefix_caching: true` and `modalities: ["text"]`).
+There is no `session.config` flag. Audio-output sessions stay on the legacy path.
+
+Input is unchanged: `video.frame` and `audio.chunk` are separate buffers. Warmup
+prefills `history + frames` as they arrive (`max_tokens=1`). When
+`use_audio_in_video` is true, warmup and query both pass that processor kwarg.
+`video.query` is the same prompt plus optional input audio and the question text,
+so the warmed vision prefix cache is reused and the query pays for the audio/text
+suffix. All buffered frames are
+submitted in arrival order (`num_frames` subsample is legacy-only). Once the buffer
+hits `max_frames`, the oldest half is dropped in one step and a warning is logged;
+until then the buffer matches the legacy path. After a turn,
+history compresses to the last two text messages and the next warmup follows the
+new prefix.
+
+### Legacy path (incremental prefill inactive)
+
+Stage-0 prefix caching is off, or `modalities` includes `"audio"`: no warmup,
+frames are re-sampled up to `num_frames` each query, and the query waits only
+for those sampled frames to decode. Prompt shape is the same
+(`image_pil` frames, optional trailing `input_audio`, text). Both paths keep the
+frame buffer across queries (`max_frames`) and compress history the same way.
+
 ## Known Limitations

- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
+- Past `max_frames`, incremental prefill keeps a suffix of the newest frames, not a resample of the whole stream. Rebuilding that longer window is left to a follow-up.
+- Incremental prefill applies only to text-only sessions with stage-0 prefix caching enabled
+  on `WS /v1/video/chat/stream`. It does not accelerate `/v1/realtime` camera input, which
+  uses duplex conversation items rather than `video.frame` / `video.query`. Audio-output
+  sessions on this endpoint rebuild the prompt each query.
+- Stage-0 prefix caching is safe for audio-output requests: on a hit, the omni prefix cache
+  restores the thinker hidden states the talker needs. A preempted request that resumes on a
+  hit does not get them back, so size the deployment to avoid preemption; see
+  [Prefix Caching](../design/feature/prefix_caching.md).
 - Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
 - If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
 - The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.
@@ -28,7 +28,9 @@ import torch
 from vllm_omni import Omni
 from vllm_omni.model_executor.models.yue2.yue2 import (
    ABC_SAMPLING,
+    CONTEXT,
    KEY_MAX_AUDIO_FRAMES,
+    KEY_MAX_HOLD_STEPS,
    KEY_MIN_TOKENS,
    KEY_PENALTY_WINDOW,
    KEY_PHASE,
@@ -41,10 +43,12 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
    KEY_TOP_P,
    SEMANTIC_SAMPLING,
    STOP_TOKEN_IDS,
+    SYNTHESIS_HOLD_STEPS,
 )
 from vllm_omni.tokenizers.yue2_prompt import (
    abc_ids_from_generated,
    abc_prefix_ids,
+    semantic_frames,
    semantic_prefix_ids,
 )
 from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
@@ -71,7 +75,13 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
        # conditioning prefix from the scheduled tokens alone.
        KEY_PREFIX_IDS: list(prompt_ids),
    }
-    params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
+    if phase == "abc":
+        params.max_tokens = preset["max_tokens"]
+    else:
+        # The song is synthesized on a side stream while the request is held;
+        # leave max_tokens headroom for the hold tokens.
+        params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - len(prompt_ids))
+        params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1)
    params.stop_token_ids = list(STOP_TOKEN_IDS)
    params.detokenize = False
    return params
@@ -144,6 +154,9 @@ def main() -> None:
        payload = {
            "prompt_token_ids": prompt_ids,
            "generated_token_ids": generated_ids,
+            # Codec frames only: without the end token and the hold tokens
+            # the request emits while its song is synthesized.
+            "semantic_frames": semantic_frames(generated_ids),
            "seed": args.seed,
            "cot": args.cot,
            "max_frames": args.max_frames,
@@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
 audio-reference generation. Reference transcripts must match the reference
 audio; they are not style instructions.

+For the optional CUDA MRV2 runner, see the
+[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
+The default Local deployment continues to use V1.
+
 ## References

 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 - No `ref_audio` required or accepted for MOSS-SoundEffect.
 - Input field maps to the `ambient_sound` parameter in the upstream processor.
 - Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
+
+## Local 1.5 MRV2 and slot attention
+
+`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
+an explicit deploy profile. The default Local profile continues to use V1.
+Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
+and the model's sampling defaults.
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
+CUDA graphs without compiling the codec with Inductor.
+
+For sustained high concurrency on a large-memory CUDA GPU, use:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+The high-concurrency profile places both stages on one GPU, sets each stage's
+capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
+mixed FULL CUDA graphs with both its token budget and capture limit set to
+512, so new prompt chunks remain within graph coverage. It also enables
+`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
+continuation control stay in GPU request slots, gathered through MRV2's
+batch mapping. Prefill conditioning and explicit per-request sampling seeds
+retain their original behavior. Published audio codes own their storage, so
+slot reuse cannot overwrite an in-flight output. Set the override to `false`
+to compare with the generic Omni model state. The profile selects
+`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
+disabled, and codec CUDA graph buckets through 128. The configuration was
+validated on one H200. The earlier capacity-256 variant used about 110 GiB of
+sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
+graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
+compilation can take several minutes; subsequent starts can reuse the AOT
+cache. Graph capture is still performed at startup.
+
+The high-concurrency profile also coalesces ready codec streams with
+`connectors.shm.extra.generation_min_batch_size: 32` and
+`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
+neither caps execution at 32 nor guarantees every codec group has 32 rows.
+First and final chunks share the bounded window at high concurrency. When
+fewer than 32 requests remain, dispatch is immediate. Pending outputs are
+retired before waiting, with the original deadline preserved, so waiting does
+not prevent the previous batch from releasing its in-flight state. Cancellation
+and input notifications keep their ordinary scheduler bookkeeping. This option
+requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
+`0` to disable coalescing; other profiles keep the immediate-dispatch default.
+
+The high-concurrency profile enables two additional Local model-state
+optimizations under the CUDA stage-0 `hf_overrides`, alongside
+`mrv2_gpu_slot_state: true`:
+
+- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
+  CPU reference-code slices and their destination positions into one pinned
+  upload. Host staging is reused only after the upload event completes;
+  GPU-resident references keep the original path.
+- `mrv2_direct_tokens: true` returns the already determined text token after
+  the Local audio/stop decision. It avoids constructing and sampling the
+  full text vocabulary. Audio-code and binary-stop sampling are unchanged.
+  Requests requiring distribution metadata or token constraints, and
+  unsupported execution modes, retain the normal sampler.
+
+Other profiles leave both flags disabled. Set both flags explicitly when
+comparing their combined effect. Deployment
+inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
+retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
+throughput: faster prompt preparation can change competition between the
+Talker and codec sharing the GPU.
+
+For the single-H200 capacity-128 throughput baseline, keep both switches on.
+In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
+rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
+tokens alone had no reliable gain. The combined path had 685 ms mean first
+audio versus 679 ms across six control rounds, with lower mean completion time.
+A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
+3.78% for the control. This sample does not establish speaker similarity or
+exclude small quality changes. At concurrency 256, the combined path had
+raised mean first-audio latency by about 0.33 s, so other profiles leave the
+switches off until their latency and quality are checked.
+
+The optional connector setting `generation_coalescing_policy: idle_wait`
+waits for an inbox notification when no stream is runnable, all receivers
+are parked and registered, and no output needs retirement. The first ready
+arrival starts a fresh coalescing window, preserving the batch budget.
+Control messages, including cancellations, can wait up to two windows
+instead of one (24 ms at the 12 ms setting). The default policy remains
+`fixed`; neither policy changes the global orchestration default.
+
+The codec backends differ in state access:
+
+- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
+- `triton` replaces the attention calculation, retaining the ring-cache
+  gather/copy and explicit mask construction.
+- `triton_slot` writes surviving K/V directly into request slots, attends
+  directly to the ring, and advances the active slot offsets in three
+  ordered kernels. Graph-padding rows do not advance persistent state.
+
+The slot path preserves the existing chunk-complete ring semantics, including
+retaining the final cache-capacity tokens when a chunk exceeds the ring.
+It does not introduce another cache owner or change request-slot lifetime.
+Zero-length padding rows skip the attention computation and emit zeros.
+
+The high-concurrency profile includes intermediate codec batch buckets 6, 12
+and 24. To compare with power-of-two buckets, change the stage-1
+`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
+Keep the maximum bucket equal to the state capacity to retain terminal-tail
+coalescing. Smaller buckets reduce padding work but require more graphs;
+measure complete serving runs before selecting them for a deployment.
+
+For a codec output-path comparison, set
+`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
+codec output path; `1` selects GPU output snapshots on a private codec
+stream. Both configurations use MRV2 and keep streaming audio responses.
+The option changes input ownership, metadata staging and snapshot handling
+as well as transfer scheduling, so its timing difference is not just D2H
+copy time.
+
+Those chunk-complete semantics truncate the causal window: every ring holds
+exactly `context` entries and a whole chunk is written before attention runs,
+so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
+tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
+400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
+its first 80 tokens attend to nothing and no chunk sees the previous one there.
+Padding a terminal tail to 15 frames therefore decodes it differently from an
+exact-length decode; the difference is deterministic (identical in fp32 and
+fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
+ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
+then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
+equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
+observed about a 2% throughput cost. It is off by default and its perceptual
+quality benefit has not been established. Terminal tails are padded
+into the regular 15-frame graph bucket in either mode.
+The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
+retain their existing attention implementation. Event-driven orchestration
+remains independently selectable and its default is unchanged.
+
+### Optional progressive chunks
+
+Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
+before steady decoding, for example:
+
+```yaml
+connectors:
+  shm:
+    extra:
+      codec_chunk_ramp: [1, 4, 15]
+```
+
+The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
+entry, chunks use `codec_chunk_frames`. Each request advances independently;
+the final partial chunk is flushed and an empty terminal carries no codec
+tokens. The codec captures every ramp length, sizes the maximum execution
+step for the largest entry, and uses the ramp's first entry for its dedicated
+first-chunk graph. Additional shapes increase startup and graph memory.
+
+Ramps are opt-in. They can reduce early playback underrun while increasing
+codec calls, first-audio latency or total generation time. The measured C128
+throughput presets retain 1→15; their published figures do not establish an
+additional benefit from enabling ramps. The processor and graph-shape tests
+run without GPU or model weights:
+
+```bash
+CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
+  tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
+  -m 'core_model and cpu' --run-level core_model
+```
+
+### Low-latency and reference-encoding options
+
+`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
+uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
+Its codec first-chunk fast path uses dedicated graphs and a stream handoff
+before the regular decoder reuses the request slot. A gate limits contention
+with regular codec work. Slot waits are bounded at 30 seconds. A failed or
+stalled decode raises an error rather than replaying a partially advanced
+slot; a failed output enqueue retains owned PCM for regular delivery.
+Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
+policy; it does not change the client concurrency or stage capacities.
+
+Reference encoding runs in the API layer, independently of MRV2. Enable
+reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
+eager encoder is retained when the option is absent. The graph path uses the
+loaded tokenizer's encoder/quantizer, length buckets and optional compilation
+(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
+attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
+Compilation and attention changes need not produce bit-identical codes.
+
+For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
+reference-code storage and, by default, a single encoder host with four
+workers. Use a dedicated directory per service, checkpoint and encoding
+configuration. Workers have their own graph resources; all graph captures
+complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
+separate encoders while sharing codes. More workers/graphs consume memory;
+they do not imply more GPUs. Example after selecting an available GPU and
+configuring any operator-managed MPS daemon:
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
+VLLM_OMNI_MOSS_REF_GRAPHS=1 \
+VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
+VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
+VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
+VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
+VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --api-server-count 4 \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+This example is a configurable serving profile, not the complete launch
+configuration of a historical benchmark. In particular, the low-latency YAML
+uses synchronous stage-1 scheduling; the retained source46 experiment used
+asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
+environment settings. Do not infer a throughput result from the YAML alone.
+
+### Cold versus hot reference measurements
+
+A cold reference misses the encoded-reference cache: the request pays for
+reference parsing/preparation and GPU encoding before speech generation.
+A hot reference reuses reference codes but still synthesizes new speech;
+generated audio is not cached. Pre-registering a reference moves encoding
+into registration and can make subsequent synthesis hot, but registration
+latency is still part of the first-use cost.
+
+Model startup/compilation is separate from reference coldness. EN1088 contains
+repeated references, so its first pass is not uniformly cold. For an all-unique
+cold pass, use its 666 first-occurrence references, keep order, warm execution
+with disjoint references, then repeat the same 666 requests. A second pass
+through independent API-local caches is not guaranteed hot; verify cache
+coverage before comparing it with a shared-cache result. Report requests/s
+alongside audio-s/s and retain failed, empty, long and near-silent outputs.
+
+### Reproduce the serving benchmark
+
+Use the complete Seed-TTS English test set, including reference audio and text.
+Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
+entries. Run the native benchmark after the server is ready:
+
+```bash
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
+export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
+export SEED_TTS_WER_EVAL=0
+
+for concurrency in 128; do
+  for phase in warm r1 r2; do
+    vllm bench serve --omni \
+      --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+      --backend openai-audio-speech --endpoint /v1/audio/speech \
+      --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
+      --seed-tts-locale en --disable-shuffle --num-prompts 1088 \
+      --num-warmups 0 --ready-check-timeout-sec 0 \
+      --output-len 256 --max-concurrency "$concurrency" \
+      --request-rate inf --seed 42 \
+      --extra-body '{"task_type":"Base","max_new_tokens":256}' \
+      --save-result --save-detailed --result-dir results/moss-local \
+      --result-filename "c${concurrency}-${phase}.json"
+  done
+done
+```
+
+Discard the complete `warm` pass and combine measured runs as total generated
+audio seconds divided by total benchmark duration. Require 1088 successes,
+zero failures and 1088 nonempty-audio metric samples in each run. The output
+cap matches the benchmark protocol; success alone does not establish speech
+quality or that every sentence ended before the cap. Dataset seed 42 does
+not fix an independent sampling seed for every request.
+
+For an attention-only comparison, copy the high-concurrency YAML beside the
+original to preserve relative `base_config` resolution. Change its
+codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
+and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
+other settings, warmup and client concurrency identical. Restart the server
+between configurations and use separate result directories. Comparing the
+compiled slot profile to uncompiled Triton includes both changes.
+
+Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
+and no model weights:
+
+```bash
+python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
+  tests/model_executor/models/moss_tts/test_streaming_attention.py \
+  -m 'core_model and cuda' --run-level=core_model
+```
@@ -40,9 +40,11 @@ user-downloaded weights and bundles none.
 | Output | 48 kHz stereo, whole song delivered when the semantic request finishes |

 Generation is one or two requests on a single AR stage: the abc phase writes
-the score, the semantic phase generates codec frames and, on its last step,
-solves the ODE and decodes the whole song in-engine (weights load once; the
-VAE loads at startup from `$YUE2_VAE` or the hub id).
+the score, the semantic phase generates codec frames, then queues the ODE and VAE on a
+high-priority CUDA stream. Other requests keep decoding while the finishing
+request emits HOLD tokens (up to 4096 steps, within its token/context budget).
+The scheduler is unchanged. Weights load once; the VAE loads at startup from
+`$YUE2_VAE` or the hub id.

 ## Running

@@ -67,15 +69,50 @@ ending.

 ## Notes

- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine
-  reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the
-  9000-frame cap (worst case), peak usage during the terminal NAR/VAE
-  finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the
-  same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70.
- **Known limitations:** 4 concurrent full-length requests are the verified
-  shape (`max_num_seqs: 4`); CUDA graph capture is a
-  follow-up; the abc phase runs eagerly
-  after prefill (its tokens are the product, not audio).
+The lifecycle e2e regression forces KV-cache preemption with chunked and
+unchunked recompute, mixes ABC and semantic sampling, and cancels a request
+after NAR work has started. It checks history alignment, released synthesis
+buffers, and successful generation after cancellation. Run with one H100/H200
+and cached YuE2-3B and YuE2-Vae weights:
+
+```bash
+export YUE2_MODEL_DIR=/path/to/YuE2-3B
+export YUE2_VAE=/path/to/YuE2-Vae
+CUDA_VISIBLE_DEVICES=0 python -m pytest tests/e2e/online_serving/test_yue2.py \
+    -k preemption_and_synthesis_abort -m 'slow and tts' --run-level full_model -q -s
+```
+
+The test retains the module's weekly TTS routing; it is not a per-PR CI gate.
+
+- **Memory:** `gpu_memory_utilization` budgets the vLLM engine; NAR K/V,
+  acoustic graph buffers and VAE activations also need room during synthesis.
+  The 0.70 / 4-slot defaults require validation on the target card. The earlier
+  RTX 4090 measurement (15.8 GiB after startup, 20.9 GiB peak at 9000 frames)
+  predates the async/compiled implementation; it does not validate this
+  version or long ABC prefixes on a 4090. Reduce the fraction if the target
+  workload needs more synthesis memory, especially for long ABC prefixes and
+  the 9000-frame cap. Releasing completed chunks bounds live NAR buffers,
+  while the shared graph allocator retains reserved storage for reuse.
+- **CUDA graphs:** the default deploy keeps the AR backbone eager. Both ABC
+  and semantic sampling use prewarmed power-of-two graph buckets. Custom
+  sampler captures have a bounded cache and fall back to eager sampling when
+  it fills. The NAR velocity pass uses a chunk graph on the serialized synthesis
+  stream. Chunk graphs share one allocator pool; completed chunks release
+  their engines and buffers after their own events. FA3 runs on Hopper; other
+  CUDA cards use SDPA, which is also warmed at startup.
+- **Serving:** run `vllm serve m-a-p/YuE2-3B --omni` with the default
+  `yue2.yaml`, which uses eager AR with 4 slots. Validate memory capacity
+  for the target card and workload.
+  For the tested single-H200 graph settings, reuse this deploy with
+  `--stage-overrides '{"0":{"max_num_seqs":32,"gpu_memory_utilization":0.5,"enforce_eager":false,"compilation_config":{"cudagraph_mode":"FULL_AND_PIECEWISE"}}}'`.
+  Decode uses FULL and mixed prefill uses PIECEWISE; VAE and the Python
+  synthesis queue remain outside these graphs.
+  Audio is delivered as a whole song; time to first audio equals completion
+  latency. Aborting a running song stops future work units and waits only for
+  the at-most-two submitted units before releasing its buffers.
+- **Preemption:** before drawing for a resumed row, model-owned sampling
+  reconciles the codec history and terminal flags with the engine's retained
+  sequence. Discarded draws cannot advance its penalty window or frame budget.
 - Audio is not expected to match the upstream torch reference bit-for-bit
  (fused vs eager kernels). Measured against the upstream torch reference
  (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token
@@ -2928,9 +2928,17 @@ class TestPlatformOverrides:

        assert deploy.stages[0].engine_extras["block_size"] == 128

-    @pytest.mark.parametrize("platform", ["cuda", "npu"])
-    def test_recommended_native_runner_platform_defaults(self, platform):
-        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+    @pytest.mark.parametrize(
+        "filename,pipeline_key",
+        [
+            ("qwen3_tts_mrv2.yaml", "qwen3_tts"),
+            ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+        ],
+    )
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
        deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
        deploy = _apply_platform_overrides(deploy, platform=platform)
        expect_v2 = platform == "cuda"
@@ -2946,6 +2954,60 @@ class TestPlatformOverrides:
            # v2 only engages the native plane on stages declaring support.
            assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)

+    def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
+        base = load_deploy_config(base_path)
+        candidate = load_deploy_config(candidate_path)
+        assert base.model_runner == "v1"
+        assert candidate.stages == base.stages
+        assert candidate.connectors == base.connectors
+        assert candidate.async_chunk == base.async_chunk
+        cuda = _apply_platform_overrides(candidate, platform="cuda")
+        stages = merge_pipeline_deploy(pipeline, cuda)
+        assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
+        codec_args = stages[1].yaml_engine_args
+        assert codec_args["max_num_seqs"] == 64
+        assert codec_args["enforce_eager"] is False
+        assert codec_args["compilation_config"] == {
+            "mode": 0,
+            "cudagraph_mode": "FULL",
+            "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
+            "cudagraph_num_of_warmups": 1,
+        }
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            # Platform resolution mutates the deployment; each server loads
+            # a fresh profile before applying its own platform overrides.
+            fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
+            original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert fallback.stages == original.stages
+        for variant in ("moss_tts_delay", "moss_tts_realtime"):
+            assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
+
+    def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
+        talker, codec = merge_pipeline_deploy(pipeline, deploy)
+        assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+        assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
+        assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
+        assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
+        compilation = codec.yaml_engine_args["compilation_config"]
+        assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
+        assert compilation["cudagraph_mode"] == "FULL"
+        assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+        assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
+        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+            expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert actual.stages == expected.stages
+
    @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
    def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
        from unittest.mock import Mock
@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
    merged = resolve_deploy_yaml(leaf)
    assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
    assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)
@@ -167,9 +167,9 @@ def test_inventory_matches_reviewed_snapshot_counts():
        EnvironmentVariableCategory.PUBLIC_OMNI: 32,
        EnvironmentVariableCategory.INHERITED_VLLM: 20,
        EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
        EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 6,
+        EnvironmentVariableCategory.INTERNAL: 7,
    }

    disposition_counts = Counter(
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
        if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
    )
    assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 46,
+        ModelEnvironmentVariableDisposition.PROMOTE: 66,
        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
        ModelEnvironmentVariableDisposition.EXTERNAL: 0,
        ModelEnvironmentVariableDisposition.INTERNALIZE: 17,
@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
        "interleave_mm_strings",
        "media_io_kwargs",
        "final_output",
+        "supports_running_prefix_cache_reset",
        "active_stream_window",
        "session_mode",
        "duplex_max_sessions",
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded readiness coalescing must not allocate or consume request state."""
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
+    events = [notification("0"), notification("1", "2", "3")]
+    s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.006)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
+
+
+@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
+def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    if pending == "in_flight":
+        s.requests["0"].num_in_flight_tokens = 1
+    elif pending == "admission":
+        s.requests["0"].status = RequestStatus.WAITING
+    elif pending == "registration":
+        s.input_coordinator.pending_chunk_registrations = [object()]
+    elif pending == "terminal":
+        s._pending_data_plane_terminal_req_ids = {"0"}
+    else:
+        s.requests.clear()
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == 10
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == pytest.approx(10.012)
+
+
+def notification(*ids, terminal=()):
+    return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
+
+
+class Inbox:
+    def __init__(self, clock, initial=(), delayed=()):
+        self.clock = clock
+        self.initial = list(initial)
+        self.delayed = list(delayed)
+        self.waits = []
+
+    def get_nowait(self):
+        if self.initial:
+            return self.initial.pop(0)
+        raise queue.Empty
+
+    def get(self, timeout):
+        self.waits.append(timeout)
+        if not self.delayed or self.delayed[0][0] > timeout:
+            self.clock[0] += timeout
+            raise queue.Empty
+        delay, value = self.delayed.pop(0)
+        self.clock[0] += delay
+        return value
+
+
+def scheduler(monkeypatch, initial=(), delayed=()):
+    clock = [10.0]
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
+    s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    s._generation_max_wait_s = 0.012
+    s._generation_min_batch_size = 4
+    s._pause_state = PauseState.UNPAUSED
+    s.requests = {
+        str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
+    }
+    s.running = list(s.requests.values())
+    s.input_coordinator = SimpleNamespace(
+        requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
+    )
+    s._latest_omni_connector_output = None
+    s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
+    return s, clock
+
+
+def test_arrivals_form_target_batch_in_order(monkeypatch):
+    events = [notification("0"), notification("1", "2"), notification("3")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.005)
+    assert not s.input_coordinator.requests_with_ready_chunks
+    assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
+
+
+def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
+    events = [notification("0", "gone", "1", "2"), notification("0")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
+    s.requests["1"].is_finished = lambda: True
+    s.requests["2"].is_finished = lambda: True
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
+
+
+@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
+def test_immediate_dispatch(monkeypatch, reason):
+    event = notification("0")
+    s, clock = scheduler(monkeypatch, [event])
+    if reason == "enough":
+        event.chunk_ready_req_ids.update(["1", "2", "3"])
+    elif reason == "empty":
+        event.chunk_ready_req_ids.clear()
+    elif reason == "low":
+        s.requests = {"0": s.requests["0"]}
+    elif reason == "paused":
+        s._pause_state = PauseState.PAUSED_ALL
+    elif reason == "off":
+        s._generation_max_wait_s = 0
+    assert s._drain_omni_connector_outputs() == [event]
+    assert not s._omni_connector_output_inbox.waits
+    assert clock[0] == 10
+
+
+def test_terminal_arrival_keeps_original_deadline(monkeypatch):
+    events = [notification("0"), notification("1", terminal=["1"])]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
+    s, _ = scheduler(monkeypatch)
+    s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
+    last = notification("3")
+    s._latest_omni_connector_output = last
+    assert s._drain_omni_connector_outputs() == [last]
+    assert s._latest_omni_connector_output is None
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.012)
+    s._omni_connector_output_inbox.initial.append(notification("3"))
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.024)
+
+
+def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
+    event = notification("0", "cancelled")
+    event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
+    s, _ = scheduler(monkeypatch, [event])
+    s._generation_max_wait_s = 0
+    s.waiting = []
+    s.input_coordinator._async_chunk = True
+    s.input_coordinator.update_request_metadata = mocker.Mock()
+    s.input_coordinator.process_pending_chunks = mocker.Mock()
+    s._consume_pending_connector_output("generation")
+    s.input_coordinator.update_request_metadata.assert_called_once_with(
+        s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
+    )
+    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+
+
+@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])
+def test_invalid_config_rejected(monkeypatch, target, wait):
+    from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
+
+    def initialize(s, *args, **kwargs):
+        s.max_num_running_reqs = 256
+        s.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                retains_state_across_chunks=True,
+                stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
+            )
+        )
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
+    monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
+    with pytest.raises(ValueError):
+        OmniGenerationScheduler()
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
+    event = notification("0", terminal=["0"] if terminal else [])
+    s, clock = scheduler(monkeypatch, [event])
+    s.running = s.running[1:]
+    assert s._drain_omni_connector_outputs() == [event]
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
+    events = [notification("0"), notification("1")]
+    s, clock = scheduler(monkeypatch, events[:1])
+    s.requests["7"].num_in_flight_tokens = 1
+    assert s._drain_omni_connector_outputs() == events[:1]
+    assert s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+    s.input_coordinator.requests_with_ready_chunks.add("0")
+    # Simulate EngineCore retiring the previous batch after five milliseconds.
+    clock[0] += 0.005
+    s.requests["7"].num_in_flight_tokens = 0
+    s._omni_connector_output_inbox.initial.append(events[1])
+    assert s._drain_omni_connector_outputs() == events[1:]
+    assert not s._generation_defer_batch
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
+    assert clock[0] == pytest.approx(10.012)
+    assert s._generation_batch_deadline is None
+
+
+def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")])
+    s.requests["7"].num_in_flight_tokens = 1
+    s._generation_batch_deadline = clock[0] - 0.001
+    assert len(s._drain_omni_connector_outputs()) == 1
+    assert not s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
+    from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
+
+    request = _chunk_request("pending")
+    s = _make_generation_scheduler(request)
+    s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
+    allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
+    output = s.schedule()
+    assert output.total_num_scheduled_tokens == 0
+    assert not output.num_scheduled_tokens
+    allocate.assert_not_called()
+    assert list(s.waiting) == [request]
+    assert request.num_computed_tokens == request.num_in_flight_tokens == 0
+    assert s.chunk_transfer_adapter.restore_called
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise generation constructor policy after actual deploy/platform resolution."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def construct_scheduler(monkeypatch, mocker):
+    # Only upstream cache allocation and connector construction are omitted.
+    # Keep the actual Omni I/O initialization, native decision and constructor.
+    mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
+
+    def upstream_init(self, config):
+        self.vllm_config = config
+        self.max_num_running_reqs = config.scheduler_config.max_num_seqs
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
+    return OmniGenerationScheduler
+
+
+def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            use_v2_model_runner=native,
+            supports_native_mrv2_data_plane=True,
+            retains_state_across_chunks=stateful,
+            async_chunk=True,
+            stage_id=1,
+            stage_connector_config={"extra": extras or {}},
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
+        scheduler_config=SimpleNamespace(max_num_seqs=capacity),
+    )
+
+
+@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+def test_moss_profile_generation_constructor_after_platform_resolution(
+    construct_scheduler, monkeypatch, mocker, profile, platform
+):
+    # merge_pipeline_deploy owns platform resolution. Mock the detected device
+    # instead of applying an override first and then resolving the worker's
+    # actual platform a second time.
+    monkeypatch.setattr(current_omni_platform, "device_name", platform)
+    deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
+    pipeline = resolve_pipeline_config("moss_tts_local")
+    codec = merge_pipeline_deploy(pipeline, deploy)[1]
+    extras = deploy.connectors["shm"]["extra"]
+    config = _config(
+        native=codec.yaml_engine_args["use_v2_model_runner"],
+        stateful=pipeline.stages[1].retains_state_across_chunks,
+        capacity=codec.yaml_engine_args["max_num_seqs"],
+        extras=extras,
+    )
+    warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
+    scheduler = construct_scheduler(config)
+    # The raw extras persist through fallback; the live constructor must cope.
+    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    if platform == "cuda":
+        assert scheduler._native_data_plane and scheduler.input_coordinator is not None
+        assert scheduler._generation_min_batch_size == 32
+        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
+        warning.assert_not_called()
+        if profile == "low_latency":
+            assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
+    else:
+        assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
+        assert scheduler._generation_min_batch_size == 1
+        assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
+        warning.assert_called_once()
+
+
+@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
+def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
+    config = _config(
+        stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
+    )
+    with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
+        construct_scheduler(config)
+
+
+@pytest.mark.parametrize(
+    "extras,message",
+    [
+        ({"generation_min_batch_size": 0}, "generation_min_batch_size"),
+        ({"generation_min_batch_size": 129}, "generation_min_batch_size"),
+        ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
+        ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
+        ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
+        ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
+    ],
+)
+def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
+    with pytest.raises(ValueError, match=message):
+        construct_scheduler(_config(native=False, extras=extras))
@@ -1,10 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace

 import pytest
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.request import RequestStatus

+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
 from vllm_omni.engine.serialization import deserialize_additional_information

@@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest:
    )


-def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
    tagged = apply_omni_final_stage_metadata(
        _make_engine_request(),
        final_stage_id=0,
@@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
        request_id="req",
        additional_information=tagged.additional_information,
    )
-    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
    metadata = deserialize_additional_information(tagged.additional_information)
-    return result, metadata
+    return omits_kv, omits_chunk, metadata


 def test_stage_zero_request_omits_kv_transfer():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)

-    assert omits_transfer
+    assert omits_kv
+    assert omits_chunk
    assert "omni_force_kv_transfer" not in metadata


 def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)

-    assert not omits_transfer
+    assert not omits_kv
+    assert omits_chunk
    assert metadata["omni_final_stage_id"] == 0
    assert metadata["omni_force_kv_transfer"] is True
+
+
+class _ChunkRequest(SimpleNamespace):
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
+
+
+def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=final_stage_id,
+        force_kv_transfer=force_kv_transfer,
+    )
+    request = _ChunkRequest(
+        request_id="req-shm-leak",
+        additional_information=tagged.additional_information,
+        num_in_flight_tokens=1,
+        num_stale_output_tokens=0,
+        sampling_params=SimpleNamespace(num_logprobs=None),
+        has_encoder_inputs=False,
+        pooling_params=None,
+        status=RequestStatus.RUNNING,
+        resumable=False,
+        num_output_placeholders=0,
+        spec_token_ids=[],
+        _output_token_ids=[],
+        num_computed_tokens=1,
+        stop_reason=None,
+        num_nans_in_logits=None,
+        client_index=0,
+        trace_headers=None,
+    )
+    request.is_finished = lambda: False
+    request.get_finished_reason = lambda: None
+    request.take_prefill_stats = lambda: None
+    request.take_events = lambda: None
+    return request
+
+
+def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
+    adapter = mocker.MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 1
+
+    sched = mocker.MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._update_request_with_output.return_value = ([42], True)
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._handle_stopped_request.return_value = True
+
+    def _free_request(req, delay_free_blocks=False):
+        sched._omits_kv_transfer_cache.pop(req.request_id, None)
+        return None, None
+
+    sched._free_request.side_effect = _free_request
+    sched.chunk_transfer_adapter = adapter
+    sched.running = [request]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
+
+    scheduler_output = mocker.MagicMock()
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = mocker.MagicMock()
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = inter_stage_outputs
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return adapter
+
+
+@pytest.mark.parametrize(
+    ("final_stage_id", "force_kv_transfer", "expect_save"),
+    [
+        pytest.param(0, False, False, id="stage-zero-final"),
+        pytest.param(1, False, True, id="downstream-final"),
+    ],
+)
+def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
+    request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
+    adapter = _run_finished_save_step(mocker, request)
+
+    assert adapter.save_async.called is expect_save
+
+
+def test_stage_zero_final_skips_inter_stage_output(mocker):
+    warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
+    request = _make_chunk_request(0)
+    adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
+
+    adapter.save_async.assert_not_called()
+    assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
+
+
+def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    scheduler.skipped_waiting = []
+    scheduler.log_stats = False
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    session = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+        status=RequestStatus.WAITING,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert "req" in scheduler._omits_kv_transfer_cache
+    update = SimpleNamespace(
+        additional_information={"omni_final_stage_id": 1},
+        model_intermediate_buffer=None,
+        arrival_time=0.0,
+        sampling_params=None,
+    )
+
+    OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
+
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert session.additional_information == {"omni_final_stage_id": 1}
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
+
+
+def test_omits_cache_misses_when_payload_object_is_replaced():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
+    request.additional_information = {"omni_final_stage_id": 1}
+
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)
@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers

 # isort: on

@@ -103,10 +104,11 @@ def _make_drain_sched(session: Request) -> MagicMock:
    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
    sched.finished_req_ids_dict = {}
    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
    return sched


-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
    """Feed one model-runner frame through update_from_output.

    Returns True when the frame's tokens were delivered (the append ran),
@@ -114,7 +116,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
    """
    scheduler_output = MagicMock(spec=SchedulerOutput)
    scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
    scheduler_output.num_invalid_spec_tokens = 0

    model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -134,6 +136,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
    return sched._update_request_with_output.called


+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
    """One in-flight decode frame at replacement: its late output drains the
    counter exactly and the new segment's first frame is delivered."""
@@ -225,3 +248,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
    assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
    assert session.num_stale_output_tokens == 0
    assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()
@@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
    OmniChunkTransferAdapter,
 )
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers

 # isort: on

@@ -175,6 +176,7 @@ def _run_resumable_segment_stop(
    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
    sched.finished_req_ids_dict = {}
    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)

    scheduler_output = MagicMock(spec=SchedulerOutput)
    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -312,6 +314,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
    sched.finished_req_ids_dict = {}
    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)

    scheduler_output = MagicMock(spec=SchedulerOutput)
    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -400,6 +403,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
    sched.finished_req_ids_dict = {}
    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)

    scheduler_output = MagicMock(spec=SchedulerOutput)
    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -126,6 +126,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
    sched.max_num_running_reqs = 8
    sched._native_data_plane = False
    sched._retains_state_across_chunks = False
+    # Unset MagicMock attributes are truthy. Keep optional scheduling modes
+    # disabled so this fixture actually schedules its one-shot request.
+    sched._generation_defer_batch = False
+    sched._generation_first_chunk_only = False
+    sched._generation_max_regular_batch = 0
    sched._pause_state = PauseState.UNPAUSED
    sched.requests = {request.request_id: request}
    sched.running = [request]
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations

 from types import SimpleNamespace
@@ -76,7 +79,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
    scheduler = OmniARScheduler.__new__(OmniARScheduler)
    scheduler.input_coordinator = coordinator
    scheduler.chunk_transfer_adapter = None
-    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
    scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
    scheduler.finished_req_ids = set()
    scheduler.finished_req_ids_dict = None
@@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():

 def test_budget_reduced_capacity_drives_runner_lru():
    pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
+    #
+    # These read 128 and 192 while the paging unit was the frame. tiny_spec
+    # declares one token per frame, so it used to page a single token at a
+    # time -- a block size no attention kernel accepts. Paging at 16 makes
+    # each page sixteen times larger. The block *counts* are unchanged, and
+    # so is what this test is about: a budget that fits fewer sessions than
+    # the pipeline asked for must drive the runner's LRU.
    runner = make_runner(
        pipeline,
-        available_bytes=128,
+        available_bytes=2048,
        gpu_memory_fraction=1.0,
    )
    kv = runner.kv_cache
@@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
    monkeypatch.setattr(runner, "execute_model", fail_if_called)
    runner._warmup_ar_rollout()
    assert execute.called is False
+
+
+# ── paging unit vs eviction unit ────────────────────────────────────────────
+
+
+def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
+    """Every resolution that runs today must page exactly as it does today."""
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for tokens_per_frame in (16, 320, 1440, 4096):
+        assert tokens_per_frame % 16 == 0
+        assert paging_block_size(tokens_per_frame) == tokens_per_frame
+
+
+def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
+    """832x480 is the checkpoint's own default and the kernel rejects it.
+
+    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
+    frame-sized block is not something FlashAttention's paged kernel accepts.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 16 == 8
+    assert paging_block_size(1560) == 16
+    # And whatever it returns is always something the kernel takes.
+    for tokens_per_frame in range(1, 200):
+        assert paging_block_size(tokens_per_frame) % 16 == 0
+
+
+# ── the kernel's constraint is data, not a constant ─────────────────────────
+
+
+def _multiple_of(base):
+    from vllm.v1.attention.backend import MultipleOf
+
+    return MultipleOf(base)
+
+
+def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
+    """AR-Diffusion calls flash_attn itself, so no backend can be asked.
+
+    The answer has to come from the module that picks the kernel, and it has to
+    have the shape vLLM uses, so the policy that consumes it does not care
+    where it came from.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+    advertised = supported_kernel_block_sizes()
+    assert advertised, "a kernel that accepts nothing cannot be paged for"
+    for entry in advertised:
+        assert isinstance(entry, int | MultipleOf)
+
+
+def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
+    """hpc_attn advertises [64] -- not a multiple, a single legal value.
+
+    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
+    be the block and the only thing left to page at is 64. Assuming multiples
+    of 16 here would hand the kernel a block size it rejects.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 64 == 24
+    assert paging_block_size(1560, [64]) == 64
+    assert paging_block_size(64, [64]) == 64
+    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size
+
+
+def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
+    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.
+
+    The same checkpoint on the two cards must therefore page differently, which
+    is exactly what a constant cannot express.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    hopper = [_multiple_of(16)]
+    blackwell = [_multiple_of(16), 128, 256]
+    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
+    assert paging_block_size(1440, blackwell) == 1440
+    assert paging_block_size(1560, hopper) == 16
+    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest
+
+
+def test_whatever_comes_back_is_always_something_the_kernel_accepts():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for advertised, check in (
+        ([_multiple_of(16)], lambda n: n % 16 == 0),
+        ([64], lambda n: n == 64),
+        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
+    ):
+        for tokens_per_frame in range(1, 300):
+            assert check(paging_block_size(tokens_per_frame, advertised))
+
+
+def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    with pytest.raises(ValueError, match="no legal block size"):
+        paging_block_size(1560, [])
@@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected():
        )


+# --- eviction against the block table ---------------------------------------
+
+
+class _RecordingManager(ChunkWindowManager):
+    """Capture the block range eviction asks for, without a real block pool.
+
+    ``remove_skipped_blocks`` reads three attributes and calls one method. A
+    pool would add nothing: the bug under test is the conversion from the
+    spec's units into the block table's, and that is entirely visible here.
+    """
+
+    def __init__(self, spec):
+        self.kv_cache_spec = spec
+        self.sliding_window = spec.sliding_window
+        self.block_size = spec.block_size
+        self.freed: tuple[int, int] | None = None
+
+    def _remove_blocks_in_range(self, request_id, first_block, last_block):
+        self.freed = (first_block, last_block)
+
+
+def test_sink_is_preserved_when_a_frame_spans_many_blocks():
+    # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is
+    # not a legal block size, so paging falls back to 16 and a frame spans 98
+    # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame
+    # is a block and a sink counted in frames indexes the block table correctly
+    # by accident -- which is why this went unnoticed.
+    chunk_size, sink_chunks, window_chunks = 1560, 9, 9
+    spec = ChunkWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+    )
+    manager = _RecordingManager(spec)
+
+    # Far enough past sink + window that eviction has definitely started.
+    manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size)
+    assert manager.freed is not None
+    first_freed, last_freed = manager.freed
+
+    # The assertion is on the token boundary rather than a block count, so it
+    # stays meaningful if the paging unit changes again.
+    sink_tokens = sink_chunks * chunk_size
+    assert first_freed * spec.block_size >= sink_tokens, (
+        f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink"
+    )
+    # And it must not over-protect: at most one straddling block beyond the sink.
+    assert (first_freed - 1) * spec.block_size < sink_tokens
+    # Only wholly-skipped blocks are released.
+    skipped = manager.get_num_skipped_tokens(30 * chunk_size)
+    assert last_freed * spec.block_size <= sink_tokens + skipped
+
+
+def test_sink_conversion_is_identity_when_a_frame_is_a_block():
+    # Every geometry that paged one frame per block must be untouched.
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1)
+    manager = _RecordingManager(spec)
+    manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK)
+    assert manager.freed is not None
+    first_freed, _ = manager.freed
+    assert first_freed == spec.sink_chunks
+
+
 def test_staged_window_reuse_matches_a_full_gather():
    """Reusing the staged history must be byte-identical to re-gathering the whole window.

@@ -23,6 +23,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState

 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -36,10 +37,26 @@ NEG = "negative"


 def make_state(
-    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+    *,
+    num_layers=1,
+    window_chunks=2,
+    chunk_size=BLOCK,
+    sink_chunks=0,
+    reset_at_boundary=False,
+    dtype=torch.float32,
+    device=torch.device("cpu"),
+    reuse_history_staging=False,
 ):
+    """Build a cache. ``chunk_size`` defaults to the block size, but the two are
+    independent -- the shipped 832x480 gives 1560 tokens per frame against
+    16-token blocks, so a frame is 97.5 blocks."""
    cfg = ARDiffusionKVConfig(
-        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+        enable=True,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        reuse_history_staging=reuse_history_staging,
    )
    kv = ARDiffusionKVCache(
        cfg,
@@ -109,11 +126,13 @@ def _commit_video_span(
    n_chunks: int,
    dtype: torch.dtype,
    device: torch.device,
+    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    span = n_chunks * chunk_size
+    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
    ctx.ensure_video_slots(device)
-    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
    st.commit_paged_context(kv_branch)
@@ -547,6 +566,644 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa
        torch._dynamo.reset()


+# ── a frame that is not a whole number of blocks ────────────────────────────
+
+# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder
+# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds
+# chunk_size == BLOCK, so nothing there can reach this path.
+RAGGED_CHUNK = 24
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.parametrize("action_len", [0, 3])
+def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
+    """Attention reads a sequence, not a set of blocks.
+
+    With one committed chunk of 24 tokens the history stops 8 slots into its
+    last block. The committing path writes straight after it and is fine. The
+    scratch path used to restart at slot zero of a fresh region, which left
+    those 8 slots unwritten -- and since the kernel consumes the block table as
+    one contiguous run, it read them as if they were tokens and shifted every
+    token after them by 8. Shapes stayed correct throughout, so only the values
+    showed it.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    # window_chunks=2 with one chunk of history keeps the whole sequence
+    # resident, so the window does not also round up here; that case is
+    # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below.
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK
+    )
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8
+
+    current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, RAGGED_CHUNK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len, query_len=query.shape[1], device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    new_k = torch.cat([history_k, current_k], dim=1)
+    new_v = torch.cat([history_v, current_v], dim=1)
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
+    """The write targets themselves must leave no hole.
+
+    This is the property the values above depend on, asserted directly so a
+    regression names the cause rather than showing a numeric mismatch.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK)
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+
+    history_tail_block = ctx.history_block_ids[-1]
+    # The chunk starts inside the history's own last block, not at a fresh one.
+    assert ctx.current_video_block_ids[0] == history_tail_block
+    first_slot = int(ctx.current_video_slot_mapping[0])
+    assert first_slot == history_tail_block * BLOCK + 8
+
+    # And the action region must not be charged for that managed block.
+    ctx.ensure_action_slots(3, device)
+    assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1
+    assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids
+
+
+def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
+    """The table's shape must not track how full the window is.
+
+    With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5
+    blocks. Deriving the width by flooring that would understate the capacity,
+    the max() against the live block count would take over, and the shape would
+    change as the window filled -- recompiling the graph. The same floor would
+    also let max_seq_len come out under the kv_len actually passed.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    # The window is deliberately not a whole number of blocks.
+    max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size
+    assert max_video_tokens % BLOCK != 0
+
+    widths: set[int] = set()
+    kv_lens: list[int] = []
+    for _ in range(6):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table(
+            action_len=0, query_len=RAGGED_CHUNK, device=device
+        )
+        widths.add(int(block_table.shape[1]))
+        # A bound handed to the kernel has to actually bound the sequence.
+        assert int(seq_lens[0]) <= max_seq_len
+        kv_lens.append(int(seq_lens[0]))
+        st.commit_paged_context(POS)
+
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+    # The window has to have actually filled, or none of the above was tested.
+    assert max(kv_lens) > min(kv_lens)
+    assert max(kv_lens) <= ctx.max_video_blocks * BLOCK
+    # Whole blocks are read, so each of the two boundaries -- the sink's end and
+    # the recent window's start -- can contribute up to BLOCK - 1 tokens the
+    # window no longer needs.
+    assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1)
+
+
+def _poison_pools(kv: ARDiffusionKVCache) -> None:
+    for pool in (*kv._k_pools, *kv._v_pools):
+        pool.fill_(float("nan"))
+
+
+def _write_prepared_chunk(
+    kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first.
+
+    Blocks are handed out whole and a reused block keeps whatever it held, so a
+    slot read without having been written would not look unusual. Every slot of
+    the chunk's blocks that this chunk does not write is set to NaN first --
+    except the leading slots of a block the history still occupies -- so reading
+    one turns the attention output NaN.
+    """
+    mapping = ctx.current_video_slot_mapping
+    first_slot = int(mapping[0])
+    for index, block in enumerate(ctx.current_video_block_ids):
+        start = first_slot if index == 0 else block * BLOCK
+        kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan")
+        kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan")
+    k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][mapping] = k
+    kv._v_pools[0][mapping] = v
+    return k, v
+
+
+def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]:
+    """Token positions attention must read at ``end``.
+
+    Every token of the current chunk, plus every written history token in a
+    block that is still resident and holds a sink token or one of the window's
+    most recent tokens.
+    """
+    start = end - window
+    blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)}
+    return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)]
+
+
+def _model_block_indices(kv, adapter) -> dict[int, int]:
+    """Each resident block's index in model positions, by block id.
+
+    Compaction removes the evicted gap after the sink's blocks and shifts the
+    rest of the table down by it, so a table index past the sink is
+    ``compacted_tokens`` behind the position it holds.
+    """
+    block_size = kv.block_size
+    sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size)
+    shift = adapter.compacted_tokens // block_size
+    return {
+        int(block): index if index < sink_blocks else index + shift
+        for index, block in enumerate(kv.block_table(adapter))
+        if block != kv.null_block_id
+    }
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
+    sink_chunks, window_chunks, reset_at_boundary
+):
+    """Once the window slides, the table must still cover the window exactly.
+
+    The sink and the recent window are separate token ranges, and each can
+    straddle a block edge on its own: 24-token chunks against 16-token blocks
+    put a boundary 8 tokens past an edge on alternating chunks. Counting blocks
+    off the ends of the resident list dropped the block holding the window's
+    first tokens whenever both boundaries straddled, and the returned length
+    ran past the last written slot. Neither changes a shape, so this compares
+    values against dense attention over the tokens the window keeps -- after
+    eviction, with every unwritten slot poisoned, on both the scratch and the
+    committing forward of every chunk.
+
+    The reset case keeps only the sink across a boundary, so the recent window
+    reaches back over blocks that are no longer resident. Those must be skipped
+    by position rather than read.
+
+    Eight chunks are enough for compaction to drop the evicted gap from the
+    table, after which storage positions run behind model positions. Removing
+    a gap that is not a whole number of chunks moves where eviction snaps, and
+    in the reset case the window then keeps tokens a boundary should have
+    dropped.
+    """
+    torch.manual_seed(0)
+    device, dtype = torch.device("cpu"), torch.float32
+    kv, st = make_state(
+        device=device,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        chunk_size=RAGGED_CHUNK,
+    )
+    _poison_pools(kv)
+    sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window)
+    committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+    evicted = compacted = False
+
+    for chunk in range(8):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + RAGGED_CHUNK
+        for commit_current in (False, True):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+            ctx.ensure_video_slots(device)
+            k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device)
+            written = dict(committed)
+            written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)})
+
+            query = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+            block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+                action_len=0, query_len=RAGGED_CHUNK, device=device
+            )
+            resident = set(_model_block_indices(kv, st.adapter(POS)).values())
+            positions = _window_positions(
+                end, sink=sink, window=window, written=written, history=history, resident=resident
+            )
+            label = f"chunk {chunk}, commit_current={commit_current}"
+            assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers"
+            paged = ar_diffusion_paged_attention(
+                query,
+                kv.key_cache(0),
+                kv.value_cache(0),
+                block_table=block_table,
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                max_query_len=max_query_len,
+                max_seq_len=max_seq_len,
+                softmax_scale=HEAD_DIM**-0.5,
+                causal=False,
+            )
+            ref = _dense_attention(
+                query,
+                torch.stack([written[p][0] for p in positions]).unsqueeze(0),
+                torch.stack([written[p][1] for p in positions]).unsqueeze(0),
+            )
+            torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}")
+        committed = written
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+        evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS))
+
+    assert evicted, "the window never slid, so nothing after eviction was tested"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_compaction_does_not_change_which_tokens_stay_resident(
+    monkeypatch, sink_chunks, window_chunks, reset_at_boundary
+):
+    """Compaction renumbers the block table; it must not change what eviction keeps.
+
+    Storage positions run ``compacted_tokens`` behind model positions, and
+    eviction snaps to chunk boundaries counted in storage positions. Removing a
+    gap that is not a whole number of chunks moves every later snap, so eviction
+    keeps a different set of tokens -- in the reset case, tokens a boundary
+    should have dropped, which the window then reads. Attention checked against
+    what is resident cannot see that, so the oracle here is the same rollout
+    with compaction switched off.
+    """
+
+    def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]:
+        with monkeypatch.context() as patch:
+            if not compact:
+                patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+            kv, st = make_state(
+                window_chunks=window_chunks,
+                sink_chunks=sink_chunks,
+                reset_at_boundary=reset_at_boundary,
+                chunk_size=RAGGED_CHUNK,
+            )
+            ticks = []
+            for _ in range(12):
+                ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+                ctx.ensure_video_slots(torch.device("cpu"))
+                ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values()))
+                st.commit_paged_context(POS)
+            return ticks, st.adapter(POS).compacted_tokens
+
+    expected, _ = resident_blocks_per_tick(compact=False)
+    actual, compacted_tokens = resident_blocks_per_tick(compact=True)
+    assert compacted_tokens > 0, "the table was never compacted, so nothing was tested"
+    for tick, (want, got) in enumerate(zip(expected, actual)):
+        assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"
+
+
+def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
+    """The staging buffers must fit the table width this module pads to.
+
+    A sink or window that is not a whole number of blocks spans one more block
+    than its token count suggests, so sizing staging from tokens left it short
+    and the first prepared forward refused to run. And because such a chunk
+    shares its first block with the history's tail, the history a probe sees is
+    not left untouched by the next probe, so it is restaged rather than reused.
+    """
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(
+        device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True
+    )
+    capacity = int(kv.history_staging[0][0].shape[0])
+    for chunk in range(6):
+        for probe in range(2):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+            ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+            assert int(ctx.max_seq_len) <= capacity
+            assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history"
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+        st.commit_paged_context(POS)
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
+    """Action K/V follows the video blocks in the table, and the kernel reads it as one run.
+
+    A 24-token chunk from an empty history ends 8 slots into its second 16-token
+    block, so the run would read those 8 unwritten slots as the first action
+    tokens and never reach the last ones. Every shape stays right, which is why
+    this is refused rather than computed.
+    """
+    _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK)
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    with pytest.raises(ValueError, match="partly written video block"):
+        ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))
+
+
+def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
+    """832x480 with the checkpoint's own sink and window, slot by slot.
+
+    A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge
+    but 28080 tokens together -- a whole number of blocks -- so a width taken
+    from their sum was one block short on every tick once the window had
+    filled. Twelve ticks of three frames, checking every slot the kernel reads
+    against the token that slot is supposed to hold. No attention is computed:
+    at this size the bookkeeping is the whole question.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    frame, frames_per_tick = (480 // 16) * (832 // 16), 3
+    block = paging_block_size(frame)
+    seq_len = frames_per_tick * frame
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.float16,
+        block_size=block,
+        max_model_len=1 << 20,
+        available_bytes=1 << 27,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=frames_per_tick,
+        max_scratch_tokens_per_branch=seq_len,
+        device=torch.device("cpu"),
+    )
+    st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1)
+    sink, window = 9 * frame, int(kv.spec.sliding_window)
+    holds: dict[int, int] = {}  # slot -> token position last written there
+    widths: set[int] = set()
+    filled = compacted = False
+
+    for tick in range(12):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + seq_len
+        ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()):
+            holds[slot] = history + offset
+        table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu"))
+        widths.add(int(table.shape[1]))
+        kv_len = int(seq_lens[0])
+
+        position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
+        read = []
+        for b in table[0].tolist()[: -(-kv_len // block)]:
+            assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
+            read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
+        read = read[:kv_len]
+        wrong = [(slot, want) for slot, want in read if holds.get(slot) != want]
+        assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies"
+
+        start = end - window
+        kept = set(range(min(sink, end))) | set(range(max(start, 0), end))
+        missed = kept - {want for _, want in read}
+        assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read"
+
+        filled = filled or end > sink + window
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+
+    assert filled, "the window never filled, so the case that fails was never reached"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+
+
+class _CountedBlockIds:
+    """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read.
+
+    Every path to a request's block ids goes through ``get_blocks``, so a copy of
+    the whole table is counted in full even if only a few entries are indexed.
+    """
+
+    def __init__(self, manager):
+        self.reads = 0
+        self._get_blocks = manager.get_blocks
+
+    def __call__(self, request_id):
+        from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+        blocks = self._get_blocks(request_id)
+        return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks))
+
+
+class _CountedBlock:
+    def __init__(self, block, counter: _CountedBlockIds):
+        self._block = block
+        self._counter = counter
+
+    @property
+    def block_id(self) -> int:
+        self._counter.reads += 1
+        return self._block.block_id
+
+
+def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
+    """However long the block table is, reading the window must not cost more.
+
+    Evicted positions stay in the table as null entries until compaction drops
+    them a whole number of chunks at a time -- up to two frames' worth of blocks
+    at 832x480 -- so the table can still be far longer than the window. Only the
+    sink's blocks and the recent window's blocks can hold a kept token, so those
+    are all the read may touch. Compaction is switched off here so the table
+    grows to many times the window, and reads are counted where the ids leave
+    vLLM's block manager.
+    """
+    monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    reads: list[int] = []
+    table_lengths: list[int] = []
+    for _ in range(60):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        counted = _CountedBlockIds(kv.manager)
+        kv.manager.get_blocks = counted
+        try:
+            ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device)
+        finally:
+            del kv.manager.get_blocks
+        reads.append(counted.reads)
+        table_lengths.append(len(kv.block_table(st.adapter(POS))))
+        st.commit_paged_context(POS)
+
+    assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested"
+    assert max(reads) <= ctx.max_video_blocks, (
+        f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window"
+    )
+
+
+def test_the_checkpoints_own_default_resolution_can_build_a_cache():
+    """832x480 is what LingBot World v2 ships as its default, and it could not run.
+
+    A frame-sized block made the resolution a kernel-compatibility question:
+    1560 tokens per frame is not a multiple of 16, so FlashAttention's paged
+    kernel rejected it and the default resolution had no realtime path at all.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+    assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block"
+
+    block_size = paging_block_size(tokens_per_frame)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=block_size,
+        max_model_len=1 << 16,
+        available_bytes=1 << 28,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=block_size,
+        device=torch.device("cpu"),
+    )
+
+    assert kv.block_size % 16 == 0
+    # The eviction unit is still the frame; only the paging unit changed.
+    assert kv.spec.chunk_size == tokens_per_frame
+    assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98
+
+
+def test_the_shipped_resolution_geometry_matches_dense_attention():
+    """The real number, not a scaled-down stand-in: 1560 tokens per frame.
+
+    This is the case the issue reported failing at ~4e-02 on the scratch path.
+    It runs on CPU because the reference is dense attention, not a kernel.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    # The non-committing path is the one that was wrong: four of the five
+    # forwards per generated block take it.
+    ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=0, query_len=tokens_per_frame, device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    ref = _dense_attention(
+        query,
+        torch.cat([history_k, current_k], dim=1),
+        torch.cat([history_v, current_v], dim=1),
+    )
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
+    """The shipped 832x480 geometry through FlashAttention's paged kernel.
+
+    The CPU test above proves the addressing is right against a dense
+    reference. This one proves the real kernel agrees, at the real 1560
+    tokens per frame, through the production path -- host prep followed by
+    the fused write+attend op, which is what a forward actually calls.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+    assert kv.block_size == 16
+
+    inputs = layer_ctx.to_layer_inputs()
+    paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
@pytest.mark.parametrize("history_chunks", [0, 1, 3])
@pytest.mark.parametrize("action_len", [0, 3])
@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_validate_fa4_head_dims = fa.validate_fa4_head_dims
+
+
+@pytest.fixture(autouse=True)
+def kernel_validator(monkeypatch):
+    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
+    # validation is exercised by test_flash_attn_compile.py.
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+
+def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
+    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
+    impl._kernel_variant = kernel_variant
+    impl.causal = causal
+    return impl
+
+
+def _resolve(
+    impl,
+    *,
+    context: ExecutionContext | None = None,
+    attn_metadata: AttentionMetadata | None = None,
+    dtype=torch.bfloat16,
+):
+    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
+    return impl.resolve_execution_path(
+        context or ExecutionContext(platform="cuda"),
+        tensor,
+        tensor,
+        tensor,
+        attn_metadata,
+    )
+
+
+class _UnmigratedBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "UNMIGRATED_TEST"
+
+    @staticmethod
+    def get_impl_cls():
+        return None
+
+    @staticmethod
+    def get_metadata_cls():
+        return None
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+def test_capability_result_is_tri_state():
+    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
+    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
+    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
+    with pytest.raises(ValueError, match="actionable reason"):
+        CapabilityResult(SupportStatus.UNSUPPORTED)
+
+
+def test_unmigrated_backend_default_is_conservative():
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _UnmigratedBackend.resolve_capabilities(context)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    requested = result.requested_support(context)
+    assert requested.status is SupportStatus.UNSUPPORTED
+    assert "no verified fullgraph declaration" in requested.reason
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("mode", list(CompilationMode))
+def test_fullgraph_request_checks_compilation_mode(mode):
+    result = replace(_resolve(_impl()), compilation_mode=mode)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
+    assert result.requested_support(context).status is expected
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED
+
+
+def test_request_preserves_verified_rejection_reason():
+    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
+    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support
+
+
+def test_exact_cuda_fa4_dense_path_is_supported():
+    context = ExecutionContext(
+        platform="cuda",
+        require_fullgraph=True,
+    )
+    result = _resolve(_impl(), context=context)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+
+
+def test_causal_fa4_path_remains_unmigrated():
+    result = _resolve(_impl(causal=True))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
+    head_dims = (80, 48)
+    validator = Mock(return_value=True)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
+    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
+    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
+    assert result.support.status is SupportStatus.SUPPORTED
+    validator.assert_called_once_with(*head_dims, 8)
+
+
+def test_fa4_reports_kernel_rejection(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "kernel dimension constraint" in result.support.reason
+    assert "another backend" in result.support.reason
+
+
+def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_fa4_rejects_mismatched_query_key_dimensions():
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "Q and K head dimensions must match" in result.support.reason
+
+
+@pytest.mark.parametrize("arch", [90, 100, 120])
+def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
+    validator = Mock()
+    monkeypatch.setitem(
+        sys.modules,
+        "flash_attn.cute.interface",
+        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
+    )
+    # Call the real adapter, not the fixture's stand-in.
+    verified = _validate_fa4_head_dims(80, 48, 8)
+    if arch == 120:
+        assert not verified
+        validator.assert_not_called()
+    else:
+        assert verified
+        validator.assert_called_once_with(80, 48, arch // 10, 8)
+
+
+def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
+    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
+    assert not _validate_fa4_head_dims(64, 64, 8)
+
+
+def test_initialized_kernel_identity_overrides_caller_claim():
+    result = _resolve(
+        _impl(kernel_variant=None),
+        context=ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+        ),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
+    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.path == "unverified"
+
+
+@pytest.mark.parametrize(
+    "changes",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_unverified_flash_attention_variants_remain_unmigrated(changes):
+    result = _resolve(
+        _impl(),
+        context=replace(ExecutionContext(platform="cuda"), **changes),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_unverified_dtype_remains_unmigrated():
+    result = _resolve(_impl(), dtype=torch.float32)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize(
+    "attn_metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
+        AttentionMetadata(
+            extra={
+                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
+                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
+                "max_seqlen_q": 16,
+                "max_seqlen_k": 16,
+            }
+        ),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_metadata_normalization_prevents_dense_support(attn_metadata):
+    result = _resolve(_impl(), attn_metadata=attn_metadata)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_unpublished_all_true_mask_remains_runtime_dependent():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "runtime_mask_dependent"
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_producer_published_noop_mask_resolves_dense_fa4():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+
+
+def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    impl = _impl()
+    dense = _resolve(impl, attn_metadata=metadata)
+    assert dense.support.status is SupportStatus.SUPPORTED
+
+    metadata.attn_mask[:, 8:] = False
+    metadata.extra["attention_mask_mode"] = "padding"
+    masked = _resolve(impl, attn_metadata=metadata)
+    assert masked.support.status is SupportStatus.UNMIGRATED
+    request = ExecutionContext(platform="cuda", require_fullgraph=True)
+    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED
+
+
+def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
+    from vllm_omni.diffusion.attention.backends import flash_attn
+    from vllm_omni.diffusion.attention.backends.utils import fa
+
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
+    impl = _impl()
+    impl.fa_deterministic = False
+    impl.softmax_scale = 0.125
+    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED
+
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
+    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
+    assert impl.forward_cuda(query, query, query, metadata) is query
+
+
+@pytest.mark.parametrize("input_index", [1, 2])
+@pytest.mark.parametrize("mismatch", ["dtype", "device"])
+def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
+    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
+    tensors[input_index] = (
+        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
+    )
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert f"{mismatch}s must match" in result.support.reason
+
+
+@pytest.mark.parametrize("skip_quant", [False, True])
+def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+
+    layer = Attention.__new__(Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = _impl(kernel_variant=None)
+    layer._hsdp_compile_boundary_enabled = False
+    layer.skip_sequence_parallel = False
+    layer._no_parallel_strategy = NoParallelAttention()
+    layer.parallel_strategy = layer._no_parallel_strategy
+    layer.use_ring = False
+    layer.paged_kv_cache_role = None
+    layer._kv_cache_dtype = "fp8"
+    layer._disable_kv_quant = False
+    layer._kv_cache_skip_steps = None
+    layer._kv_cache_skip_layers = {0} if skip_quant else None
+    layer.layer_idx = 0
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
+    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
+    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert metadata.extra == {"kv_cache_dtype": "fp8"}
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Stateful contract example, independent of an optional attention kernel.
+
+The prepared module owns tensor state and passes it explicitly to an opaque op.
+This deliberately does not model opaque Python/C++ planning handles.
+"""
+
+import gc
+import weakref
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=())
+def _planned_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: torch.Tensor,
+) -> torch.Tensor:
+    return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous()
+
+
+@_planned_attention.register_fake
+def _planned_attention_fake(query, key, value, query_scale):
+    return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+class _PreparedAttention(torch.nn.Module):
+    """Initialization/planning runs eagerly; forward only executes tensors."""
+
+    def __init__(self, query_scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("query_scale", query_scale.detach().clone())
+
+    def forward(self, query, key, value):
+        return _planned_attention(query, key, value, self.query_scale)
+
+
+def _inputs(length):
+    generator = torch.Generator().manual_seed(length)
+    return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3))
+
+
+def _reference(plan, inputs):
+    query, key, value = inputs
+    return F.scaled_dot_product_attention(query * plan.query_scale, key, value)
+
+
+def test_dynamic_replay_uses_each_instances_state_without_recompilation():
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    # Same tensor metadata, different values: state must be a graph input,
+    # never a Python instance identifier or a captured first-instance constant.
+    plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)]
+    compiled = torch.compile(
+        lambda plan, q, k, v: plan(q, k, v),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    for length in (5, 9, 13):
+        inputs = _inputs(length)
+        for plan in plans:
+            torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs))
+        assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs))
+    assert compile_count == 1
+
+
+def test_compiled_callable_retains_state_until_released():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    plan_ref = weakref.ref(plan)
+    state_ref = weakref.ref(plan.query_scale)
+    inputs = _inputs(5)
+    expected = _reference(plan, inputs)
+    compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True)
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    # Deleting the caller's reference must not invalidate subsequent replay.
+    del plan
+    gc.collect()
+    assert plan_ref() is not None
+    assert state_ref() is not None
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    del compiled
+    gc.collect()
+    assert plan_ref() is None
+    assert state_ref() is None
+
+
+def test_stateful_custom_op_schema_and_fake_implementation():
+    inputs = _inputs(5)
+    # Unequal Q/K lengths and V head dimensions exercise the fake output shape.
+    query, _, _ = inputs
+    key = torch.randn(2, 3, 7, 8)
+    value = torch.randn(2, 3, 7, 4)
+    scale = torch.full((1, 3, 1, 1), 0.5)
+    torch.library.opcheck(_planned_attention, (query, key, value, scale))
+
+
+def test_stateful_attention_compiles_with_inductor():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    compiled = torch.compile(plan, fullgraph=True, dynamic=True)
+    for length in (5, 9):
+        inputs = _inputs(length)
+        torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs))
@@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
+from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus
 from vllm_omni.platforms import current_omni_platform

 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch):
    assert cu_seqlens_k.tolist() == [0, 2, 5]


+@pytest.mark.cpu
+def test_published_noop_mask_uses_dense_path(monkeypatch):
+    calls = []
+
+    def fake_dense_func(q, k, v, **_kwargs):
+        calls.append("dense")
+        return q
+
+    def fail_varlen(*_args, **_kwargs):
+        raise AssertionError("published no-op mask must not use varlen attention")
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 4, 2, 8)
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 4), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert calls == ["dense"]
+
+
@pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
    calls = []
@@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
    return fake


+@pytest.fixture
+def npu_contract(monkeypatch, mocker):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4"))
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0)))
+    _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed)
+    return _npu_impl(), query, dense, packed
+
+
+def _assert_platform_path(impl, platform, query, metadata, path):
+    context = ExecutionContext(platform=platform, require_fullgraph=True)
+    result = impl.resolve_execution_path(context, query, query, query, metadata)
+    assert result.path == path
+    assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform]
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+def _real_pad_metadata():
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    return AttentionMetadata(
+        extra={
+            "npu_attn_varlen": True,
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 5,
+            "max_seqlen_k": 5,
+        }
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"])
+def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_dense")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_mask_layout(npu_contract):
+    impl, query, dense, packed = npu_contract
+    mask = torch.arange(8)[None] < 5
+    metadata = AttentionMetadata(attn_mask=mask)
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8))
+
+
+@pytest.mark.cpu
+def test_npu_contract_packed_dispatch(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8])
+
+
+@pytest.mark.cpu
+def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    # The same metadata resolves differently after the kernel choice changes.
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.args[1], query[:, :5])
+    torch.testing.assert_close(dense.call_args.args[2], query[:, :5])
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract):
+    impl, query, dense, packed = npu_contract
+    # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask.
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8)
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask)
+
+
+@pytest.mark.cpu
+def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    with pytest.raises(ValueError, match="refusing to run unmasked attention"):
+        impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_not_called()
+
+
+@pytest.fixture
+def rocm_contract(monkeypatch, mocker):
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"]))
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", True)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", dense)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", packed)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4"))
+    return _npu_impl(), query, dense, packed
+
+
+@pytest.mark.cpu
+def test_rocm_contract_dense_dispatch(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    _assert_platform_path(impl, "rocm", query, None, "rocm_dense")
+    impl.forward_hip(query, query, query)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_masked_unpadding(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"})
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[0].shape == (5, 2, 4)
+    torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32))
+
+
+@pytest.mark.cpu
+def test_rocm_contract_multiple_packed_documents(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 4,
+            "max_seqlen_k": 4,
+        }
+    )
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen")
+    # NPU's [real, pad] path cannot consume this packing contract.
+    assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.kwargs["cu_seqlens_q"] is cu
+    assert packed.call_args.kwargs["cu_seqlens_k"] is cu
+
+
+@pytest.mark.cpu
+def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    # Resolution may inspect shape/metadata, but must not read mask values.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values"))
+        _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_requires_initialized_aiter(monkeypatch):
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", False)
+    impl = _npu_impl()
+    query = torch.zeros(1, 8, 2, 4)
+    result = impl.resolve_execution_path(
+        ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None
+    )
+    assert result.path == "rocm_unverified"
+    assert result.kernel_variant is None
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------


@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+@pytest.fixture(autouse=True)
+def isolated_compiler_cache():
+    # Each case changes kernel/configuration on the same Attention.forward code
+    # object. Keep unrelated cases from exhausting its Dynamo recompile limit;
+    # retain the cache across all shapes and replays within an individual test.
+    torch.compiler.reset()
+    yield
+    torch.compiler.reset()
+
+
+@dataclass
+class _ParallelStrategyStub:
+    name: str
+    enabled: bool = True
+
+
+def _make_attention_layer(monkeypatch, *, head_size=64):
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    monkeypatch.setattr(
+        attention_layer,
+        "get_attn_backend_for_role",
+        lambda **_kwargs: (FlashAttentionBackend, None),
+    )
+    monkeypatch.setattr(
+        attention_layer,
+        "build_parallel_attention_strategy",
+        lambda **_kwargs: NoParallelAttention(),
+    )
+    return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False)
+
+
+@pytest.fixture
+def fake_fa4(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded", encoding="utf-8")
+
+    @functools.cache
+    def cached_kernel_loader():
+        with open(marker, encoding="utf-8") as handle:
+            handle.read()
+
+    def fake_attention(query, key, value, **kwargs):
+        cached_kernel_loader()
+        return (
+            torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                key.transpose(1, 2),
+                value.transpose(1, 2),
+                scale=kwargs["softmax_scale"],
+                is_causal=kwargs["causal"],
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_attention)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous):
+    """Fast Dynamo/custom-op regression; Inductor coverage is separate."""
+    impl = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q = q.transpose(1, 2).contiguous().transpose(1, 2)
+        assert not q.is_contiguous()
+    path = impl.resolve_execution_path(
+        ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+            dtype="bfloat16",
+            require_fullgraph=True,
+        ),
+        q,
+        q,
+        q,
+        None,
+    )
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    compiled = torch.compile(
+        lambda query, key, value: impl.forward_cuda(query, key, value),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    out = compiled(q, q, q)
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q2 = q2.transpose(1, 2).contiguous().transpose(1, 2)
+    out2 = compiled(q2, q2, q2)
+    assert out.is_contiguous()
+    assert out2.is_contiguous()
+
+    assert out.shape == q.shape
+    assert out2.shape == q2.shape
+    torch.testing.assert_close(out, impl.forward_cuda(q, q, q))
+    torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2))
+    assert compile_count == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fa4_production_attention_entry_compiles_with_inductor(
+    monkeypatch,
+    fake_fa4,
+):
+    """Compile through Attention.forward with the default Inductor backend."""
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    layer = _make_attention_layer(monkeypatch)
+    layer.attention = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+
+    resolved_contexts = []
+    resolve_execution_path = fake_fa4.resolve_execution_path
+
+    def record_execution_path(context, query, key, value, attn_metadata):
+        resolved_contexts.append(context)
+        return resolve_execution_path(
+            context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
+    monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path)
+    context = ExecutionContext(
+        platform="cuda",
+        kernel_variant="fa4",
+        dtype="bfloat16",
+        paged_kv=True,
+        parallel_strategy=ParallelStrategy.ULYSSES,
+        require_fullgraph=True,
+    )
+    path = layer.resolve_execution_path(context, q, q, q, None)
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+    assert resolved_contexts[-1].paged_kv is False
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    out = compiled(q, q, q)
+
+    assert out.shape == q.shape
+    torch.testing.assert_close(out, layer(q, q, q))
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2))
+
+    layer.paged_kv_cache_role = "self"
+    monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object())
+    paged_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].paged_kv is True
+    assert paged_path.support.status is SupportStatus.UNMIGRATED
+    layer.paged_kv_cache_role = None
+
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ulysses")
+    hybrid_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING
+    assert hybrid_path.support.status is SupportStatus.UNMIGRATED
+    layer.use_ring = False
+    layer.parallel_strategy = NoParallelAttention()
+
+    layer._hsdp_compile_boundary_enabled = True
+    hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert hsdp_path.support.status is SupportStatus.UNMIGRATED
+    assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY
+
+    layer._hsdp_compile_boundary_enabled = False
+    runtime_context = ForwardContext(
+        omni_diffusion_config=OmniDiffusionConfig(
+            parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1),
+        ),
+    )
+    monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context)
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ring")
+    runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+    assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)])
+def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims):
+    """Validate real FA4 numerics through the complete compiled attention layer."""
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major < 10:
+        pytest.skip("Requires Blackwell")
+    if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None:
+        pytest.skip("Requires CuTe FlashAttention-4")
+
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    # Tile boundaries, cross-attention lengths, batching, and a longer sequence.
+    for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)):
+        query, key, value = (
+            torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator)
+            for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size))
+        )
+        path = layer.resolve_execution_path(context, query, key, value, None)
+        assert path.requested_support(context).status is SupportStatus.SUPPORTED
+        with sdpa_kernel(SDPBackend.MATH):
+            reference = torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2).float(),
+                key.transpose(1, 2).float(),
+                value.transpose(1, 2).float(),
+                scale=head_size**-0.5,
+            ).transpose(1, 2)
+        actual = compiled(query, key, value)
+        assert actual.dtype == query.dtype
+        assert actual.device == query.device
+        torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(actual, layer(query, key, value))
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)])
+def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims):
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4")
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    path = layer.resolve_execution_path(context, query, query, value, None)
+    assert path.requested_support(context).status is SupportStatus.UNSUPPORTED
+    with pytest.raises(AssertionError) as error:
+        layer(query, query, value)
+    assert str(error.value) in path.support.reason
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous):
+    if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires CuTe FlashAttention-4")
+    query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value))
+        assert not query.is_contiguous()
+    torch.library.opcheck(
+        torch.ops.vllm_omni.fa4_dense_attention.default,
+        (query, key, value, 80**-0.5, False, False),
+    )
@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP."""
+
+from contextlib import ExitStack
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu]
+
+
+class _SDPAAttention(nn.Module):
+    """Use deterministic native attention while testing loading and FSDP."""
+
+    def __init__(self, *, causal, softmax_scale, **kwargs):
+        super().__init__()
+        self.causal = causal
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value):
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge):
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+
+    from safetensors.torch import save_file
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+    from vllm.config.load import LoadConfig
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.diffusion.distributed import hsdp
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.platforms import current_omni_platform
+
+    device = torch.device("cuda", rank)
+    current_omni_platform.set_device(device)
+    dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120))
+    try:
+        with ExitStack() as stack:
+            stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))))
+            stack.enter_context(set_default_torch_dtype(torch.bfloat16))
+            # HSDP spans both workers, while every vLLM linear uses TP=1.
+            for module in (linear, parameter):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0))
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+            for module in (transformer_cosmos3, transformer_cosmos3_edge):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+                stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention))
+            stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None)))
+            stack.enter_context(
+                patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank))
+            )
+
+            model_cls = (
+                transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer
+                if edge
+                else transformer_cosmos3.Cosmos3VFMTransformer
+            )
+            tf_config = dict(
+                hidden_size=16,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=8,
+                intermediate_size=32,
+                vocab_size=32,
+                latent_patch_size=2,
+                latent_channel=48,
+                temporal_compression_factor=4,
+                rope_scaling={"mrope_section": [2, 1, 1]},
+            )
+            if edge:
+                tf_config.update(
+                    backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE,
+                    qk_norm_for_text=False,
+                    use_und_k_norm_for_gen=True,
+                )
+            config = SimpleNamespace(
+                tf_model_config=tf_config,
+                dtype=torch.bfloat16,
+                quantization_config=None,
+                lora_path=None,
+                num_weight_load_threads=1,
+                hsdp_weight_load_strategy="full",
+                parallel_config=SimpleNamespace(
+                    use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1
+                ),
+            )
+
+            class Pipeline(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.transformer = model_cls(config)
+                    self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+                    self.vae = nn.Linear(2, 2)
+                    self.weights_sources = [
+                        DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False)
+                    ]
+
+                def load_weights(self, weights):
+                    state = dict(weights)
+                    self.transformer.load_state_dict(
+                        {name.removeprefix("transformer."): t for name, t in state.items()}
+                    )
+                    self.transformer.post_load_weights()
+                    self.transformer.eval()
+                    self.transformer.validate_loaded_weights(set(state))
+                    return set(state)
+
+            full = Pipeline()
+            if rank == 0:
+                torch.manual_seed(42)
+                with torch.no_grad():
+                    for name, tensor in full.transformer.named_parameters():
+                        tensor.copy_(torch.randn_like(tensor) * 0.02)
+                        if "norm" in name:
+                            tensor.fill_(1)
+                save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors"))
+            dist.barrier()
+
+            def load(model):
+                loader = DiffusersPipelineLoader(LoadConfig(), config)
+                with patch.object(loader, "_init_from_load_format", return_value=model):
+                    return loader._load_model_with_hsdp(device)
+
+            full = load(full)
+            config.hsdp_weight_load_strategy = "pre_sharded"
+            pre_sharded = load(Pipeline())
+            full_parameters = dict(full.transformer.named_parameters())
+            from torch.distributed.tensor import DTensor
+
+            for name, actual in pre_sharded.transformer.named_parameters():
+                expected = full_parameters[name]
+                if name.startswith("time_embedder."):
+                    assert not isinstance(actual, DTensor)
+                    assert actual.dtype == torch.float32
+                else:
+                    assert isinstance(actual, DTensor)
+                    actual = actual.full_tensor()
+                    expected = expected.full_tensor()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            assert pre_sharded.vae.weight.device == device
+            torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch)
+            assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers())
+
+            inputs = dict(
+                hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2),
+                timestep=torch.tensor([1.0], device=device),
+                text_ids=torch.tensor([[1, 2]], device=device),
+                text_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                video_shape=(1, 2, 2),
+                fps=24.0,
+            )
+            with torch.no_grad():
+                # Also exercise the cached UND path on the second denoising step.
+                for step in (1.0, 0.5):
+                    inputs["timestep"].fill_(step)
+                    expected = full.transformer(**inputs)
+                    actual = pre_sharded.transformer(**inputs)
+                    assert torch.isfinite(actual).all()
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(),
+    reason="requires two CUDA GPUs and NCCL",
+)
+def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge):
+    torch.multiprocessing.spawn(
+        _load_cosmos_worker,
+        args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge),
+        nprocs=2,
+        join=True,
+    )
@@ -3,11 +3,13 @@
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""

 import gc
+from types import SimpleNamespace

 import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DeviceMesh, DTensor
 from vllm.utils.network_utils import get_file_store_init_method

@@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import (
    _unshardable_parameters,
    shard_model,
 )
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader

 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]

@@ -338,3 +341,92 @@ class TestHSDPShardConditions:
                matched.append(name)
        assert "blocks.0" in matched
        assert "blocks.1" in matched
+
+
+def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group):
+    with torch.device("meta"):
+        model = nn.Module()
+        model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False))
+    model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False)
+    shard_model(
+        model.transformer,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "0"],
+    )
+
+    load_plan = SimpleNamespace(
+        roots=[SimpleNamespace(name="transformer", module=model.transformer)],
+        groups=[SimpleNamespace(module=model.transformer[0])],
+        binding_names=frozenset({"transformer.0.weight"}),
+    )
+    sharded_weight = model.transformer[0].weight
+
+    DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu"))
+
+    fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0]
+    assert model.transformer[0].weight is sharded_weight
+    assert model.transformer[0].weight.to_local().device.type == "cpu"
+    assert fsdp_param._sharded_param_data.device.type == "cpu"
+    assert (
+        fsdp_param._sharded_param_data.untyped_storage().data_ptr()
+        == model.transformer[0].weight.to_local().untyped_storage().data_ptr()
+    )
+    assert torch.equal(model.transformer.scratch, torch.tensor([3.0]))
+
+
+def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch):
+    from safetensors.torch import save_file
+    from vllm.config.load import LoadConfig
+
+    from vllm_omni.diffusion.distributed import hsdp as hsdp_module
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    class Transformer(nn.Module):
+        _hsdp_shard_conditions = [lambda name, module: name == "block"]
+        _hsdp_ignored_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.block = nn.Linear(3, 4, bias=False)
+            self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16)
+            self.root_weight = nn.Parameter(torch.zeros(2))
+            self.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+            release_module_parameters_to_meta(self.block)
+
+        def post_load_weights(self):
+            self.time_embedder.to(torch.float32)
+
+    model = nn.Module()
+    model.transformer = Transformer()
+    model.vae = nn.Linear(2, 2)
+    tensors = {
+        "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3),
+        "root_weight": torch.tensor([2.0, 3.0]),
+        "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(tensors, str(tmp_path / "model.safetensors"))
+    model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)]
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization_config=None,
+        lora_path=None,
+        hsdp_weight_load_strategy="pre_sharded",
+        num_weight_load_threads=1,
+        parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model)
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0))
+    monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+    loaded = loader._load_model_with_hsdp(torch.device("cpu"))
+    assert loaded is model
+    for name, expected in tensors.items():
+        actual = dict(model.transformer.named_parameters())[name]
+        if name.startswith("time_embedder."):
+            assert not isinstance(actual, DTensor)
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected.float())
+        else:
+            assert isinstance(actual, DTensor)
+            torch.testing.assert_close(actual.to_local(), expected)
+    torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0]))
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
    HostWeightPlan,
    HostWeightPlanResult,
    TensorBinding,
+    build_checkpoint_binding_plan,
 )
 from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
@@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
            )
        assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
        assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))
+
+
+def test_pre_sharded_hsdp_rejects_quantization_before_loading():
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=object(),
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    with pytest.raises(ValueError, match="does not support quantization"):
+        loader._load_model_with_hsdp(torch.device("cpu"))
+
+
+def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    discovered = PipelineModules(
+        dits=[model.transformer],
+        dit_names=["transformer"],
+        vaes=[],
+        encoders=[],
+        encoder_names=[],
+        resident_modules=[],
+        resident_names=[],
+    )
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
+    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
+    pre_sharded_load = mocker.patch.object(
+        loader,
+        "_load_model_with_pre_sharded_hsdp",
+        return_value=model,
+    )
+
+    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
+    pre_sharded_load.assert_called_once()
+
+
+def test_pre_sharded_hsdp_rejects_tensor_transform():
+    def transform(tensor):
+        return tensor.transpose(-2, -1)
+
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+                transform=transform,
+            )
+        }
+    )
+
+    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
+        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+            )
+        }
+    )
+
+    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
+    target = torch.empty(8, 4, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+
+    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": target},
+    )
+    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]
+
+    assert checkpoint_target is target
+    assert checkpoint_target.shape == (8, 4, 3)
+    assert checkpoint_target.device.type == "cpu"
+    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
+    assert local_bytes == target.numel() * target.element_size()
+
+
+def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
+    runtime_target = torch.empty(2, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": runtime_target},
+    )
+    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)
+
+    assert torch.equal(runtime_target, checkpoint_tensor)
+
+
+def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
+    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.first": binding,
+            "transformer.second": binding,
+        }
+    )
+
+    with pytest.raises(ValueError, match="multiple runtime tensors"):
+        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
+        )
+
+
+def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
+    root = nn.Module()
+    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+
+    names = {module: name or "root" for name, module in root.named_modules()}
+    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+        root,
+        nested_hsdp_roots={root.left},
+    )
+
+    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]
+
+
+@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
+def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
+    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod
+
+    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)
+
+    class Pipeline(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+
+        @staticmethod
+        def remap_checkpoint_key(name):
+            return "transformer.weight" if name == "transformer.renamed" else name
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
+    source = SimpleNamespace(
+        model_or_path=str(tmp_path),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    model = Pipeline()
+    result = build_checkpoint_binding_plan(
+        model,
+        dit_modules=(("transformer", model.transformer),),
+        sources=(source,),
+        model_path=None,
+        tensor_parallel_size=1,
+        online_quantization=False,
+    )
+
+    if checkpoint_key == "renamed":
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
+    else:
+        assert result.plan is None
+        assert "no checkpoint binding" in result.fallback_reason
@@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa

    # The conditioning is prepared once per request and reused by every step.
    assert prepare_spy.call_count == 1
-    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+    assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
+    # The time grid reaches prepare, so the adaLN modulations are computed once.
+    assert prepare_spy.call_args.args[7].numel() == 2
    assert step_spy.call_count == 2
    capture_spy.assert_not_called()
    assert not wrapper._cache
@@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
    assert padded[4].shape == (1, 100, 4)
    assert padded[5].shape == (1, 100)
    assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
+    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
    assert wrapper.max_graphs == 32


@@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
            inputs["ref"],
            inputs["ref_mask"],
        )
-        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        # sample_latents passes the grid, so the key carries its step count (two steps here).
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
        assert key in wrapper._cache

        # The static context was refreshed for this request's conditioning.
        entry = wrapper._cache[key]
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        grid = torch.tensor([0.0, 0.4], device="cuda")
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
        for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
            if want is None:
                assert static is None
            else:
                torch.testing.assert_close(static, want)
        torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+        assert int(entry.static_step) == 1

    assert len(wrapper._cache) == 1

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT fast paths keep the reference numerics.
+
+Two paths are checked: the grouped conv position embedding as one batched
+GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
+under the per-step CUDA graph.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk import auk_transformer as T
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str = "cpu") -> AuKTransformer:
+    torch.manual_seed(5)
+    return (
+        AuKTransformer(
+            dim=32,
+            heads=2,
+            dim_head=16,
+            ff_mult=2,
+            latent_dim=4,
+            text_hidden_dim=8,
+            num_layers=2,
+            num_single_layers=2,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+    """The channels-first conv stack as the reference implementation writes it."""
+    keep = None if mask is None else mask.unsqueeze(1)
+    x = x.transpose(1, 2)
+    if keep is not None:
+        x = x.masked_fill(~keep, 0.0)
+    for layer in module.conv1d:
+        x = layer(x)
+        if keep is not None and isinstance(layer, torch.nn.Conv1d):
+            x = x.masked_fill(~keep, 0.0)
+    return x.transpose(1, 2)
+
+
+@torch.inference_mode()
+def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
+    torch.manual_seed(0)
+    conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
+    x = torch.randn(2, 37, 64)
+    expected = conv(x.transpose(1, 2)).transpose(1, 2)
+    torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("masked", [False, True])
+def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
+    torch.manual_seed(1)
+    module = T.ConvPosEmbedding(64).eval()
+    x = torch.randn(2, 37, 64)
+    mask = None
+    if masked:
+        mask = torch.ones(2, 37, dtype=torch.bool)
+        mask[1, 30:] = False
+    torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_infer", [False, True])
+def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
+    dit = _make_dit()
+    text = torch.randn(1, 5, 8)
+    ref = torch.randn(1, 3, 4)
+    x = torch.randn(1, 6, 4)
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
+    plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
+    tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
+    assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
+    for i in range(4):
+        expected = dit.step(x, grid[i], plain)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
+        index = torch.tensor(i)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
+
+
+def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(3)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    c_mask[:, 5:] = False
+    text[:, 5:] = 0.0
+    ref = torch.randn(1, 6, 4, generator=generator)
+    ref_mask = torch.ones(1, 6, dtype=torch.bool)
+    ref_mask[:, 4:] = False
+    mask = torch.ones(1, 9, dtype=torch.bool)
+    mask[:, 7:] = False
+    inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
+    return {name: value.to(device) for name, value in inputs.items()}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    inputs = _padded_inputs("cuda")
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
+    x = torch.randn(1, 9, 4, device="cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    common = dict(
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=inputs["ref_mask"],
+        cfg_strength=cfg_strength,
+        timesteps=grid[:-1],
+    )
+    for i in range(4):
+        graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
+        ctx = dit.prepare(
+            inputs["text"],
+            target_len=9,
+            c_mask=inputs["c_mask"],
+            ref=inputs["ref"],
+            ref_mask=inputs["ref_mask"],
+            cfg_infer=cfg_strength >= 1e-5,
+            timesteps=grid[:-1],
+        )
+        eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
+        torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
+    # One graph serves every step of the grid.
+    assert len(wrapper._cache) == 1
+    assert next(iter(wrapper._cache))[-1] == 4
@@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
            v = dit(x, time=timesteps[i], **inputs)
        x = x + (timesteps[i + 1] - timesteps[i]) * v

-    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    # sample_latents precomputes the adaLN modulations of the whole grid in one
+    # GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
+    torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
    assert dit.text_cond is None and dit.text_uncond is None


@@ -12,7 +12,8 @@ import pytest
 import torch

 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
+from vllm_omni.diffusion.models.auk import auk_vae
+from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
 from vllm_omni.platforms import current_omni_platform

@@ -246,6 +247,72 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
        AuKVAEDecodeGraph(vae, tile_frames=60)


+@pytest.mark.cpu
+@torch.inference_mode()
+def test_fused_activation_is_skipped_off_cuda() -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
+    x = torch.randn(1, 4, 16)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@torch.inference_mode()
+def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
+    x = torch.randn(1, 4, 16, device="cuda")
+    assert activation._fused_eligible(x)
+    monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+def _activation_input(layout: str, length: int) -> torch.Tensor:
+    """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
+    if layout == "contiguous":
+        x = torch.randn(2, 6, length, device="cuda")
+    elif layout == "transposed":
+        x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
+    elif layout == "sliced":
+        x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
+    else:
+        raise ValueError(layout)
+    assert x.shape == (2, 6, length)
+    if length > 1:  # contiguity of a single-sample view depends on the layout
+        assert x.is_contiguous() == (layout == "contiguous")
+    return x * 3
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
+@pytest.mark.parametrize("causal", [False, True])
+# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
+# signal crosses a block boundary at 512 input samples and the decimated one at
+# 1024, so both sides of each boundary are covered, plus short and odd lengths.
+@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
+@torch.inference_mode()
+def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
+    torch.manual_seed(0)
+    activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
+    with torch.no_grad():
+        activation.act.alpha.normal_(0.0, 0.3)
+        activation.act.beta.normal_(0.0, 0.3)
+    x = _activation_input(layout, length)
+    original = x.clone()
+
+    eager = activation.downsample(activation.act(activation.upsample(x)))
+    assert activation._fused_eligible(x)
+    fused = activation(x)
+
+    # Same fp32 arithmetic, only the order in which the taps are summed differs.
+    assert fused.shape == eager.shape
+    torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
+    # Neither path writes into its (possibly strided) input.
+    torch.testing.assert_close(x, original, atol=0, rtol=0)
+
+
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
@torch.inference_mode()
@@ -291,6 +291,18 @@ class TestRequestParsing:
        with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
            build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})

+    @pytest.mark.parametrize(
+        "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
+    )
+    def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        assert pipeline._ref_cache_size == expected
+
+    @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
+    def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
+        with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
+            build_pipeline(model_config={"auk_ref_cache_size": size})
+
    def test_pipeline_declares_audio_output(self, build_pipeline):
        pipeline, _ = build_pipeline()

@@ -492,6 +504,40 @@ class TestRequestParsing:
        assert call["sample"] is True
        assert isinstance(call["generator"], torch.Generator)

+    def test_reference_latents_are_cached_by_content(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+        clip = _silence(1.0)
+
+        for _ in range(3):
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        # A different clip is a new entry, the first one still hits.
+        other = (np.full_like(clip[0], 0.25), clip[1])
+        pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
+        pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert len(pipeline.vae.encode_calls) == 2
+        assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
+
+    def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        for _ in range(2):
+            pipeline.forward(
+                _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
+            )
+        assert len(pipeline.vae.encode_calls) == 2
+
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
+        for _ in range(2):
+            pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline.vae.encode_calls) == 2
+
+    def test_reference_cache_is_bounded(self, build_pipeline):
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
+        for level in (0.1, 0.2, 0.3):
+            clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline._ref_cache) == 2
+
    def test_one_request_per_forward(self, build_pipeline):
        pipeline, _ = build_pipeline()
        batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)
@@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None:
    assert hasattr(model.gen_layers[0].mlp, "gate_proj")


+@pytest.mark.parametrize("edge", [False, True])
+@pytest.mark.parametrize(
+    "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
+)
+def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    released = []
+
+    def release_and_check(block):
+        parameters = dict(block.named_parameters())
+        metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
+        release_module_parameters_to_meta(block)
+        for name, parameter in block.named_parameters():
+            assert parameter is parameters[name]
+            assert parameter.device.type == "meta"
+            assert parameter.__dict__ == metadata[name]
+        released.append(block)
+
+    monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
+    monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
+    model_cls = (
+        transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
+    )
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+            hsdp_weight_load_strategy=strategy,
+        )
+    )
+    blocks = [*model.language_model.layers, *model.gen_layers]
+    expected_meta = use_hsdp and strategy == "pre_sharded"
+    assert len(released) == (len(blocks) if expected_meta else 0)
+    assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
+    assert model.time_embedder.linear_1.weight.device.type != "meta"
+    assert all(buffer.device.type != "meta" for buffer in model.buffers())
+
+
+@pytest.mark.parametrize("edge", [False, True])
+def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
+
+    model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=True),
+            hsdp_weight_load_strategy="pre_sharded",
+        )
+    )
+    groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
+    plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
+    processed = []
+    # Exercise model selection with real vLLM linears, without invoking the
+    # CPU-specific weight packing hook on meta tensors in CPU-only CI.
+    monkeypatch.setattr(
+        UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
+    )
+    DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
+    assert expected
+    assert processed == expected
+
+    model._hsdp_pre_sharded_meta_post_load = False
+    with pytest.raises(ValueError, match="has not declared"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    model._hsdp_pre_sharded_meta_post_load = True
+
+    class UnsupportedMethod(UnquantizedLinearMethod):
+        pass
+
+    expected[-1].quant_method = UnsupportedMethod()
+    processed.clear()
+    with pytest.raises(ValueError, match="UnsupportedMethod"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    assert not processed
+
+
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3

@@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
    assert array.shape == (2, 4, 4, 3)
    video = post({"video": frames}, SimpleNamespace(output_type="pt"))
    assert set(video) == {"video"}
-    assert video["video"] is frames
+    assert isinstance(video["video"], np.ndarray)
+    np.testing.assert_array_equal(video["video"], frames.numpy())
    image = post(
        {"image": torch.ones(4, 4, 3)},
        SimpleNamespace(output_type="pt"),
@@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
    assert latent_output["image"] is latent


+@pytest.mark.parametrize("output_type", [None, "pt", "np"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+    from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames
+
+    # Include black and midtones: all-white frames cannot expose double normalization.
+    frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype)
+    frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone()
+    expected = frames.float().numpy().copy()
+    sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type)
+    post = get_lingbot_video_post_process_func(None)
+
+    processed = post({"video": frames}, sampling)
+    encoded_frames = np.stack(_coerce_video_to_frames(processed["video"]))
+
+    assert set(processed) == {"video"}
+    assert encoded_frames.dtype == np.float32
+    np.testing.assert_array_equal(encoded_frames, expected)
+    np.testing.assert_array_equal(frames.float().numpy(), expected)
+
+
+@pytest.mark.parametrize("output_key", ["image", "video"])
+def test_postprocess_preserves_latent_payload(output_key: str) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1)
+    sampling = OmniDiffusionSamplingParams(output_type="latent")
+    processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling)
+
+    assert set(processed) == {output_key}
+    assert processed[output_key] is latents
+
+
 def test_postprocess_preserves_image_contract_for_serving():
    from vllm_omni.diffusion.data import DiffusionOutput
    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
@@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving():
        sampling_params=sampling,
        request_id="lingbot-image-contract-test",
    )
-    raw_output = {"image": torch.ones(8, 12, 3)}
+    image = torch.zeros(8, 12, 3)
+    image[4:] = 1.0
+    raw_output = {"image": image}
    processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
    normalized = normalize_diffusion_postprocess_output(processed)

@@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving():
    endpoint_images = _extract_images_from_result(result)
    assert len(endpoint_images) == 1
    assert endpoint_images[0].size == (12, 8)
+    assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0)
+    assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255)


 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
@@ -13,7 +13,6 @@ import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory

 from tests.helpers.mark import hardware_marks
-from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -55,6 +54,7 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
 }

 _OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test"
+_OFFLOAD_MEMORY_PROBE = "vllm_omni_layerwise_offload_memory_test"
 # A second CI build can reverse the same exact-head measurement without a
 # source edit; the default keeps the existing baseline-first workload.
 _MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER"
@@ -65,7 +65,7 @@ _MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = {


 class OffloadStateProbe:
-    """Test-only worker extension exposing the configured offloader state."""
+    """Test-only worker extension exposing offloader and allocator state."""

    model_runner: Any

@@ -86,6 +86,21 @@ class OffloadStateProbe:
            "block_count": sum(group_sizes),
        }

+    def reset_peak_memory_for_test(self) -> dict[str, str]:
+        """Start a process-local inference peak after model initialization."""
+        current_omni_platform.synchronize()
+        torch.accelerator.reset_peak_memory_stats()
+        return {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}
+
+    def get_peak_memory_for_test(self) -> dict[str, str | float]:
+        """Return this worker's allocator peak without counting other jobs."""
+        current_omni_platform.synchronize()
+        return {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": torch.accelerator.max_memory_allocated() / (1024**2),
+        }
+

 def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool:
    if not np.allclose(audio1, audio2, atol=atol):
@@ -96,13 +111,6 @@ def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float
    return True


-def _device_used_mb(device_index: int) -> float:
-    current_omni_platform.synchronize()
-    with current_omni_platform.device(device_index):
-        free_bytes, total_bytes = current_omni_platform.mem_get_info()
-    return (total_bytes - free_bytes) / (1024**2)
-
-
 def _extract_audio(output: Any) -> np.ndarray | None:
    if not output:
        return None
@@ -117,17 +125,29 @@ def _extract_audio(output: Any) -> np.ndarray | None:
    return np.asarray(audio)


-def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+def _collect_probe_results(value: Any, probe: str) -> list[dict[str, Any]]:
    if isinstance(value, dict):
-        return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else []
+        return [value] if value.get("probe") == probe else []
    if isinstance(value, (list, tuple)):
-        states: list[dict[str, Any]] = []
+        results: list[dict[str, Any]] = []
        for item in value:
-            states.extend(_collect_offload_states(item))
-        return states
+            results.extend(_collect_probe_results(item, probe))
+        return results
    return []


+def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+    return _collect_probe_results(value, _OFFLOAD_STATE_PROBE)
+
+
+def _collect_memory_results(value: Any, operation: str) -> list[dict[str, Any]]:
+    return [
+        result
+        for result in _collect_probe_results(value, _OFFLOAD_MEMORY_PROBE)
+        if result.get("operation") == operation
+    ]
+
+
 def _measurement_order() -> tuple[bool, bool]:
    """Return the requested order without changing the default CI workload."""
    configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first")
@@ -144,8 +164,6 @@ def run_inference(
    num_inference_steps: int = 3,
 ) -> dict[str, Any]:
    current_omni_platform.empty_cache()
-    device_index = current_omni_platform.current_device()
-    initial_used_mb = _device_used_mb(device_index)

    if model_name in AUDIO_MODEL:
        params = AUDIO_MODEL_PARAMS
@@ -166,52 +184,45 @@ def run_inference(
        if not offload_states:
            raise AssertionError("The offload-state worker probe returned no diffusion worker results")

-        # Measure steady-state inference memory, not model construction. Enabling
-        # layerwise offload first loads the model and then replaces each block's
-        # device storage with CPU-backed weights.  Monitoring that transition
-        # captures both the original model and temporary staging allocations,
-        # which is not representative of layerwise-offloaded inference.
-        monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-        current_omni_platform.reset_peak_memory_stats()
-        monitor.start()
+        # Reset and read allocator peaks in the diffusion worker process. A
+        # device-wide mem_get_info() sample includes unrelated sibling jobs on
+        # the same GPU and produced negative "savings" in shared ROCm CI.
+        reset_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="reset_peak_memory_for_test", timeout=60),
+            "reset",
+        )
+        if not reset_results:
+            raise AssertionError("The worker peak-memory probe returned no reset acknowledgements")

-        try:
-            # Refer to tests/e2e/offline_inference/test_wan22.py
-            # Use minimal settings for testing
-            output = runner.omni.generate(
-                "A cat sitting on a table",
-                OmniDiffusionSamplingParams(
-                    generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-                    guidance_scale=1.0,
-                    num_inference_steps=num_inference_steps,
-                    **params["sampler_params"],
-                ),
-            )
-        finally:
-            monitor.stop()
+        # Refer to tests/e2e/offline_inference/test_wan22.py.
+        # Use minimal settings for testing.
+        output = runner.omni.generate(
+            "A cat sitting on a table",
+            OmniDiffusionSamplingParams(
+                generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+                guidance_scale=1.0,
+                num_inference_steps=num_inference_steps,
+                **params["sampler_params"],
+            ),
+        )
+        memory_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="get_peak_memory_for_test", timeout=60),
+            "snapshot",
+        )
+        if not memory_results:
+            raise AssertionError("The worker peak-memory probe returned no snapshots")

        audio = _extract_audio(output)
        del output

-    # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
-    # starting usage. Each mode runs in its own spawned process below, so model
-    # objects and process-local allocator/compiler/backend workspace state cannot
-    # carry into the other measurement. Shared filesystem caches can persist;
-    # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that.
-    peak_used_mb = monitor.peak_used_mb
-    incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb)
-
    del runner
    gc.collect()
    cleanup_dist_env_and_memory()
    current_omni_platform.empty_cache()
-    post_cleanup_used_mb = _device_used_mb(device_index)

    return {
-        "initial_used_mb": initial_used_mb,
-        "peak_used_mb": peak_used_mb,
-        "incremental_peak_mb": incremental_peak_mb,
-        "post_cleanup_used_mb": post_cleanup_used_mb,
+        "peak_allocated_mb": max(float(result["peak_allocated_mb"]) for result in memory_results),
+        "memory_results": memory_results,
        "audio": audio,
        "offload_states": offload_states,
    }
@@ -259,12 +270,8 @@ def _assert_audio_outputs(


 def _print_measurement(label: str, measurement: dict[str, Any]) -> None:
-    print(
-        f"{label}: initial={measurement['initial_used_mb']:.1f} MB, "
-        f"peak={measurement['peak_used_mb']:.1f} MB, "
-        f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, "
-        f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB"
-    )
+    print(f"{label}: worker_peak_allocated={measurement['peak_allocated_mb']:.1f} MB")
+    print(f"{label} worker memory: {measurement['memory_results']}")
    print(f"{label} offload state: {measurement['offload_states']}")


@@ -298,6 +305,33 @@ def test_measurement_order_rejects_unknown_value(monkeypatch) -> None:
        _measurement_order()


+@pytest.mark.diffusion
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_collect_memory_results_filters_nested_worker_responses() -> None:
+    value = [
+        None,
+        [
+            {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"},
+            {
+                "probe": _OFFLOAD_MEMORY_PROBE,
+                "operation": "snapshot",
+                "peak_allocated_mb": 123.0,
+            },
+        ],
+        {"probe": _OFFLOAD_STATE_PROBE, "operation": "snapshot"},
+    ]
+
+    assert _collect_memory_results(value, "reset") == [{"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}]
+    assert _collect_memory_results(value, "snapshot") == [
+        {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": 123.0,
+        }
+    ]
+
+
@pytest.mark.diffusion
@pytest.mark.core_model
@pytest.mark.cpu
@@ -366,8 +400,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):

    # Verify that layerwise offloading significantly reduces memory usage
    # Passes only if the actual savings meets the expected savings
-    no_offload_peak_memory = no_offload["incremental_peak_mb"]
-    layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"]
+    no_offload_peak_memory = no_offload["peak_allocated_mb"]
+    layerwise_offload_peak_memory = layerwise_offload["peak_allocated_mb"]
    actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
    assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
        f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
@@ -77,6 +77,7 @@ class _PipelineBackedEngine:
    def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
        self.pipeline = pipeline
        self.executor = SimpleNamespace(
+            is_dead=False,
            register_failure_callback=MagicMock(),
            check_health=MagicMock(),
        )
@@ -127,9 +128,9 @@ def _streaming_diffusion_output(

 def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
    envelope = output.output if isinstance(output.output, dict) else {}
-    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
-    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
-    return payload, metadata
+    payload = envelope.get("payload")
+    metadata = envelope.get("metadata")
+    return payload if isinstance(payload, dict) else {}, metadata if isinstance(metadata, dict) else {}


 class TestPipelineStreamingOutputToStageDiffusionClient:
@@ -230,7 +231,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
        client._response_socket.bind(response_address)
        client._encoder = OmniMsgpackEncoder()
        client._decoder = OmniMsgpackDecoder()
-        client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        client._output_queue = asyncio.Queue[OmniRequestOutput]()
        client._rpc_results = {}
        client._pending_rpcs = set()
        client._tasks = {}
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 from types import SimpleNamespace

@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

 def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
    attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
    calls = []

    def _boundary(query, key, value, attn_metadata=None):
@@ -38,6 +39,7 @@ def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):

 def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
    attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
    calls = []

    def _boundary(query, key, value, attn_metadata=None):
@@ -60,8 +62,10 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
    assert calls == ["impl"]


-def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
+@pytest.mark.parametrize("initialized_with_hsdp", [False, True])
+def test_attention_compile_boundary_without_diffusion_config(monkeypatch, initialized_with_hsdp):
    attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = initialized_with_hsdp
    calls = []

    def _boundary(query, key, value, attn_metadata=None):
@@ -80,4 +84,4 @@ def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
    with set_forward_context(vllm_config=None):
        assert Attention.forward(attention, query, query, query) is query

-    assert calls == ["impl"]
+    assert calls == (["boundary"] if initialized_with_hsdp else ["impl"])
@@ -15,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus

+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
@@ -95,6 +96,30 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc
    assert not adapter.code_prompt_token_ids


+def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
+    adapter, connector = shm_sender
+    ext_id = f"release_{uuid.uuid4().hex}"
+    request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
+    sibling_key = f"{ext_id}_1_0_0"
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    for _ in range(3):
+        adapter.save_async(None, request)
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    # The consumer read only the first chunk before its own request ended.
+    assert connector.get("0", "1", f"{ext_id}_0_0") is not None
+
+    adapter.release_shm_resources(ext_id)
+    # The release is queued behind sends; nothing is unlinked on the caller.
+    assert f"{ext_id}_0_1" in connector._pending_keys
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+
+    assert adapter._reclaimed_shm_total == 2
+    assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
+    assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
@pytest.mark.parametrize("fail_terminal", [False, True])
 def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
    adapter, connector = shm_sender
@@ -2620,10 +2645,9 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
    adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))

-    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
    scheduler = mocker.MagicMock()
    scheduler.chunk_transfer_adapter = adapter_mock
+    bind_omits_transfer_helpers(scheduler)
    scheduler.connector = None
    scheduler.perf_metrics = None
    scheduler.log_stats = False
@@ -213,7 +213,7 @@ class TestCleanup:
        connector.put("s0", "s1", "cleanup_req_42", data)
        assert "cleanup_req_42" in connector._pending_keys

-        connector.cleanup("cleanup_req_42")
+        assert connector.cleanup("cleanup_req_42") is True
        assert "cleanup_req_42" not in connector._pending_keys

        result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
@@ -224,7 +224,7 @@ class TestCleanup:
        connector.put("s0", "s1", "consumed_req_99", data)
        connector.get("s0", "s1", "consumed_req_99", metadata=None)

-        connector.cleanup("consumed_req_99")
+        assert connector.cleanup("consumed_req_99") is False
        assert "consumed_req_99" not in connector._pending_keys

    def test_close_cleans_all_pending(self, connector):
@@ -296,6 +296,24 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
    assert key not in connector._pending_keys


+@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"])
+def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix):
+    request_id = f"prefix_{uuid.uuid4().hex}"
+    own_keys = [f"{request_id}_0_{i}" for i in range(3)]
+    sibling_key = f"{request_id}{suffix}_0_0"
+    for key in own_keys:
+        assert connector.put("0", "1", key, key)[0]
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    # A consumed chunk is no longer counted as reclaimed.
+    assert connector.get("0", "1", own_keys[0])[0] == own_keys[0]
+
+    assert connector.cleanup_prefix(f"{request_id}_0_") == 2
+    for key in own_keys:
+        assert key not in connector._pending_keys
+        assert connector.get("0", "1", key) is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
 # ── Arrival wakeups ───────────────────────────────────────────────────


@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Observe YuE2 lifecycle contracts inside a real engine worker."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class Yue2LifecycleWorkerExtension:
+    model_runner: "GPUARModelRunner"
+
+    def start_yue2_probe(self):
+        from unittest.mock import patch
+
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+        from vllm.config import CUDAGraphMode
+        from vllm.forward_context import get_forward_context, is_forward_context_available
+
+        from vllm_omni.model_executor.models.yue2.yue2 import ABC_END, HOLD_TOKEN, MUSIC_END
+
+        model = self.model_runner.get_model()
+        self.yue2_probe = {
+            "preemptions": 0,
+            "rollbacks": 0,
+            "checked_rows": 0,
+            "mixed_steps": 0,
+            "synthesis_started": False,
+            "ar_full_replays": 0,
+            "ar_piecewise_replays": 0,
+        }
+        graph_call = CUDAGraphWrapper.__call__
+
+        def count_ar_replays(wrapper, *args, **kwargs):
+            if is_forward_context_available():
+                context = get_forward_context()
+                mode = context.cudagraph_runtime_mode
+                entry = wrapper.concrete_cudagraph_entries.get(context.batch_descriptor)
+                if mode == wrapper.runtime_mode and entry is not None and entry.cudagraph is not None:
+                    if mode == CUDAGraphMode.FULL:
+                        self.yue2_probe["ar_full_replays"] += 1
+                    elif mode == CUDAGraphMode.PIECEWISE:
+                        self.yue2_probe["ar_piecewise_replays"] += 1
+            return graph_call(wrapper, *args, **kwargs)
+
+        # This worker is created for this test and exits with the engine.
+        # Count vLLM AR replays, separately from model-owned sampler/NAR graphs.
+        self.yue2_graph_patch = patch.object(CUDAGraphWrapper, "__call__", count_ar_replays)
+        self.yue2_graph_patch.start()
+        self.yue2_cancel_job = None
+        previous: dict[str, int] = {}
+        prepare = model.prepare_runner_inputs
+        reconcile = model._reconcile_history
+        queue = model._queue_synthesis
+        sample = model.sample
+        resolve = model._resolve_pending
+        emitted: dict[str, int] = {}
+
+        def resolve_pending():
+            if model._pending is not None:
+                event, host, states = model._pending
+                event.synchronize()
+                for state, token in zip(states, host[0].tolist()):
+                    end = ABC_END if state.constants.phase == "abc" else MUSIC_END
+                    expected = HOLD_TOKEN if token == end and not state.constants.skip_synthesis else token
+                    assert emitted[state.request_id] == expected, "engine token and model history disagree"
+            resolve()
+
+        def sample_tokens(logits, sampling_metadata):
+            result = sample(logits, sampling_metadata)
+            phases = {model._states[rid].constants.phase for rid, _, _ in model._step_rows if rid in model._states}
+            if len(phases) == 2:
+                self.yue2_probe["mixed_steps"] += 1
+            for (rid, _, _), row in zip(model._step_rows, result.sampled_token_ids.tolist()):
+                emitted[rid] = row[0]
+            return result
+
+        def prepare_inputs(**kwargs):
+            for rid, computed in zip(kwargs["req_ids"], kwargs["num_computed_tokens"]):
+                if int(computed) < previous.get(rid, 0):
+                    self.yue2_probe["preemptions"] += 1
+                previous[rid] = int(computed)
+            return prepare(**kwargs)
+
+        def reconcile_history(state, accepted):
+            if len(state.history) + int(state.end_drawn) > accepted:
+                self.yue2_probe["rollbacks"] += 1
+            reconcile(state, accepted)
+            assert len(state.history) + int(state.end_drawn) <= accepted
+            if not state.finish_ready and state.job is None and not state.finished:
+                assert len(state.history) == accepted
+                assert not state.truncated and not state.end_drawn
+            self.yue2_probe["checked_rows"] += 1
+
+        def queue_synthesis(state):
+            queue(state)
+            if "cancel-synthesis" in state.request_id:
+                job = state.job
+                assert job.started and job.work is not None and job.keep
+                assert job.event is None, "test must intercept an active synthesis job"
+                # Keep an actual NAR job active until the driver cancels it.
+                # Its first work units ran normally; stop pumping later units
+                # so cancellation cannot race against a short song finishing.
+                self.yue2_cancel_advance = job.advance
+                job.advance = lambda: None
+                self.yue2_cancel_job = job
+                self.yue2_probe["synthesis_started"] = True
+
+        model.prepare_runner_inputs = prepare_inputs
+        model._reconcile_history = reconcile_history
+        model._queue_synthesis = queue_synthesis
+        model.sample = sample_tokens
+        model._resolve_pending = resolve_pending
+        return True
+
+    def get_yue2_probe(self):
+        result = dict(self.yue2_probe)
+        job = self.yue2_cancel_job
+        if job is not None:
+            if job.cancelled:
+                # Exercise the original pump after cleanup: it must not submit
+                # further units even though this test retains the job object.
+                self.yue2_cancel_advance()
+            model = self.model_runner.get_model()
+            result["cancelled"] = job.cancelled
+            result["released"] = job._released and job.work is None and not job.keep and job.stream is None
+            result["removed"] = (
+                job.request_id not in model._states
+                and model._synthesis.active is not job
+                and job not in model._synthesis.waiting
+            )
+        return result
@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request.
+
+``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few
+token buckets. A replay reuses static input buffers and the graph pool, so a
+request could in principle read conditioning left behind by an earlier, longer
+prompt. This test runs the same request sequence through an eager encoder and
+through the graphed one and compares the text condition stage 0 hands to
+stage 1 (the fused thinker hidden states) for every request. The sequence
+alternates prompts of different lengths inside one capture bucket, crosses
+into a larger bucket, and then returns to the first prompt.
+
+Like ``test_auk.py`` it needs an assembled checkpoint directory:
+
+    VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py
+"""
+
+from __future__ import annotations
+
+import bisect
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+import yaml
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor
+from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params
+
+MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR"
+
+# Text-only instruct prompts: the token count, and so the capture bucket,
+# follows the text length. SHORT_A and SHORT_B differ in length but share a
+# bucket; LONG lands in a larger one. The test checks these properties on the
+# recorded condition lengths, so a tokenizer change fails loudly instead of
+# silently shrinking the coverage.
+_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".'
+SHORT_A = _TEMPLATE.format(
+    desc="A calm young woman speaking warmly.",
+    text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.",
+)
+SHORT_B = _TEMPLATE.format(
+    desc="An older man with a deep, slightly raspy voice, speaking slowly.",
+    text=(
+        "The train to the coast leaves at seven, so please be on time. Bring a warm coat, "
+        "some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain."
+    ),
+)
+LONG = _TEMPLATE.format(
+    desc=(
+        "A cheerful radio host in her thirties, bright and energetic, with clear articulation, "
+        "a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone"
+    ),
+    text=(
+        "Good morning and welcome to the weekend edition. Today we are talking about small gardens, "
+        "the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into "
+        "a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. "
+        "Stay with us, because after the news we will open the phone lines and hear from you."
+    ),
+)
+SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A]
+
+# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also
+# runs eager in both configs, see _deploy).
+GEN_SECONDS = 1.0
+NFE = 1
+
+# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs
+# the same kernels on the real tokens, so the two paths agree to bf16 rounding.
+# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the
+# condition's max magnitude for every request; the bounds leave about 3x
+# headroom for other GPUs.
+MIN_COSINE = 0.9999
+MAX_REL_ABS = 1e-2
+
+_model_dir = os.environ.get(MODEL_DIR_ENV)
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"),
+]
+
+
+def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str:
+    """``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped.
+
+    Only the encoder differs between the two runs. The conditions compared here
+    are recorded before stage 1 runs, so its DiT compile, graph capture and
+    codec bucket warm-up would only add startup time; eager stage 1 skips them
+    and still takes the real stage handoff.
+    """
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    for stage in deploy["stages"]:
+        if stage["stage_id"] == 1:
+            stage["enforce_eager"] = True
+        elif not graph_encoder:
+            stage["enforce_eager"] = True
+            stage.pop("compilation_config", None)
+    path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml"
+    path.write_text(yaml.safe_dump(deploy))
+    return str(path)
+
+
+def _graph_capture_sizes() -> list[int]:
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    stage = next(s for s in deploy["stages"] if s["stage_id"] == 0)
+    assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test"
+    return sorted(stage["compilation_config"]["cudagraph_capture_sizes"])
+
+
+def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]:
+    """Run SEQUENCE one request at a time and return each request's text condition."""
+    conditions: list[torch.Tensor] = []
+    original = auk_input_processor.encoder2dit
+
+    def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None:
+        result = original(*args, **kwargs)
+        if result is not None:
+            conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone())
+        return result
+
+    # The orchestrator runs in this process and resolves the stage input
+    # function by module attribute at startup, so the patch must precede Omni.
+    with monkeypatch.context() as patch:
+        patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit)
+        with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner:
+            for text in SEQUENCE:
+                prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS)
+                outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0))
+                assert outputs, "no outputs returned"
+    assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests"
+    return conditions
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    capture_sizes = _graph_capture_sizes()
+    eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch)
+    graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch)
+
+    lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)}
+    for i, (ref, got) in enumerate(zip(eager, graphed)):
+        assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}"
+        assert torch.isfinite(got).all(), f"request {i}: non-finite condition"
+        cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item()
+        rel_abs = ((got - ref).abs().max() / ref.abs().max()).item()
+        assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, (
+            f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} "
+            "of the condition's magnitude"
+        )
+
+    # The sequence must exercise what it claims to: two different lengths in
+    # one bucket, a second bucket, and every prompt inside the graphed range.
+    assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size"
+    bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()}
+    assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket)
+    assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket)
+
+    # Repeats of a prompt reproduce its first condition exactly, eager or graphed:
+    # nothing from the requests in between leaks into the replay.
+    for run in (eager, graphed):
+        first: dict[str, torch.Tensor] = {}
+        for text, cond in zip(SEQUENCE, run):
+            if text in first:
+                torch.testing.assert_close(cond, first[text], atol=0, rtol=0)
+            else:
+                first[text] = cond
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final).
+
+The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with
+``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the
+thinker→talker SharedMemoryConnector path.
+
+A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream
+ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still
+``put()`` a finished marker on finish and nothing unlinked it, so every such
+request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``.
+Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any
+undrained segment on request cleanup.
+
+This module starts the production async-chunk deploy and asserts the three
+text-only request ids leave no connector lockfile or POSIX segment in
+``/dev/shm``.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
+_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct")
+_SHM_DIR = Path("/dev/shm")
+_NUM_TEXT_REQUESTS = 3
+_SETTLE_S = 15.0
+
+# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0
+# still uses SharedMemoryConnector. Production yaml is already async_chunk:true.
+_ASYNC_CHUNK_DEPLOY = modify_stage_config(
+    get_deploy_config_path("ci/qwen3_omni_moe.yaml"),
+    updates={"async_chunk": True},
+)
+
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_ASYNC_CHUNK_DEPLOY,
+            use_stage_cli=True,
+        ),
+        id="async_chunk",
+    )
+]
+
+
+def _system_prompt() -> dict:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _shm_used_mib() -> str:
+    try:
+        st = os.statvfs(_SHM_DIR)
+        used = (st.f_blocks - st.f_bavail) * st.f_frsize
+        return f"{used / (1024 * 1024):.1f}MiB"
+    except OSError:
+        return "unknown"
+
+
+def _is_entry_for_request(name: str, request_id: str) -> bool:
+    """Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id."""
+    return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_"))
+
+
+def _entries_for_request_ids(request_ids: list[str]) -> list[str]:
+    if not _SHM_DIR.is_dir():
+        pytest.skip("/dev/shm is not available")
+    found: list[str] = []
+    with os.scandir(_SHM_DIR) as entries:
+        for entry in entries:
+            if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids):
+                found.append(entry.name)
+    return found
+
+
+def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]:
+    """Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*."""
+    deadline = time.monotonic() + timeout_s
+    time.sleep(1.0)
+    leftover: list[str] = _entries_for_request_ids(request_ids)
+    while leftover:
+        if time.monotonic() >= deadline:
+            return leftover
+        time.sleep(2.0)
+        leftover = _entries_for_request_ids(request_ids)
+    return leftover
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None:
+    """Stage-0-final (text-only) requests must not leave connector SHM behind."""
+    print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True)
+
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_system_prompt(),
+        content_text="What is the capital of China? Answer in 20 words.",
+    )
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+    responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS)
+    produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)]
+    assert len(produced) == _NUM_TEXT_REQUESTS, (
+        f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; "
+        f"got {len(produced)} non-empty texts. A preprocess/engine miss would "
+        f"leave /dev/shm unchanged and hide a leak."
+    )
+    request_ids = [resp.request_id for resp in responses]
+    assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, (
+        f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}"
+    )
+
+    leftover = _wait_no_entries_for_request_ids(request_ids)
+    print(
+        f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: "
+        f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}",
+        flush=True,
+    )
+    assert not leftover, (
+        f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}"
+    )
@@ -32,9 +32,9 @@ above the byte floor) and passes on that alone, following the MiniMax Music 3
 precedent.

 Weekly rather than per-PR: each request is a 16 s song and the verification
-checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090
-(24 GB); the hardware mark targets an H100 because the CI fleet has no
-consumer cards.
+checkpoints add ~5.5 GB of downloads. The async/graph implementation was
+validated on H200; the hardware mark targets the CI fleet's H100. The older
+RTX 4090 result does not establish this implementation's memory requirements.
 """

 import concurrent.futures
@@ -60,8 +60,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0
 # 400 frames = 16 s of song. The byte floor sits far below the observed
 # payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS.
 _FRAMES = 400
-# Concurrent pair stays short: it exists to exercise row re-indexing, not
-# to double the wall time.
+# Short concurrent requests exercise row re-indexing as budgets differ.
 _CONCURRENT_FRAMES = 250
 _MIN_BYTES = 1_000_000

@@ -173,7 +172,7 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
    # CI runs without the verification assets, so the structural checks below
    # are the only guard against the two real bugs this PR fixed during review:
    # mono/half-speed output (WAV header + duration bounds) and the int32
-    # crash under concurrent row re-indexing (two simultaneous requests).
+    # crash under concurrent row re-indexing (simultaneous requests).
    # Duration bounds: the end token is masked for the first 200 steps
    # (min_tokens), and the 400-frame budget caps the song, so a healthy
    # request lands in [8, 16.5] s; the half-speed bug produced ~32 s.
@@ -190,13 +189,15 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
    # (see _http_speech); the OpenAI SDK response does not expose headers.
    _http_speech(omni_server, seed=SEED, frames=_FRAMES)

-    # Two concurrent requests with different seeds: both must succeed and
-    # produce different songs (the model-owned sampler re-indexes rows across
-    # concurrent requests; a dtype mismatch here crashed the engine before).
-    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
-        futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)]
-        pair = [f.result() for f in futures]
-    assert pair[0] != pair[1], "different seeds produced identical audio"
+    # Different budgets exercise row re-indexing as shorter requests finish.
+    # Compare the same-budget pair so length cannot hide an ignored seed.
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
+        futures = [
+            pool.submit(_http_speech, omni_server, seed=SEED + i, frames=frames)
+            for i, frames in enumerate((_CONCURRENT_FRAMES, _FRAMES, _CONCURRENT_FRAMES))
+        ]
+        audios = [f.result() for f in futures]
+    assert audios[0] != audios[2], "different seeds produced identical audio at the same frame budget"

    # Without the verification assets the structural assertions above
    # (non-empty payload, byte floor) are the whole test — pass on them.
@@ -235,3 +236,199 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
    assert melody_score > control_score, (
        f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}"
    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_abort_then_next_request_succeeds(omni_server) -> None:
+    """Serving smoke check after a client timeout, not proof of job cleanup."""
+    # The chosen seed and budget time out on the tested cards. Require the
+    # timeout so this cannot silently become another successful request;
+    # whether it lands in AR or synthesis depends on the card.
+    with pytest.raises(httpx.TimeoutException):
+        httpx.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+            json={
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED + 2,
+                "max_new_tokens": 3000,
+                "stream": False,
+                "response_format": "wav",
+                "extra_params": {"cot": "off"},
+            },
+            timeout=1.0,
+        )
+    _http_speech(omni_server, seed=SEED + 3, frames=_CONCURRENT_FRAMES)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("chunked", [True, False], ids=["chunked", "unchunked"])
+@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "ar-graphs"])
+async def test_yue2_preemption_and_synthesis_abort(chunked, enforce_eager, monkeypatch) -> None:
+    """Force real preemption and abort only after NAR work has been submitted."""
+    import asyncio
+    from pathlib import Path
+
+    import torch
+    from vllm import SamplingParams
+
+    from tests.helpers.runtime import get_model_prefix
+    from tests.helpers.stage_config import modify_stage_config
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.model_executor.models.yue2.yue2 import (
+        ABC_END,
+        EOD,
+        KEY_MAX_AUDIO_FRAMES,
+        KEY_MAX_HOLD_STEPS,
+        KEY_MIN_TOKENS,
+        KEY_PHASE,
+        KEY_PREFIX_IDS,
+        KEY_SEED,
+        KEY_SKIP_SYNTHESIS,
+        STOP_TOKEN_IDS,
+    )
+    from vllm_omni.tokenizers.yue2_prompt import abc_prefix_ids, semantic_frames, semantic_prefix_ids
+    from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+    from vllm_omni.transformers_utils.repo_utils import hf_api
+
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    model_path = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + MODEL))
+    if not model_path.is_dir():
+        model_path = Path(hf_api().snapshot_download(str(model_path)))
+    tokenizer = YuE2TextTokenizer(model_path / "qwen.tiktoken")
+    prefix = semantic_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "off")
+    frames = 256
+    # Each song fits alone, but four full songs cannot fit in 64 KV blocks.
+    # Chunked recompute also replays output tokens past the original prompt.
+    deploy = modify_stage_config(
+        get_deploy_config_path("yue2.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enforce_eager": enforce_eager,
+                    "compilation_config": {
+                        "cudagraph_mode": "NONE" if enforce_eager else "FULL_AND_PIECEWISE",
+                    },
+                    "enable_chunked_prefill": chunked,
+                    "async_scheduling": True,
+                    "max_num_batched_tokens": 256 if chunked else 1024,
+                    "max_model_len": 1000,
+                    "num_gpu_blocks_override": 64,
+                    "block_size": 16,
+                    "max_num_seqs": 4,
+                    "tokenizer": str(model_path),
+                }
+            }
+        },
+    )
+    engine = AsyncOmni(
+        model=str(model_path),
+        deploy_config=deploy,
+        worker_extension_cls="tests.e2e.features.yue2.worker_extension.Yue2LifecycleWorkerExtension",
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+
+    async def probe():
+        return (await engine.collective_rpc("get_yue2_probe", stage_ids=[0]))[0][0]
+
+    async def song(request_id, seed):
+        params = SamplingParams(
+            max_tokens=1000 - len(prefix),
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PREFIX_IDS: prefix,
+                KEY_SEED: seed,
+                KEY_MAX_AUDIO_FRAMES: frames,
+                KEY_MIN_TOKENS: frames,
+                KEY_MAX_HOLD_STEPS: 1000 - len(prefix) - frames - 1,
+            },
+        )
+        audio = []
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": prefix}, request_id=request_id, sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+            multimodal = getattr(output.outputs[0], "multimodal_output", None)
+            if multimodal is not None and "audio" in multimodal:
+                data = multimodal["audio"]
+                audio.extend(data if isinstance(data, list) else [data])
+        assert audio, "request completed without audio"
+        waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio])
+        assert torch.isfinite(waveform).all() and waveform.abs().max() > 0
+        assert len(semantic_frames(tokens)) == frames
+        # The VAE convolution trims a small number of boundary samples.
+        assert waveform.numel() == pytest.approx(frames * 1920 * 2, abs=256)
+
+    async def plan():
+        abc_prefix = abc_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "melody")
+        params = SamplingParams(
+            max_tokens=64,
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PHASE: "abc",
+                KEY_PREFIX_IDS: abc_prefix,
+                KEY_SEED: SEED,
+                KEY_SKIP_SYNTHESIS: True,
+            },
+        )
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": abc_prefix}, request_id="abc-plan", sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+        assert tokens and all(t < EOD or t == ABC_END for t in tokens)
+
+    tasks = []
+    try:
+        assert await engine.collective_rpc("start_yue2_probe", stage_ids=[0]) == [[True]]
+        tasks = [asyncio.create_task(plan())]
+        tasks.extend(asyncio.create_task(song(f"preempt-{i}", SEED + i)) for i in range(4))
+        await asyncio.wait_for(asyncio.gather(*tasks), timeout=600)
+        observations = await probe()
+        assert observations["preemptions"] > 0, observations
+        assert observations["rollbacks"] > 0, observations
+        assert observations["checked_rows"] >= 4 * frames, observations
+        assert observations["mixed_steps"] > 0, observations
+        if enforce_eager:
+            assert observations["ar_full_replays"] == 0, observations
+        else:
+            assert observations["ar_full_replays"] > 0, observations
+
+        cancelled = asyncio.create_task(song("cancel-synthesis", SEED + 7))
+        tasks.append(cancelled)
+        deadline = asyncio.get_running_loop().time() + 120
+        while not (await probe())["synthesis_started"]:
+            assert not cancelled.done(), "song finished before the active-job cancellation probe"
+            assert asyncio.get_running_loop().time() < deadline, "NAR job never started"
+            await asyncio.sleep(0.05)
+        cancelled.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await cancelled
+        deadline = asyncio.get_running_loop().time() + 30
+        while True:
+            observations = await probe()
+            if observations.get("removed"):
+                assert observations["cancelled"] and observations["released"], observations
+                break
+            assert asyncio.get_running_loop().time() < deadline, observations
+            await asyncio.sleep(0.05)
+        await asyncio.wait_for(song("after-cancel", SEED + 8), timeout=120)
+        observations = await probe()
+        assert observations["released"]
+        print("YuE2 lifecycle:", observations)
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        engine.shutdown()
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Deterministic tests for the CFG-companion output lifecycle.

 Regression suite for the companion-output race behind the nightly Bagel
@@ -94,6 +94,9 @@ class _FakePool:
    def release_bindings(self, request_ids):
        pass

+    async def release_request_resources(self, request_ids):
+        pass
+

 def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
    orch = Orchestrator.__new__(Orchestrator)
@@ -13,13 +13,15 @@ substitutable with the base class — including base-style calls that pass

 import inspect

+import msgspec
 import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.request import Request

-from vllm_omni.engine import PromptEmbedsPayload
+from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -111,3 +113,24 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
    assert update is not None
    assert request.model_intermediate_buffer is buffer
    assert update.model_intermediate_buffer is buffer
+
+
+@pytest.mark.parametrize("session_id", [None, "session"])
+def test_session_survives_omni_upgrade_wire_transport_and_scheduler_conversion(session_id):
+    core = EngineCoreRequest(
+        request_id="session-request",
+        prompt_token_ids=[1, 2],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=2),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        session_id=session_id,
+    )
+    upgraded = OmniEngineCoreRequest.from_request(core)
+    decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
+    scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
+    assert decoded.session_id == session_id
+    assert scheduled.session_id == session_id
@@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
    )


-async def _codec_output(obj, raw):
-    replica = obj.stage_pools[1].replica_id
-    processed = await obj._process_llm_stage_outputs(1, replica, raw, set())
-    await obj._handle_processed_outputs(1, replica, processed)
+async def _codec_output(obj, raw, stage=1):
+    replica = obj.stage_pools[stage].replica_id
+    processed = await obj._process_llm_stage_outputs(stage, replica, raw, set())
+    await obj._handle_processed_outputs(stage, replica, processed)


 def _messages(obj):
@@ -132,6 +132,57 @@ def _audio(output):
    return torch.cat(audio) if isinstance(audio, list) else audio


+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", list(RequestOutputKind))
+@pytest.mark.parametrize("suffix_overtook", [False, True])
+async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook):
+    obj = _orchestrator(kind, final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    suffix = _raw(required=True, samples=(20, 21))
+    terminal = _raw(required=True, terminal=True, samples=(30, 31))
+    if suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+        assert obj.output_async_queue.empty()
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11)))
+    if not suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+    messages = _messages(obj)
+    expected = {
+        RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
+        RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
+        RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
+    }
+    assert [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind]
+    assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages)
+    assert messages[-1].finished
+    assert not obj.request_states["r"].pending_first_audio_outputs
+
+
+@pytest.mark.asyncio
+async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    assert obj.request_states["r"].pending_first_audio_outputs
+    await obj._cleanup_request_ids(["r"], abort=True)
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    assert not obj.request_states and obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage) and error.stage_id == 0
+    assert not obj.request_states
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", list(RequestOutputKind))
@pytest.mark.parametrize("codec_overtook", [False, True])
@@ -1,9 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import patch

+import pytest
 from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.executor.uniproc_executor import UniProcExecutor

-from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc
+from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "first_decoder,stream_decoder,stream_first_audio,expected",
+    [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)],
+)
+def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
+    mocker,
+    first_decoder,
+    stream_decoder,
+    stream_first_audio,
+    expected,
+):
+    executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
+    )
+    executor.driver_worker = mocker.Mock()
+    runner = executor.driver_worker.worker.model_runner
+    runner.model.first_frame_decoder = object() if first_decoder else None
+    runner.model.stream_decoder = object() if stream_decoder else None
+    runner.model.stream_first_audio = stream_first_audio
+    runner._omni_data_plane = None
+    assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
+    assert runner.model_state.set_first_audio_sink.call_count == int(expected)


 def test_preprocess_add_request_preserves_omni_fields():
@@ -26,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields():
    assert current_wave == 3
    assert result.external_req_id == "external"
    assert result.additional_information == {"conditioning": "payload"}
+
+
+def test_codec_uses_generic_first_audio_binding(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+
+    def hook(*args):
+        return False
+
+    def bind(sink):
+        calls["sink"] = sink
+        return hook
+
+    plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
+    model = SimpleNamespace(bind_first_chunk_fast_path=bind)
+    runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    assert calls["hook"] is hook
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+
+
+def test_first_audio_binding_preserves_talker_marker(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+    runner = SimpleNamespace(
+        model=SimpleNamespace(first_frame_decoder=object()),
+        model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
+    )
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
+    executor.vllm_config.parallel_config.tensor_parallel_size = 2
+    assert not module._bind_first_audio_sink(executor, outputs, scheduler)
@@ -154,6 +154,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
    asyncio.run(run())


+@pytest.mark.cpu
+def test_release_request_resources_skips_without_async_chunk():
+    async def run() -> None:
+        call = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=call)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
+        )
+
+        await pool.release_request_resources(["req-1"])
+        call.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
+    async def run() -> None:
+        async def hang(*_args):
+            await asyncio.Event().wait()
+
+        live = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
+        )
+        monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
+
+        release = asyncio.create_task(pool.release_request_resources(["req-1"]))
+        await asyncio.sleep(0.05)
+        live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
+        assert not release.done()
+
+        await asyncio.wait_for(release, timeout=1.0)
+
+    asyncio.run(run())
+
+
@pytest.mark.cpu
@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
@pytest.mark.parametrize("failure", ["error", "timeout", "missing"])
@@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt):
            assert audio_file.samplerate == 22050
    else:
        assert len(raw) == frame_count * 2
+
+
+def test_float32_pcm_fast_path_matches_soundfile():
+    from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes
+
+    rng = np.random.default_rng(0)
+    audio = np.concatenate(
+        [
+            rng.uniform(-1.2, 1.2, 4096),
+            np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]),
+        ]
+    ).astype(np.float32)
+    with BytesIO() as buffer:
+        soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16")
+        expected = buffer.getvalue()
+    assert _float32_to_pcm16_bytes(audio) == expected
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Processor:
+    """Segment API of the model's processor with a deterministic tokenizer."""
+
+    def __init__(self):
+        self.model_config = SimpleNamespace(
+            n_vq=3,
+            im_start_token_id=1,
+            im_end_token_id=2,
+            audio_start_token_id=3,
+            audio_pad_token_id=1024,
+            audio_assistant_slot_token_id=5,
+        )
+        self.encoded = []
+
+    def _encode_text(self, text):
+        self.encoded.append(text)
+        return [100 + ord(c) % 50 for c in text]
+
+    def _user_prompt_prefix_ids(self):
+        return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
+
+    def _user_prompt_after_reference_ids(self, language, fields):
+        return self._encode_text(f"lang={fields.get('language', language)}")
+
+    def _assistant_prompt_prefix_ids(self):
+        return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
+
+    def reference(self, text, language, codes):
+        # Mirrors the processor's _build_continuation_codes + _pad (one item).
+        fields = {} if language is None else {"language": language}
+        ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language, fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [3]
+        )
+        text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
+        text_rows[:, 0] = torch.tensor(ids)
+        audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
+        audio_rows[:, 0] = 5
+        audio_rows[:, 1:] = codes
+        unified = torch.cat([text_rows, audio_rows])
+        return unified[:, 0].tolist(), unified[:, 1:].contiguous()
+
+
+@pytest.mark.parametrize("language", [None, "English"])
+def test_matches_processor_and_caches_constant_segments(language):
+    proc = _Processor()
+    codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
+    kwargs = {"text": "ref words target words"}
+    if language is not None:
+        kwargs["language"] = language
+    ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
+    expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
+    assert ids == expected_ids
+    assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
+
+    proc.encoded.clear()
+    _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
+    assert proc.encoded == ["other"]
+
+
+def test_unsupported_layouts_fall_back():
+    proc = _Processor()
+    codes = torch.zeros((2, 3), dtype=torch.long)
+    assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None
@@ -7,10 +7,12 @@ from __future__ import annotations
 import asyncio
 import base64
 import gc
+import hashlib
 import io
 import json
 import threading
 import weakref
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock

@@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import (
    QwenOmniStreamingVideoHandler,
    StreamingVideoSessionConfig,
 )
-from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler
+from vllm_omni.entrypoints.openai.video_stream_base import (
+    OmniStreamingVideoHandler,
+)
 from vllm_omni.outputs import OmniRequestOutput

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
        pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"),
    ],
 )
-async def test_video_frame_selection_preserves_uncached_frames(
-    monkeypatch, num_frames, frame_ids, repeat_image, expected_indices
-):
-    """Prewarm misses still use image URLs, with metadata paired by buffer position."""
+async def test_video_frame_selection_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices):
+    """Queries wait for prewarmed PIL frames, with metadata paired by buffer position."""
    engine_prompts = []
-    pending_decode = asyncio.Event()
-
-    async def blocked_to_thread(*args, **kwargs):
-        await pending_decode.wait()
-
-    monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)

    class OneOutputEngine:
        async def generate(self, **kwargs):
@@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames(
    assert not [message for message in ws.sent if message["type"] == "error"]
    assert len(engine_prompts) == 1
    content = engine_prompts[0]["messages"][-1]["content"]
-    assert content == [
-        *[
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}}
-            for index in expected_indices
-        ],
-        {"type": "text", "text": "describe"},
-    ]
+    assert content[-1] == {"type": "text", "text": "describe"}
+    image_parts = content[:-1]
+    assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices)
+    for part, index in zip(image_parts, expected_indices):
+        raw = base64.b64decode(frames[index])
+        assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes()
+        assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest()
    consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"]
    if not any(frame_ids):
        assert consumed_events == []
@@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc

    class BlockingEngine:
        async def generate(self, *, request_id, **kwargs):
+            submitted = kwargs.get("submitted")
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
            if request_ids:
                assert aborted[-1] == request_ids[-1]
            request_ids.append(request_id)
@@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query():


@pytest.mark.asyncio
-async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
+async def test_query_waits_for_frame_prewarm(monkeypatch):
    decode_started = threading.Event()
    release_decode = threading.Event()
    query_started = asyncio.Event()
@@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
    assert decode_started.is_set()

    ws.put({"type": "video.query", "text": "describe"})
-    await asyncio.wait_for(query_started.wait(), timeout=2.0)
+    await asyncio.sleep(0.05)
+    assert not query_started.is_set()

    release_decode.set()
+    await asyncio.wait_for(query_started.wait(), timeout=2.0)
    ws.put({"type": "video.done"})
    await asyncio.wait_for(task, timeout=2.0)
    assert "session.done" in ws.sent_types()


+@pytest.mark.asyncio
+async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch):
+    """APC off and num_frames=1: a blocked older frame must not delay the query."""
+    frame_a = _make_jpeg(1, 0, 0)
+    frame_b = _make_jpeg(0, 1, 0)
+    decode_a_started = threading.Event()
+    release_a = threading.Event()
+    decode_b_done = threading.Event()
+    query_started = asyncio.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == frame_a:
+            decode_a_started.set()
+            release_a.wait(timeout=5.0)
+        image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        if raw_bytes == frame_b:
+            decode_b_done.set()
+        return image
+
+    class RecordingHandler(QwenOmniStreamingVideoHandler):
+        async def _process_query(self, *args, **kwargs):
+            query_started.set()
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    ws = TimedWebSocket()
+    handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0)
+    task = asyncio.create_task(handler.handle_session(ws))
+    try:
+        ws.put(
+            {
+                "type": "session.config",
+                "model": "test",
+                "modalities": ["text"],
+                "num_frames": 1,
+                "enable_frame_filter": False,
+            }
+        )
+        await asyncio.sleep(0)
+        ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"})
+        for _ in range(100):
+            if decode_a_started.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_a_started.is_set()
+        ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"})
+        for _ in range(200):
+            if decode_b_done.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_b_done.is_set()
+
+        ws.put({"type": "video.query", "text": "describe"})
+        await asyncio.wait_for(query_started.wait(), timeout=2.0)
+        assert not release_a.is_set()
+        release_a.set()
+        ws.put({"type": "video.done"})
+        await asyncio.wait_for(task, timeout=2.0)
+    finally:
+        release_a.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"])
 async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames):
@@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only():
    assert messages[1] == {"role": "assistant", "content": "recent answer"}
    assert messages[2] == user_message
    assert user_message["content"][-1] == {"type": "text", "text": "current question"}
+
+
+# ---------------------------------------------------------------------------
+# Incremental prefill (warmup) lifecycle
+# ---------------------------------------------------------------------------
+
+
+class RecordingReuseEngine:
+    """Engine stub with stage-0 prefix caching on; records generate() and abort()."""
+
+    def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0):
+        self.requests: list[dict[str, Any]] = []
+        self.aborted: list[str] = []
+        self._warmup_delay = warmup_delay
+        self._query_delay = query_delay
+        self.engine = SimpleNamespace(
+            stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))]
+        )
+
+    def warmup_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")]
+
+    def query_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")]
+
+    def generate(self, **kwargs):
+        self.requests.append(kwargs)
+        submitted = kwargs.get("submitted")
+        if submitted is not None and not submitted.done():
+            submitted.set_result(None)
+        is_warmup = kwargs["request_id"].startswith("video-warmup-")
+        delay = self._warmup_delay if is_warmup else self._query_delay
+
+        async def _gen():
+            if delay:
+                await asyncio.sleep(delay)
+            yield _text_result("ok")
+
+        return _gen()
+
+    async def abort(self, request_id: str) -> None:
+        self.aborted.append(request_id)
+
+
+class PromptRecordingHandler(QwenOmniStreamingVideoHandler):
+    """Records each query turn's user-message content and processor kwargs."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.query_contents: list[list[dict[str, Any]]] = []
+        self.warmup_contents: list[list[dict[str, Any]]] = []
+        self.mm_processor_kwargs: list[Any] = []
+
+    async def _preprocess_to_engine_prompt(self, request):
+        self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None))
+        return {"prompt": "p"}
+
+    def build_engine_prompt(
+        self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None
+    ):
+        messages, user_message = super().build_engine_prompt(
+            config,
+            frame_buffer,
+            audio_buffer,
+            message_history,
+            query_text,
+            prewarmed_frames,
+            frame_indices=frame_indices,
+        )
+        self.query_contents.append(list(user_message["content"]))
+        return messages, user_message
+
+    def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames):
+        messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames)
+        if messages:
+            self.warmup_contents.append(list(messages[-1]["content"]))
+        return messages
+
+
+async def _poll(predicate, timeout: float = 3.0) -> bool:
+    for _ in range(int(timeout / 0.01)):
+        if predicate():
+            return True
+        await asyncio.sleep(0.01)
+    return bool(predicate())
+
+
+@pytest.mark.asyncio
+async def test_warmup_fires_on_frame_and_query_cancels_it():
+    engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+
+    assert await _poll(lambda: engine.warmup_ids())
+    warmup_id = engine.warmup_ids()[0]
+    assert engine.requests[0]["sampling_params"].max_tokens == 1
+
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    # in-flight warmup was cancelled + aborted, and not restarted mid-query
+    assert warmup_id in engine.aborted
+    assert len(engine.warmup_ids()) == 1
+    assert handler.mm_processor_kwargs[:2] == [
+        {"use_audio_in_video": True},
+        {"use_audio_in_video": True},
+    ]
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    # the committed turn changed the context -> a fresh warmup fires
+    assert await _poll(lambda: len(engine.warmup_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_warmup_reaches_stage0_before_abort():
+    """Cancel during submit must not abort a warmup the engine never received."""
+
+    class _HoldSubmitEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldSubmitEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(engine.started.is_set)
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.05)
+    assert engine.requests == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.aborted)
+    assert engine.warmup_ids()
+    assert engine.aborted[0] == engine.warmup_ids()[0]
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_query_reaches_stage0_before_abort():
+    """Cancel during query submit must not abort a request the engine never received."""
+
+    class _HoldQueryEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.query_started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            if str(kwargs.get("request_id", "")).startswith("video-warmup-"):
+                return super().generate(**kwargs)
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.query_started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldQueryEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.query", "text": "first"})
+    assert await _poll(engine.query_started.is_set)
+    ws.put({"type": "video.query", "text": "second"})
+    await asyncio.sleep(0.05)
+    assert engine.query_ids() == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.query_ids())
+    first_query = engine.query_ids()[0]
+    assert await _poll(lambda: first_query in engine.aborted)
+    assert await _poll(lambda: len(engine.query_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_audio_session_never_starts_warmup():
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))})
+    await asyncio.sleep(0.3)
+
+    assert engine.warmup_ids() == []
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_warmup_starts_only_after_frames_ready(monkeypatch):
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        decode_started.set()
+        release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(decode_started.is_set)
+    await asyncio.sleep(0.1)
+    assert engine.warmup_ids() == []
+
+    release_decode.set()
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch):
+    """Cancel warmup during prefill must not drop a later frame still decoding."""
+    frame1 = _b64(_make_jpeg(1, 1, 1))
+    frame2 = _b64(_make_jpeg(2, 2, 2))
+    raw2 = base64.b64decode(frame2)
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == raw2:
+            decode_started.set()
+            release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine(warmup_delay=10.0)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame1})
+    assert await _poll(lambda: engine.warmup_ids())
+
+    ws.put({"type": "video.frame", "data": frame2})
+    assert await _poll(decode_started.is_set)
+
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.1)
+    release_decode.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    assert len(image_parts) == 2
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_keeps_frames_evicted_while_pinned(monkeypatch):
+    """Regression: frames FIFO-evicted while an in-flight query awaits prewarm
+    must keep their PIL entries and stay in the query prompt."""
+    frame_a = _b64(_make_jpeg(1, 1, 1))
+    frame_a2 = _b64(_make_jpeg(2, 2, 2))
+    frame_b = _b64(_make_jpeg(3, 3, 3))
+    raw_a = base64.b64decode(frame_a)
+    decoded: list[bytes] = []
+    release_others = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes != raw_a:
+            release_others.wait(timeout=5.0)
+        decoded.append(raw_bytes)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "max_frames": 2,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame_a})
+    assert await _poll(lambda: raw_a in decoded)
+    ws.put({"type": "video.frame", "data": frame_a2})  # decode gated -> query will wait
+
+    ws.put({"type": "video.query", "text": "describe"})  # pins [a, a2]
+    await asyncio.sleep(0.1)
+    ws.put({"type": "video.frame", "data": frame_b})  # buffer full -> evicts pinned a
+    await asyncio.sleep(0.1)
+    release_others.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)]
+    assert uuids == expected
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "max_frames": max_frames,
+            "num_frames": max_frames,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    for frame in frames:
+        ws.put({"type": "video.frame", "data": frame})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: handler.query_contents)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"]
+
+
+def _frame_uuid(frame_b64: str) -> str:
+    return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest()
+
+
+@pytest.mark.asyncio
+async def test_below_max_frames_incremental_matches_legacy_window():
+    frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)]
+    expected = [_frame_uuid(frame) for frame in frames]
+    assert await _query_image_uuids(["text"], frames, max_frames=4) == expected
+    assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected
+
+
+@pytest.mark.asyncio
+async def test_incremental_full_window_drops_oldest_half():
+    frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)]
+    kept = await _query_image_uuids(["text"], frames, max_frames=4)
+    legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4)
+    assert kept == [_frame_uuid(frame) for frame in frames[2:]]
+    assert legacy == [_frame_uuid(frame) for frame in frames[1:]]
+
+
+async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame})
+    if "audio" not in modalities:
+        assert await _poll(lambda: engine.warmup_ids())
+    else:
+        await asyncio.sleep(0.05)
+        assert engine.warmup_ids() == []
+    ws.put({"type": "audio.chunk", "data": _b64(pcm)})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return handler
+
+
+@pytest.mark.asyncio
+async def test_reuse_query_with_audio_matches_legacy_prompt():
+    """Incremental and legacy query prompts match; both set the processor kwarg."""
+    frame = _b64(_make_jpeg())
+    pcm = b"\x01\x00\x02\x00"
+    reuse = await _run_frame_audio_query(["text"], frame, pcm)
+    legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm)
+
+    assert reuse.warmup_contents
+    assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0])
+    assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"]
+
+    assert reuse.query_contents and legacy.query_contents
+    reuse_q = reuse.query_contents[0]
+    legacy_q = legacy.query_contents[0]
+    assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"]
+    assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"]
+    assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"]
+    assert reuse_q[1] == legacy_q[1]
+    assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"}
+
+    assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True}
+    assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True}
+    assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True}
@@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
    audio_tokenizer.to.assert_called_once()


+@pytest.mark.parametrize(
+    "variant,graphs_env,prepares",
+    [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
+)
+def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
+    if graphs_env is None:
+        monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    adapter._moss_variant = variant
+    encoder = mocker.Mock()
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    assert encoder.prepare.called is prepares
+
+
+def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = RuntimeError("capture failed")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    encoder.prepare.assert_called_once_with()
+
+
+def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    with pytest.raises(SharedReferenceEncoderStartupError):
+        asyncio.run(adapter.warmup())
+
+
 def test_qwen3_tts_metadata():
    assert Qwen3TTSAdapter.backend == "ar"
    assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
@@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client(
        produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))

    class Engine:
-        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
+        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None):
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
            if explicit_sampling:
                assert sampling_params_list is not None
                assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA
@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
+from vllm_omni.config.stage_config import build_stage_runtime_overrides
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
    assert launched["model"] == "generic-diffusion"
    assert launched["stage_id"] == 0
    assert runtime.stage_pools[0].clients == [client]
+
+
+def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0]
+
+    assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("prefix", ["", "stage_0_"])
+def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None:
+    overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"})
+
+    assert overrides["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("strategy", ["full", "pre_sharded"])
+def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None:
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"hsdp_weight_load_strategy": strategy}
+    )
+
+    assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy
+
+
+@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"])
+def test_removed_hsdp_loading_strategies_are_rejected(strategy):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="hsdp_weight_load_strategy"):
+        OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)
@@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch):
        def build_and_log_summary(self):
            return "Fake summary"

+        def log_timing_summary(self):
+            pass
+
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni.OrchestratorAggregator",
        _FakeOrchestratorAggregator,
@@ -97,6 +97,7 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[
 class OmniResponse:
    """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""

+    request_id: str | None = None
    text_content: str | None = None
    audio_data: list[str] | None = None
    audio_content: str | None = None
@@ -421,6 +422,8 @@ class OnlineOmniClient:
            text_content = ""
            audio_data = []
            for chunk in chat_completion:
+                if result.request_id is None:
+                    result.request_id = getattr(chunk, "id", None)
                for choice in chunk.choices:
                    content = getattr(getattr(choice, "delta", None), "content", None)
                    modality = getattr(chunk, "modality", None)
@@ -453,6 +456,7 @@ class OnlineOmniClient:
        """Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
        result = OmniResponse()
        try:
+            result.request_id = getattr(chat_completion, "id", None)
            audio_data = None
            text_content = None
            for choice in chat_completion.choices:
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Scheduler test helpers shared by AR / chunk-transfer unit tests."""
+
+from __future__ import annotations
+
+from types import MethodType
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+
+def bind_omits_transfer_helpers(sched) -> None:
+    """Bind the real omit helpers so a MagicMock scheduler does not skip save_async.
+
+    MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage``
+    silently skips the chunk put.
+    """
+    sched._omits_kv_transfer_cache = {}
+    sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched)
+    sched._request_omits_kv_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_kv_transfer_to_next_stage,
+        sched,
+    )
+    sched._request_omits_chunk_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_chunk_transfer_to_next_stage,
+        sched,
+    )
@@ -111,7 +111,7 @@ class TestFailureCounterWiring:
        obj.request_states["req-1"] = SimpleNamespace(
            metrics=SimpleNamespace(
                e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
            ),
            failure_recorded=False,
        )
@@ -127,7 +127,7 @@ class TestFailureCounterWiring:
        obj.request_states["req-1"] = SimpleNamespace(
            metrics=SimpleNamespace(
                e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
            ),
            failure_recorded=False,
        )
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations

 import pytest
@@ -155,3 +158,39 @@ def test_build_and_log_summary_multiple_requests() -> None:
    r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2")
    assert len(r1_stage_entry["stages"]) == 2
    assert len(r2_stage_entry["stages"]) == 1
+
+
+def test_log_timing_summary_skips_tables_unless_debug(monkeypatch) -> None:
+    from vllm_omni.metrics import stats
+
+    agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0)
+    agg.stage_first_ts[0] = 0.0
+    agg.stage_last_ts[0] = 0.02
+    agg.on_stage_metrics(
+        0,
+        "r1",
+        StageRequestStats(
+            batch_id=1,
+            batch_size=1,
+            num_tokens_in=3,
+            num_tokens_out=3,
+            stage_gen_time_ms=20.0,
+            rx_transfer_bytes=0,
+            rx_decode_time_ms=0.0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(),
+        ),
+    )
+    agg.on_finalize_request(0, "r1", req_start_ts=0.0)
+    tables, lines = [], []
+
+    def record_table(*args, **kwargs):
+        tables.append(args)
+        return ""
+
+    monkeypatch.setattr(stats, "_format_table", record_table)
+    monkeypatch.setattr(stats.logger, "info", lambda msg, *args: lines.append(msg % args))
+    monkeypatch.setattr(stats.logger, "isEnabledFor", lambda level: False)
+    agg.log_timing_summary()
+    assert not tables
+    assert len(lines) == 1 and lines[0].startswith("[OmniTiming] req=r1 total=")
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _reference(logits, uniforms, top_k, top_p):
+    if top_k > 0:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = logits.sort(dim=-1, descending=True)
+        probs = values.softmax(dim=-1, dtype=torch.float32)
+        values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
+        logits = values.scatter(1, indices, values)
+    return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+
+
+@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
+def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(173)
+    logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
+    # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
+    uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
+    expected = _reference(logits, uniforms, top_k, top_p)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    logits.add_(torch.randn_like(logits) * 0.1)
+    uniforms.uniform_(1e-7, 1 - 1e-7)
+    graph.replay()
+    torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
+def test_threshold_ties_and_degenerate_rows(top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    logits = torch.tensor(
+        [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_nucleus_boundaries(dtype):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(1923)
+    logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
+    values, indices = logits.sort(descending=True)
+    probabilities = values.softmax(-1, dtype=torch.float32)
+    before = probabilities.cumsum(-1) - probabilities
+    for rank in [1, 19, 49, 1000, 2000]:
+        boundary = before[0, rank]
+        uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
+        uniforms[0, indices[0, rank]] = 1 - 1e-7
+        for direction in [-torch.inf, torch.inf]:
+            top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
+            actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
+            torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_p", [0.8, 1.0])
+def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    for seed in range(16):
+        torch.manual_seed(seed)
+        logits = torch.randn(128, 2050, device="cuda")[:, ::2]
+        uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
+        actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
+        torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)
@@ -16,7 +16,7 @@ pytestmark = [
@pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
@pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
 def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
-    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+    from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention

    with torch.inference_mode():
        torch.manual_seed(42)
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HiFT graphs consume live inputs and return independently owned waveforms."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("finalize", [False, True])
+@torch.inference_mode()
+def test_hift_graph_preserves_current_and_previous_request(finalize):
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.f0_predictor.remove_weight_norm()
+    hift.enable_decode_graphs()
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    saved: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for _ in range(5):
+        mel = torch.randn(1, 80, 30, device="cuda", generator=generator)
+        source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator)
+        expected = hift._decode_eager(mel, source, finalize=finalize)
+        actual = hift.decode(mel, source, finalize=finalize)
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+        for previous, snapshot in saved:
+            assert torch.equal(previous, snapshot)
+        if hift._decode_graphs.graphs:
+            saved.append((actual, actual.clone()))
+    assert len(hift._decode_graphs.graphs) == 1
+    hift._decode_graphs.max_graphs = 1
+    mel = torch.randn(1, 80, 31, device="cuda", generator=generator)
+    source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator)
+    actual = hift.decode(mel, source, finalize=finalize)
+    expected = hift._decode_eager(mel, source, finalize=finalize)
+    assert torch.equal(actual, expected)
+    assert len(hift._decode_graphs.graphs) == 1
@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import
    HiggsAudioV3TalkerForConditionalGeneration,
 )
 from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -34,6 +35,8 @@ def make_state():
        _audio_continuation_id=99,
        update_decode_step_metadata=Mock(),
    )
+    state._stream_pos = {}
+    state._eager_state = EagerMTPState(state)
    return state


@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Waveform ownership across graph replays and asynchronous output copies."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self, channels):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+        self.channels = channels
+        self.config = SimpleNamespace(codebook_size=1024)
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch, device=self.weight.device)
+
+    def reset_decoder_state_slots(self, slots):
+        self.state[slots] = 0
+
+    def decode_streaming_batch(self, codes, lengths, slots, valid_rows):
+        audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None]
+        self.state[slots] = audio[:, -1]
+        channels = torch.arange(1, self.channels + 1, device=audio.device)
+        return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths)
+
+
+class ReplayGraph:
+    """A real CUDA graph with deliberately reused output storage."""
+
+    batch_sizes = [1]
+
+    def __init__(self, codec):
+        self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long)
+        self.slots = torch.zeros(1, device="cuda", dtype=torch.long)
+        self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long)
+        self.valid = torch.ones(1, device="cuda", dtype=torch.bool)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        codec.state.zero_()
+
+    def decode(self, codes, slots, *, allow_frame_padding):
+        if codes.shape[1:] != (1, 3):
+            return None
+        self.codes.copy_(codes)
+        self.slots.copy_(slots)
+        self.graph.replay()
+        return self.output.audio, self.output.audio_lengths, 1
+
+
+def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type=model_type),
+            async_chunk=streaming,
+            use_v2_model_runner=v2,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    decoder._codec = CausalCodec(channels).to(device)
+    decoder._n_channels = channels
+    decoder._stream_max_step_frames = 3
+    return decoder
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])])
+def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            stage_connector_config={
+                "extra": {
+                    "initial_codec_chunk_frames": 1,
+                    "codec_chunk_frames": 15,
+                    "codec_chunk_ramp": ramp,
+                }
+            },
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    assert decoder._streaming_graph_frame_sizes == frames
+    assert decoder._stream_max_step_frames == 15
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "v2,streaming,model_type",
+    [
+        (False, True, "moss_tts_local"),
+        (True, False, "moss_tts_local"),
+        (True, True, "moss_tts_realtime"),
+        (True, True, "moss_tts_local"),
+    ],
+)
+def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type):
+    decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type)
+    assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local")
+    session = decoder._ensure_stream_session()
+    slot = session.acquire()
+    assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu"
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("v2", [False, True])
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("graph", [False, True])
+@torch.no_grad()
+def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph):
+    from vllm_omni.outputs import OmniModelRunnerOutput
+    from vllm_omni.worker_v2.omni_generation_model_runner import (
+        OmniGenerationAsyncOutput,
+        OmniGenerationModelRunner,
+        _contains_cuda_tensor,
+    )
+
+    decoder = make_decoder(v2=v2, channels=channels, device="cuda")
+    # Real Local-v1.5 weights replace the class's 24 kHz default at load time.
+    decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32)
+    session = decoder._ensure_stream_session()
+    if graph:
+        session._cudagraph_wrapper = ReplayGraph(decoder._codec)
+    copy_stream = torch.cuda.Stream()
+    pending = []
+    # Do not materialize any output until later forwards have replayed the
+    # graph, reset terminal/cancelled state, and reused those state slots.
+    for round_index in range(4):
+        value = round_index + 1
+        codes = torch.cat(
+            [
+                torch.full((14,), value, device="cuda", dtype=torch.long),
+                torch.full((6,), value + 10, device="cuda", dtype=torch.long),
+            ]
+        )
+        info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]]
+        output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6])
+        # The runner may overwrite its input buffer before codec completion.
+        codes.fill_(-100)
+        payload = output.multimodal_outputs
+        assert _contains_cuda_tensor(payload) == v2
+        assert isinstance(payload, PackedOutputSnapshot) == v2
+        if v2:
+            # Exactly one waveform slab, regardless of request/chunk count.
+            assert sum(slab.is_cuda for _, slab in payload._groups) == 1
+            result = OmniModelRunnerOutput(
+                req_ids=["tail", "empty", "live"],
+                req_id_to_index={"tail": 0, "empty": 1, "live": 2},
+                sampled_token_ids=[[], [], []],
+            )
+            pending.append(
+                OmniGenerationAsyncOutput(
+                    model_runner_output=result,
+                    multimodal_outputs=payload,
+                    num_reqs=3,
+                    main_stream=torch.cuda.current_stream(),
+                    copy_stream=copy_stream,
+                )
+            )
+        else:
+            pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3))
+        decoder.on_requests_finished(["live"])
+        assert decoder._stream_req_slots == {}
+        assert not session._leased_slots
+
+    for round_index, item in enumerate(pending):
+        rows = item.get_output().multimodal_outputs if v2 else item
+        for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]):
+            expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)
+            if channels > 1:
+                expected = expected[None] * torch.arange(1, channels + 1)[:, None]
+            torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0)
+            assert rows[index]["sr"].item() == 48_000
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_owned_input_copy_clamps_entire_ragged_batch():
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=2)
+    codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda")
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")],
+        seq_token_counts=[6, 0, 2],
+    )
+    codes.fill_(-100)
+    cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True)
+    torch.accelerator.synchronize()
+    expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]])
+    torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0)
+    assert cpu["model_outputs"][1].numel() == 0
+    torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0)
+    decoder.on_requests_finished(["a", "b"])
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals():
+    decoder = make_decoder(device="cuda")
+    session = decoder._ensure_stream_session()
+    # Force reuse every call, with an outstanding transfer on the first call.
+    # This delay tests host-buffer ownership, never performance or overlap.
+    session._metadata_ring = session._metadata_ring[:1]
+    slots = [session.acquire() for _ in range(3)]
+    pending = []
+    for step in range(8):
+        order = slots[step % 3 :] + slots[: step % 3]
+        terminal = {order[1]}
+        torch.cuda._sleep(2_000_000)
+        ids, tails = session._stage_slot_ids(order, terminal)
+        pending.append((ids.clone(), tails.clone(), order, list(terminal)))
+    session.close()
+    for ids, tails, expected_ids, expected_tails in pending:
+        assert ids.cpu().tolist() == expected_ids
+        assert tails.cpu().tolist() == expected_tails
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_mixed_terminal_reset_does_not_reset_live_slot():
+    decoder = make_decoder(device="cuda", channels=1)
+    session = decoder._ensure_stream_session()
+    first, second = session.acquire(), session.acquire()
+    codes = torch.ones(2, 3, device="cuda")
+    initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first})
+    following = session.step({second: codes, first: codes})
+    torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("explicit_copy_stream", [False, True])
+@torch.no_grad()
+def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream):
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=1)
+    decoder._codec_stream = torch.cuda.Stream()
+    with torch.cuda.stream(decoder._codec_stream):
+        torch.cuda._sleep(20_000_000)
+    codes = torch.ones(6, device="cuda", dtype=torch.long)
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": "live"}],
+        seq_token_counts=[6],
+    )
+    codes.fill_(-100)
+    payload = output.multimodal_outputs
+    assert isinstance(payload, PackedOutputSnapshot)
+    assert payload.producer_event is not None
+    # No explicit consumer wait here: the shared copier must honor readiness.
+    copy_stream = torch.cuda.Stream()
+    with torch.cuda.stream(copy_stream):
+        cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True)
+    decoder.on_requests_finished(["live"])
+    copy_stream.synchronize()
+    torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@torch.no_grad()
+def test_short_terminal_row_crops_each_channel_into_output_slab(channels):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    short, live = session.acquire(), session.acquire()
+    codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")}
+    slab = torch.empty(channels * 7, device="cuda")
+    buffers = {
+        short: slab[: channels * 2].view(channels, 2),
+        live: slab[channels * 2 :].view(channels, 5),
+    }
+    result = session.step(codes, terminal_slots={short}, output_buffers=buffers)
+    for slot, frames, value in [(short, 2, 1), (live, 5, 3)]:
+        expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"])
+@torch.no_grad()
+def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    slots = [session.acquire() for _ in range(4)]
+    slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda")
+    if layout == "gapped":
+        buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)}
+        frames = 6
+    elif layout == "adjacent":
+        buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)}
+        frames = 6
+    else:
+        # Interleaved with other requests' output: the short views must not
+        # overwrite the untouched part of another row/channel.
+        buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)}
+        frames = 3
+    codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)}
+    result = session.step(codes, output_buffers=buffers)
+    for row, slot in enumerate(slots):
+        expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    if layout == "strided":
+        assert torch.all(slab[:, :, 3:] == -1)
+    elif layout == "gapped":
+        assert torch.all(slab[[2, 5]] == -1)
+    session.close()
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Chunked streaming decode must match whole-sequence decode.
+
+``RingKVCache.complete`` writes a whole chunk before attending. With a ring of
+exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its
+own queries should still see; for T > context it evicts part of the chunk
+itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of
+headroom per layer so streaming reproduces the non-streaming decoder.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TRANSFORMER = {
+    "module_type": "Transformer",
+    "d_model": 16,
+    "num_heads": 2,
+    "num_layers": 1,
+    "dim_feedforward": 32,
+    "causal": True,
+    "norm": "layer_norm",
+    "positional_embedding": "rope",
+    "max_period": 10000,
+    "gating": "none",
+    "layer_scale": 1.0,
+    "conv_layout": True,
+}
+
+
+def tiny_codec() -> MossAudioTokenizerModel:
+    # 64 Hz mono input, 4-sample code frames (16 code frames per second).
+    # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer
+    # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12
+    # tokens at the second transformer, more than its 8-token ring.
+    config = MossAudioTokenizerConfig(
+        sampling_rate=64,
+        downsample_rate=4,
+        number_channels=1,
+        enable_channel_interleave=False,
+        encoder_kwargs=[
+            {"module_type": "PatchedPretransform", "patch_size": 4},
+            {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5},
+        ],
+        decoder_kwargs=[
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+        ],
+        quantizer_kwargs={
+            "input_dim": 8,
+            "rvq_dim": 8,
+            "output_dim": 8,
+            "num_quantizers": 2,
+            "codebook_size": 16,
+            "codebook_dim": 4,
+            "quantizer_type": "rlfq",
+        },
+    )
+    torch.manual_seed(0)
+    model = MossAudioTokenizerModel(config).eval()
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn_like(parameter) * 0.5)
+    model.quantizer.build_decode_lut(2, dtype=torch.float32)
+    return model
+
+
+def stream_decode(model, codes, chunk_frames, *, headroom):
+    n_q, _, total = codes.shape
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0)
+    try:
+        slot = torch.zeros(1, dtype=torch.long)
+        valid = torch.ones(1, dtype=torch.bool)
+        parts = []
+        for start in range(0, total, chunk_frames):
+            chunk = codes[:, :, start : start + chunk_frames]
+            lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long)
+            parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio)
+        return torch.cat(parts, dim=-1)
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_ring_headroom_sizes_each_layer_for_one_chunk():
+    model = tiny_codec()
+    transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")]
+    assert [m.tokens_per_input_frame for m in transformers] == [1, 2]
+    assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8]
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8 + 6, 8 + 12]
+    finally:
+        model.close_decoder_state_pool()
+    model.initialize_decoder_state_pool(1, 0)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8, 8]
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_chunked_streaming_matches_whole_sequence_only_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(1)
+    codes = torch.randint(0, 16, (2, 1, 18))
+    reference = model._decode_frame(codes).audio
+    with_headroom = stream_decode(model, codes, 6, headroom=True)
+    legacy = stream_decode(model, codes, 6, headroom=False)
+    assert with_headroom.shape == reference.shape == legacy.shape
+    torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5)
+    # The legacy context-sized ring lets the 12-token chunk evict keys inside
+    # its own window, so streaming diverges from whole-sequence decoding.
+    assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3)
+
+
+@torch.no_grad()
+def test_padded_terminal_tail_matches_exact_tail_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(2)
+    codes = torch.randint(0, 16, (2, 1, 14))
+    history, tail = codes[:, :, :12], codes[:, :, 12:]
+    padded_tail = torch.nn.functional.pad(tail, (0, 4))  # 2 real frames padded to a 6-frame step
+    outputs = {}
+    for name, last in [("exact", tail), ("padded", padded_tail)]:
+        model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+        try:
+            slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool)
+            for chunk in (history[:, :, :6], history[:, :, 6:]):
+                model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid)
+            audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio
+            outputs[name] = audio[..., : 2 * model.downsample_rate]
+        finally:
+            model.close_decoder_state_pool()
+    torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5)
@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The first-chunk fast path emits audio out of band; the main path resumes the stream."""
+
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+pytestmark = pytest.mark.core_model
+
+N_VQ, SAMPLES, CHANNELS = 2, 4, 2
+
+
+class _Session:
+    def __init__(self, capacity):
+        self.free = list(range(capacity))
+        self.reset_waits = []
+
+    def acquire(self, *, oldest=False):
+        return self.free.pop(0 if oldest else -1) if self.free else None
+
+    def order_after_reset(self, slot, stream):
+        self.reset_waits.append(slot)
+
+    def release(self, slot, **kwargs):
+        self.free.append(slot)
+
+
+class _Wrapper:
+    batch_sizes = [1, 2, 4, 8]
+
+    def __init__(self, device):
+        self.state = torch.zeros(16, device=device)
+        self.calls = []
+        self.gate = threading.Event()
+        self.gate.set()
+
+    def decode(self, codes, slots):
+        self.gate.wait()
+        self.calls.append(int(codes.shape[1]))
+        base = codes.sum(0)[:, 0].float() + 100 * slots.float()
+        channel = torch.arange(CHANNELS, device=codes.device).float()
+        audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device)
+        self.state[slots] = 1
+        return audio.to(torch.bfloat16), None, int(codes.shape[1])
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda")
+
+
+def _fast(device, capacity=4):
+    session, wrapper = _Session(capacity), _Wrapper(device)
+    fast = MossFirstChunkFastPath(
+        session,
+        wrapper,
+        n_vq=N_VQ,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=SAMPLES,
+        n_channels=CHANNELS,
+        sample_rate=torch.tensor(24000, dtype=torch.int32),
+        device=device,
+    )
+    return fast, session, wrapper
+
+
+class _CallbackSink:
+    def __init__(self, callback):
+        self.callback = callback
+
+    def prepare(self, request_ids):
+        callback = self.callback
+
+        class Delivery:
+            routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids}
+
+            def __call__(self, ids, rows, sr):
+                for rid, row in zip(ids, rows, strict=True):
+                    callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr})
+
+            def fail(self, ids):
+                pass
+
+        return Delivery()
+
+
+@pytest.mark.cuda
+def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda):
+    fast, session, wrapper = _fast(cuda)
+    emitted = {}
+    done = threading.Event()
+
+    def sink(client, req_id, payload):
+        emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"]))
+        if len(emitted) == 2:
+            done.set()
+
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # unbound: main path owns it
+    fast.bind(_CallbackSink(sink))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots)
+    assert fast.submit("r1", "k1", [5, 7], slots)
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # already has a slot
+    assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots)  # not a 1-frame chunk
+    assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3]
+    wrapper.gate.set()
+    assert done.wait(10)
+    fast.close()
+
+    assert wrapper.calls in ([1, 1], [2])
+    for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)):
+        got_client, wav, sr = emitted[req_id]
+        assert got_client == client and sr == 24000 and wav.device.type == "cpu"
+        expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES)
+        torch.testing.assert_close(wav, expected)
+    assert sorted(session.reset_waits) == [0, 1]
+
+    # The main path consumes the marker once and is ordered after the decode.
+    assert fast.take_decoded("k0") and not fast.take_decoded("k0")
+    fast.order_after(0)
+    fast.order_after(0)  # no pending handoff: no-op
+    assert wrapper.state[:2].tolist() == [1, 1]
+
+
+@pytest.mark.cuda
+def test_order_after_blocks_until_decode_finishes(cuda):
+    fast, _, wrapper = _fast(cuda)
+    fast.bind(_CallbackSink(lambda *args: None))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", [1, 2], {})
+    finished = threading.Event()
+
+    def wait_for_slot():
+        fast.order_after(0)
+        finished.set()
+
+    waiter = threading.Thread(target=wait_for_slot)
+    waiter.start()
+    assert not finished.wait(0.2)
+    wrapper.gate.set()
+    assert finished.wait(10)
+    waiter.join()
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_no_free_slot_leaves_chunk_to_main_path(cuda):
+    fast, _, _ = _fast(cuda, capacity=0)
+    fast.bind(_CallbackSink(lambda *args: None))
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)
+    assert slots == {} and not fast.take_decoded("k0")
+    fast.close()
+
+
+class _FakeFast:
+    frames = 1
+
+    def __init__(self, decoded, unsent=None):
+        self.decoded = set(decoded)
+        self.unsent = dict(unsent or {})
+        self.ordered = []
+        self.forgotten = []
+
+    def take_decoded(self, key):
+        if key in self.decoded:
+            self.decoded.remove(key)
+            return True
+        return False
+
+    def get_request_slot(self, key, slots):
+        return slots.get(key)
+
+    def order_after(self, slot):
+        self.ordered.append(slot)
+
+    def take_unsent_audio(self, key):
+        return self.unsent.pop(key, None)
+
+    def forget(self, key):
+        self.forgotten.append(key)
+
+    def gate(self):
+        self.gated = getattr(self, "gated", 0) + 1
+
+
+def _decoder(s, fast):
+    decoder = object.__new__(MossTTSCodecDecoder)
+    nn.Module.__init__(decoder)
+    decoder._stream_max_step_frames = 15
+    decoder._stream_state_capacity = 8
+    decoder._stream_req_slots = {}
+    decoder._ensure_stream_session = lambda: s
+    decoder._first_chunk_fast_path = fast
+    decoder._stream_session = s
+    decoder._codec_stream = None
+    return decoder
+
+
+@pytest.mark.cpu
+def test_main_path_skips_fast_decoded_chunk_and_resumes_state():
+    s = cpu_session()
+    reference = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    # The fast path leased slot 0 and decoded frame 0 of "a" into it.
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+    ref_slot = reference.acquire()
+    reference.step({ref_slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)])
+    assert 0 not in out and out[1].tolist() == [[2.0]]
+    assert fast.ordered == [slot]
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)])
+    expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot]
+    torch.testing.assert_close(out[0], expected)
+    assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot]
+
+
+@pytest.mark.cpu
+def test_abort_orders_release_after_fast_decode():
+    s = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True)
+    decoder.on_requests_finished({"a"})
+    assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {}
+    assert slot in s._free_stream_slots
+
+
+@pytest.mark.cpu
+def test_main_path_returns_unsent_first_chunk_without_decoding_twice():
+    s = cpu_session()
+    first = torch.full((1, 1), 7.0)
+    fast = _FakeFast({"a"}, {"a": first})
+    decoder = _decoder(s, fast)
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)])
+    torch.testing.assert_close(out[0], first)
+    assert fast.unsent == {} and fast.ordered == [slot]
+
+    # A payload containing a later frame must concatenate it after the
+    # preserved first samples, while advancing state only for the later frame.
+    first = torch.full((1, 1), 7.0)
+    fast.decoded.add("b")
+    fast.unsent["b"] = first
+    slot_b = s.acquire(oldest=True)
+    decoder._stream_req_slots["b"] = slot_b
+    s.step({slot_b: torch.full((2, 1), 3.0)})
+    out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)])
+    torch.testing.assert_close(out[0][..., :1], first)
+    assert out[0].shape[-1] == 2
+
+
+@pytest.mark.cpu
+def test_oldest_acquire_reuses_least_recently_released_slot():
+    s = cpu_session()
+    a, b = s.acquire(), s.acquire()
+    s.release(a)
+    s.release(b)
+    assert s.acquire() == b
+    assert s.acquire(oldest=True) == 7  # never-used slots were freed before any release
+    assert s._free_stream_slots[-1] == a
+
+
+@pytest.mark.cuda
+def test_failed_decode_reports_error_without_replaying_state(cuda):
+    fast, _, wrapper = _fast(cuda)
+
+    def fail(codes, slots):
+        raise RuntimeError("boom")
+
+    wrapper.decode = fail
+    fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected")))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    with pytest.raises(RuntimeError, match="first-chunk decode failed"):
+        fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_failed_output_handoff_retains_decoded_samples(cuda):
+    fast, _, _ = _fast(cuda)
+
+    def reject(*args):
+        raise RuntimeError("output queue closed")
+
+    fast.bind(_CallbackSink(reject))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    wav = fast.take_unsent_audio("k0")
+    assert wav is not None and wav.shape == (CHANNELS, SAMPLES)
+    assert fast.take_unsent_audio("k0") is None
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_gate_orders_main_stream_after_inflight_fast_decode(cuda):
+    fast, _, _ = _fast(cuda)
+    side = torch.cuda.Stream()
+    with torch.cuda.stream(side):
+        torch.cuda._sleep(200_000_000)
+        pending = torch.cuda.Event()
+        pending.record(side)
+    fast._inflight = pending
+    fast.gate()  # disabled: no wait
+    unordered = torch.cuda.Event()
+    unordered.record()
+    fast._gate_main = True
+    fast.gate()
+    ordered = torch.cuda.Event()
+    ordered.record()
+    unordered.synchronize()
+    assert not pending.query() and not ordered.query()
+    ordered.synchronize()
+    assert pending.query()
+
+
+def _cpu_fast(monkeypatch, *, scheduler=None):
+    # Replace only CUDA allocations, retaining the real lifecycle and worker.
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor)
+    monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: None)
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None))
+    session = _Session(2)
+    fast = MossFirstChunkFastPath(
+        session,
+        SimpleNamespace(batch_sizes=[1, 2]),
+        n_vq=2,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=4,
+        n_channels=1,
+        sample_rate=torch.tensor(24000),
+        device=torch.device("cpu:0"),
+        handoff_timeout_s=0.01,
+    )
+    outputs: queue.Queue = queue.Queue()
+    scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False)
+    return fast, session, sink, outputs
+
+
+def _cancel_during_admission(fast, session, scheduler, sink):
+    entered, release = threading.Event(), threading.Event()
+
+    class PausedSink:
+        def prepare(self, ids):
+            delivery = sink.prepare(ids)
+            if delivery.routes:
+                entered.set()
+                assert release.wait(5)
+            return delivery
+
+    fast.bind(PausedSink())
+    decoder = _decoder(session, fast)
+    errors = []
+    finished = threading.Event()
+    finishing = threading.Event()
+
+    def submit():
+        try:
+            assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+        except BaseException as error:
+            errors.append(error)
+
+    def cancel():
+        try:
+            finishing.set()
+            decoder.on_requests_finished({"r"})
+            finished.set()
+        except BaseException as error:
+            errors.append(error)
+
+    producer = threading.Thread(target=submit)
+    cleaner = threading.Thread(target=cancel)
+    producer.start()
+    try:
+        assert entered.wait(5)
+        assert decoder._stream_req_slots == {}
+        # Scheduler cancellation retires the route before the runner hook.
+        scheduler.requests.pop("r")
+        cleaner.start()
+        assert finishing.wait(5)
+        assert not finished.wait(0.1)
+        release.set()
+        producer.join(5)
+        cleaner.join(5)
+        assert not producer.is_alive() and not cleaner.is_alive()
+        assert not errors, errors
+        assert finished.is_set()
+        assert decoder._stream_req_slots == {}
+        assert fast._decoded == {} and fast._handoffs == {}
+        assert sorted(session.free) == [0, 1]
+        assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+
+        # A fresh request can use and return the same capacity immediately.
+        scheduler.requests["reuse"] = SimpleNamespace(client_index=3)
+        release.set()
+        assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots)
+        scheduler.requests.pop("reuse")
+        decoder.on_requests_finished({"reuse"})
+        assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1]
+    finally:
+        release.set()
+        producer.join(5)
+        if cleaner.ident is not None:
+            cleaner.join(5)
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch):
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler)
+
+    # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode.
+    def decode(jobs):
+        for job in jobs:
+            job.handoff.done.set()
+
+    monkeypatch.setattr(fast, "_decode", decode)
+    fast._handoff_timeout_s = 5.0
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cuda
+def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda):
+    fast, session, _ = _fast(cuda, capacity=2)
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False)
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cpu
+def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch):
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    handoff = _SlotHandoff()
+    fast._handoffs[0] = handoff
+    for _ in range(2):
+        with pytest.raises(TimeoutError, match="slot 0"):
+            fast.order_after(0)
+        assert fast._handoffs[0] is handoff
+    handoff.done.set()
+    fast.order_after(0)
+    assert fast._handoffs == {}
+
+
+@pytest.mark.cpu
+def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch):
+    fast, session, sink, _ = _cpu_fast(monkeypatch)
+    fast.bind(sink)
+    fast.close()
+    slots: dict[str, int] = {}
+    assert not fast.submit("r", "external", [1, 2], slots)
+    assert slots == {} and session.free == [0, 1]
+    with pytest.raises(RuntimeError, match="closed"):
+        fast.bind(sink)
+
+
+@pytest.mark.cpu
+def test_missing_route_does_not_claim_slot(monkeypatch):
+    fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={}))
+    fast.bind(sink)
+    try:
+        assert not fast.submit("gone", "external", [1, 2], {})
+        assert session.free == [0, 1] and outputs.empty()
+    finally:
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch):
+    from vllm.v1.engine import FinishReason
+
+    fast, _, sink, outputs = _cpu_fast(monkeypatch)
+    entered, release = threading.Event(), threading.Event()
+
+    def fail(jobs):
+        entered.set()
+        assert release.wait(2)
+        raise RuntimeError("partial replay failed")
+
+    fast._decode = fail
+    fast.bind(sink)
+    slots: dict[str, int] = {}
+    try:
+        assert fast.submit("r", "a", [1, 2], slots)
+        assert entered.wait(2)
+        assert fast.submit("r", "b", [3, 4], slots)
+        release.set()
+        fast._thread.join(2)
+        assert not fast._thread.is_alive()
+        for slot in slots.values():
+            with pytest.raises(RuntimeError, match="decode failed"):
+                fast.order_after(slot)
+        assert all(fast.take_decoded(key) for key in slots)
+        assert not fast.submit("r", "c", [1, 2], slots)
+        assert outputs.qsize() == 2
+        while not outputs.empty():
+            assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR
+    finally:
+        release.set()
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_platform_failure_wakes_pending_handoff(monkeypatch):
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    fast._handoffs[0] = _SlotHandoff()
+
+    def fail(device):
+        raise RuntimeError("device unavailable")
+
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail))
+    fast._run()
+    assert fast._closed
+    with pytest.raises(RuntimeError, match="decode failed"):
+        fast.order_after(0)
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-request code rows keep the runner payload layout with one host copy."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_rows_map_to_batch_positions_and_copy_once():
+    codes = torch.arange(6, dtype=torch.long).reshape(2, 3)
+    snapshot = _CodeRowsSnapshot(codes, [2, 0], 4)
+    assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None
+    rows = snapshot["codes"]["audio"]
+    assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)]
+    assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:])
+
+    copies = []
+
+    def copy(tensor):
+        copies.append(tensor)
+        return tensor.clone()
+
+    host = snapshot.copy_to_cpu(copy)["codes"]["audio"]
+    assert len(copies) == 1 and copies[0] is codes
+    codes.zero_()
+    assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]]
+    assert host[1].numel() == 0 and host[1].dtype == torch.long
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+def test_explicit_generator_bypasses_compiled_audio_sampler():
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval()
+    heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16()
+    embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16()
+    binary = torch.nn.Linear(80, 2).cuda().bfloat16()
+    hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16)
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("compiled path selected")
+
+    model._compiled_audio_sampler = forbidden
+    generator.manual_seed(17)
+    actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+    assert all(torch.equal(a, b) for a, b in zip(expected, actual))
+    with pytest.raises(AssertionError, match="compiled path selected"):
+        model.generate_frame(hidden, heads, embeddings, binary, n_vq=2)
@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Determined text tokens retain the upstream sampling/counting contract."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.states import RequestState
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("temperature", [0.0, 1.7])
+def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    for i, length in enumerate([4, 1, 3]):
+        reqs.add_request(str(i), length, [2] * length, 0, 8)
+        slot = reqs.req_id_to_index[str(i)]
+        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+    reqs.apply_staged_writes()
+    sampler.apply_staged_writes()
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    state.model._batch_state = None
+    batch = _batch(device, [2, 4, 3], [1, 2, 1])
+    batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
+    batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
+    batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
+    batch.expanded_idx_mapping = batch.idx_mapping
+    batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device)
+    batch.logits_indices = batch.query_start_loc[1:].long() - 1
+    batch.positions = torch.arange(batch.num_tokens, device=device)
+    state.model._batch_should_continue = torch.tensor([False, True, True], device=device)
+    with torch.inference_mode():
+        logits = state.model.compute_logits(torch.ones(3, 4, device=device))
+        expected = sampler(logits, batch)
+        actual = state.sample_determined_tokens(batch, sampler)
+    assert actual is not None
+    for field in ["sampled_token_ids", "num_sampled", "num_rejected"]:
+        torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0)
+    assert actual.num_sampled.tolist() == [1, 0, 1]
+    snapshot = actual.sampled_token_ids.clone()
+    state.model._batch_should_continue.logical_not_()
+    torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    "constraint",
+    [
+        {"min_tokens": 2, "stop_token_ids": [1]},
+        {"allowed_token_ids": [1, 2]},
+        {"logit_bias": {1: -10.0}},
+    ],
+)
+def test_real_sampler_constraints_disable_determined_tokens(constraint):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    reqs.add_request("constrained", 1, [2], 0, 8)
+    slot = reqs.req_id_to_index["constrained"]
+    sampler.add_request(slot, 1, SamplingParams(**constraint))
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    # Use the actual upstream request registration, including its combined
+    # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves.
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
+def test_distribution_features_fall_back_before_using_gpu(feature):
+    sampler = object.__new__(Sampler)
+    sampler.compute_nans = feature == "nans"
+    sampler.return_sampling_mask = feature == "mask"
+    sampler.trace_replay_state = object() if feature == "trace" else None
+    sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
+    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
+    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
+    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
+    sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
+    state = object.__new__(MossLocalModelState)
+    state._direct_tokens = True
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"])
+def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature):
+    hook = mocker.Mock(return_value=None)
+
+    class State:
+        def sample_determined_tokens(self, batch, sampler):
+            return hook(batch, sampler)
+
+    runner = object.__new__(OmniARModelRunner)
+    runner.model_state = State()
+    runner.sampler = object()
+    runner.batch_sharder = object() if feature == "sharding" else None
+    runner.pp_handler = object() if feature == "pipeline" else None
+    batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1)
+    grammar = object() if feature == "grammar" else None
+    fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream")
+    hidden = object()
+    assert runner.sample(hidden, batch, grammar) == "upstream"
+    fallback.assert_called_once_with(hidden, batch, grammar)
+    assert hook.call_count == int(feature == "declined")
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eager MTP emits the same frames one step earlier with identical backbone inputs."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _run(state, device, schedule, stop_step):
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    embeds = []
+    frames: dict[int, list[torch.Tensor]] = {3: [], 0: []}
+    for index, (slots, counts) in enumerate(schedule):
+        state.model.force_stop = index == stop_step
+        embed, payload, _ = _step(state, _batch(device, slots, counts), req_states)
+        embeds.append(embed)
+        rows = payload.get("codes", {}).get("audio", [])
+        for slot, row in zip(slots, rows):
+            if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()):
+                frames[slot].append((index, row.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+    return embeds, frames
+
+
+def test_eager_frames_match_canonical_one_step_earlier(device):
+    schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])]
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 3, "long", 17)
+        _admit(state, 0, "short", 17)
+    # Stop both streams at the canonical step that would draw the 4th step's frames.
+    c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4)
+    e_embeds, e_frames = _run(eager, device, schedule, stop_step=3)
+
+    for c, e in zip(c_embeds, e_embeds):
+        torch.testing.assert_close(c, e, rtol=0, atol=0)
+    for slot in (3, 0):
+        assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot])
+        for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]):
+            assert ei == ci - 1
+            torch.testing.assert_close(cf, ef, rtol=0, atol=0)
+
+
+def test_stopped_stream_does_not_emit_again(device):
+    state = _state(MossLocalModelState, device)
+    state._local_eager_mtp = True
+    _admit(state, 0, "short", 17)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    emitted = []
+    for index, count in enumerate([2, 1, 1, 1]):
+        state.model.force_stop = index == 1
+        _, payload, _ = _step(state, _batch(device, [0], [count]), req_states)
+        row = payload["codes"]["audio"][0]
+        emitted.append(bool(row.ne(state.model.audio_pad_token_id).any()))
+        computed[0] += count
+        state.model.force_stop = False
+    # Frame at the prefill step, stop at step 1, then asynchronous overrun steps.
+    assert emitted == [True, False, False, False]
+
+
+def test_unmarked_decode_row_ignores_stale_slot_state(device):
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 0, "short", 17)
+    # A decode row whose completing prefill was not observed in this state.
+    eager._keep_pool[0] = False
+    eager._eager_emb_pool[0].fill_(3.0)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0], computed[0] = 2, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    for state in (canonical, eager):
+        state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False}
+    c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states)
+    e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states)
+    torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0)
+    torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker):
+    """The Local config must not enable the parent's incompatible eager hook."""
+    from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+    from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+
+    device = torch.device("cpu")
+    template = _state(MossLocalModelState, device)
+    model = template.model
+    model.config = SimpleNamespace(mrv2_eager_mtp=True)
+
+    def initialize_parent(state, *args):
+        state.__dict__.update(template.__dict__)
+        state._eager_mtp, state._eager_rows = False, None
+        state._eager_state = EagerMTPState(state)
+
+    monkeypatch.setattr(OmniModelState, "__init__", initialize_parent)
+    state = MossLocalModelState(SimpleNamespace(), model, None, device)
+    _admit(state, 0, "req", 17)
+    batch = _batch(device, [0], [2])
+    batch.req_ids = ["req"]
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed))
+    hook = mocker.spy(state, "run_eager_mtp")
+
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.execute_model_state = SimpleNamespace(
+        input_batch=batch,
+        hidden_states=torch.zeros(2, 4),
+        finished_req_ids=set(),
+        ec_connector_output=None,
+        routed_experts=None,
+    )
+    runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
+    runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.model_config = SimpleNamespace(async_chunk=False)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
+    runner.model_state, runner.model = state, model
+    runner.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)),
+        num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)),
+        prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)),
+    )
+    runner.main_stream = runner.output_copy_stream = mocker.Mock()
+    runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock()
+    runner.sampler = None
+    runner.sample = mocker.Mock(
+        return_value=(
+            SimpleNamespace(sampled_token_ids=torch.tensor([[2]])),
+            torch.ones(1),
+            torch.zeros(1),
+        )
+    )
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
+    runner.postprocess_sampled = mocker.Mock()
+    runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
+    mock_out = SimpleNamespace(copy_event=None)
+    monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out))
+
+    assert runner.sample_tokens(None) is mock_out
+    hook.assert_called_once()
+    assert state._eager_mtp is False and state._eager_rows is None
+    assert state._local_eager_mtp is True
@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare slot lifecycle with canonical Local hooks, including output ownership."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker_v2.model_states import init_omni_model_state
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _model(device):
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    model.model = torch.nn.Module()
+    weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight)
+    model.hidden_size, model.n_vq = 4, 2
+    model.audio_pad_token_id, model.text_vocab_size = 8, 8
+    model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")}
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8
+    model.force_stop = False
+
+    def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs):
+        assert (temperature, top_k, top_p) == (1.7, 25, 0.8)
+        noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
+        codes = (hidden[:, :2].mul(16).long() + noise) % 8
+        return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes
+
+    model.local_transformer = SimpleNamespace(generate_frame=frame)
+    return model
+
+
+def _state(cls, device):
+    state = object.__new__(cls)
+    state.model = _model(device)
+    state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype
+    state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True
+    state.scheduler_config = SimpleNamespace(max_num_seqs=5)
+    state.vllm_config = SimpleNamespace()
+    state.intermediate_buffer = OmniIntermediateBuffer(5)
+    state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype)
+    state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long)
+    state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long)
+    for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"):
+        setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype))
+    state._mtp_runner = state._mtp_sample_uniforms = None
+    state._mtp_generators = {}
+    state._stream_pos = {}
+    state._stream_decode_event = None
+    state._eager_state = EagerMTPState(state)
+    if cls is MossLocalModelState:
+        state._init_slot_buffers(5, 4, device, state.dtype)
+    return state
+
+
+def _admit(state, slot, name, seed):
+    params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None)
+    with patch.object(DefaultModelState, "add_request", return_value=None):
+        state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params))
+    state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])}
+
+
+def _batch(device, slots, counts):
+    starts = np.array([0, *np.cumsum(counts)], dtype=np.int32)
+    return SimpleNamespace(
+        idx_mapping_np=np.array(slots),
+        idx_mapping=torch.tensor(slots, device=device),
+        num_reqs=len(slots),
+        num_tokens=int(starts[-1]),
+        num_scheduled_tokens=np.array(counts),
+        query_start_loc_np=starts,
+        query_start_loc=torch.tensor(starts, device=device),
+        # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs.
+        input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device),
+    )
+
+
+def _step(state, batch, req_states, dispatcher=None):
+    inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]}
+    with (
+        patch("vllm.forward_context.set_forward_context", return_value=nullcontext()),
+        patch(
+            "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext()
+        ),
+        torch.inference_mode(),
+    ):
+        state.run_preprocess(batch, inputs, req_states, dispatcher)
+        hidden = inputs["inputs_embeds"] + 0.25
+        state.run_postprocess(hidden, batch)
+        _, payload = state.postprocess_model_output(hidden, batch, req_states)
+        last = batch.query_start_loc[1:] - 1
+        tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1)
+        result = inputs["inputs_embeds"].clone(), payload, tokens
+        hidden.fill_(-100)  # next graph replay can overwrite its output immediately
+    return result
+
+
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("batch_prefill", [False, True])
+def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill):
+    reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState))
+    candidate._batch_prefill = batch_prefill
+    for state in (reference, candidate):
+        _admit(state, 3, "long", seed)
+        _admit(state, 0, "short", seed)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    saved = []
+    for index, (slots, counts) in enumerate(
+        [
+            ([3, 0], [2, 2]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+        ]
+    ):
+        if index == 4:
+            # Cancellation/removal followed by reuse must not retain hidden,
+            # reference offset, control, or the old request's RNG stream.
+            for state in (reference, candidate):
+                state.remove_request("short")
+                assert "short" not in state._mtp_generators
+                _admit(state, 0, "replacement", seed)
+            computed[0], prompt[0] = 0, 1
+        outputs = []
+        for state in (reference, candidate):
+            state.model.force_stop = index == 3
+            torch.manual_seed(91 + index)
+            outputs.append(_step(state, _batch(device, slots, counts), req_states))
+        a, b = outputs
+        torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
+        torch.testing.assert_close(a[2], b[2], rtol=0, atol=0)
+        if index == 3:
+            assert b[2].tolist() == [3, 3]
+        assert a[1].keys() == b[1].keys()
+        for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(ar, br, rtol=0, atol=0)
+            saved.append((br, br.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+            old = reference.intermediate_buffer.buffers[slot]
+            new = candidate.intermediate_buffer.buffers[slot]
+            assert old["ref_offset"] == new["ref_offset"]
+            torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0)
+            assert "hidden_states" not in new and "audio_codes" not in new
+    candidate._codes_pool.fill_(8)
+    candidate._hidden_pool.zero_()
+    for current, snapshot in saved:
+        torch.testing.assert_close(current, snapshot, rtol=0, atol=0)
+
+
+def test_slot_honors_external_stopping_control(device):
+    states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+        _admit(state, 0, "stop", None)
+        req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5))
+        _step(state, _batch(device, [0], [1]), req)
+        state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True
+        req.num_computed_tokens[:] = 1
+        output = _step(state, _batch(device, [0], [1]), req)
+        assert output[2].tolist() == [3]
+        assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]]
+
+
+@pytest.mark.cuda
+def test_slot_mtp_graph_padding_reorder_and_owned_output():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+
+        def frame(hidden, *_args, **_kwargs):
+            return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8
+
+        state.model.local_transformer.generate_frame = frame
+        for slot in [0, 3, 4]:
+            _admit(state, slot, f"req{slot}", None)
+    candidate = states[1]
+    buffers = [
+        candidate._mtp_input_ids[:4],
+        candidate._mtp_input_embeds[:4],
+        candidate._mtp_hidden[:4],
+        candidate._mtp_text_step[:4],
+    ]
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream), torch.inference_mode():
+        for _ in range(3):
+            candidate.model.mtp(*buffers)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph), torch.inference_mode():
+        graph_output = candidate.model.mtp(*buffers)
+
+    def replay(*args, **kwargs):
+        assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers))
+        assert args[0].shape[0] == 4
+        graph.replay()
+        return graph_output
+
+    candidate._mtp_runner = replay
+    candidate._is_mtp_graph_runner = lambda: True
+    from vllm.config import CUDAGraphMode
+
+    def dispatcher(bsz):
+        return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL)
+
+    req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32))
+    retained = []
+    for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]:
+        out = [
+            _step(
+                state,
+                _batch(torch.device("cuda"), slots, [1] * len(slots)),
+                req,
+                dispatcher if state is candidate else None,
+            )
+            for state in states
+        ]
+        torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0)
+        torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0)
+        for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+            retained.append((b, b.clone()))
+        req.num_computed_tokens[slots] += 1
+    with torch.inference_mode():
+        for tensor in graph_output:
+            tensor.fill_(7)
+    for tensor, copy in retained:
+        torch.testing.assert_close(tensor, copy, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None)
+    assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState)
+    dynamic.create_mrv2_model_state.assert_not_called()
+
+    class Invalid:
+        has_preprocess = True
+
+        def create_mrv2_model_state(self, *args):
+            return object()
+
+    with pytest.raises(TypeError, match="must return a ModelState"):
+        init_omni_model_state(None, Invalid(), None, torch.device("cpu"))
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled", [False, True])
+def test_local_factory_default_and_option(monkeypatch, enabled):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None)
+    model = _model(torch.device("cpu"))
+    model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled)
+    state = init_omni_model_state(None, model, None, torch.device("cpu"))
+    assert type(state) is (MossLocalModelState if enabled else OmniModelState)
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference batching preserves coordinates and asynchronous source ownership."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("flat", [False, True])
+def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = True
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    for slot, offset in [(3, 1), (0, 2)]:
+        state.intermediate_buffer.buffers[slot] = {
+            "req_id": str(slot),
+            "codes": {"ref": codes.flatten() if flat else codes},
+            "ref_offset": offset,
+        }
+    batch = _batch(torch.device("cpu"), [3, 0], [2, 1])
+    req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    embeds = state._static_inputs_embeds[:3]
+    expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]])
+    embed = mocker.spy(state.model.model.embed_tokens, "forward")
+    with torch.inference_mode():
+        state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+    torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
+    assert embed.call_count == 1
+    assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3
+
+
+@pytest.mark.cuda
+def test_pending_reference_uploads_survive_staging_reuse_and_overflow():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    outputs = []
+    with torch.cuda.stream(stream), torch.inference_mode():
+        # Keep the first H2D pending while CPU prepares all subsequent batches.
+        torch.cuda._sleep(100_000_000)
+        for value in range(1, 8):
+            embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype)
+            codes = torch.full((2, 2), value)
+            state._apply_reference_batch(embeds, [(1, codes)])
+            outputs.append((embeds, value))
+            codes.fill_(-1)
+    stream.synchronize()
+    assert len(state._prefill_staging) <= 3
+    for output, value in outputs:
+        expected = torch.zeros_like(output)
+        expected[1:3] = value / 8
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    model = state.model
+    model.n_vq = 12
+    model.audio_vocab_size = model.audio_pad_token_id = 1024
+    model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16)
+    model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model)
+    state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16)
+    references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]]
+    with torch.inference_mode():
+        actual = torch.randn_like(state._static_inputs_embeds)
+        expected = actual.clone()
+        for start, codes in references:
+            expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda"))
+        state._apply_reference_batch(actual, references)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)])
+def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim):
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    torch.manual_seed(43)
+    query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    for position in range(12):
+        # Non-contiguous prefixes retain the stride of the full 12-slot cache.
+        k, v = key[:, :, : position + 1], value[:, :, : position + 1]
+        before_k, before_v = key.clone(), value.clone()
+        expected = F.scaled_dot_product_attention(query, k, v)
+        actual = local_short_attention(query, k, v)
+        assert torch.equal(key, before_k) and torch.equal(value, before_v)
+        if position == 0:
+            assert torch.equal(actual, v)
+        torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01)
+        # Future cache contents must never affect an earlier position.
+        key[:, :, position + 1 :] += 100
+        value[:, :, position + 1 :] -= 100
+        assert torch.equal(local_short_attention(query, k, v), actual)
+        key.copy_(before_k)
+        value.copy_(before_v)
+
+
+def test_compile_and_graph_replay_reads_fresh_inputs():
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    compiled = torch.compile(local_short_attention, fullgraph=True)
+    for _ in range(3):
+        compiled(query, key, value)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        result = compiled(query, key, value)
+    for _ in range(3):
+        query.normal_()
+        key.normal_()
+        value.normal_()
+        graph.replay()
+        torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01)
+
+
+def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch):
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1")
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval()
+    attn = model.h[0].attn
+    attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16)
+    cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True)
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16)
+            # The second frame overwrites the previous frame's K/V before use.
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = forward(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02)
@@ -6,6 +6,7 @@ single-flight, and micro-batched encoding."""
 import asyncio
 import threading
 import time
+from tempfile import TemporaryDirectory

 import pytest
 import torch
@@ -598,3 +599,239 @@ def test_reference_resampler_is_cached_and_matches_functional_resample():
        assert info.misses == 1 and info.hits == 1
    finally:
        _reference_resampler.cache_clear()
+
+
+# --------------------------------------------------------------------------- #
+# Inline data: URI fast path                                                  #
+# --------------------------------------------------------------------------- #
+
+
+class _DigestAudio(_FakeAudio):
+    """Resolver whose key for data: URIs is the SHA-1 digest, like serving."""
+
+    def __init__(self, *, digest_keys=True):
+        super().__init__()
+        self._digest_keys = digest_keys
+
+    async def resolve(self, ref_str):
+        wav_list, sr, _ = await super().resolve(ref_str)
+        if not self._digest_keys:
+            return wav_list, sr, "rk:" + ref_str
+        import hashlib
+
+        return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+    def artifact_key(self, cache_key):
+        return None
+
+
+async def test_inline_reference_skips_repeat_resolution(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio()
+    ref = "data:audio/wav;base64,QUJD"
+    audio.register(ref, 5)
+    enc = make_encoder(proc)
+
+    first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    first.fill_(0)  # callers own their copy
+    second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+
+    assert audio.resolve_calls == [ref]
+    assert key1 == key2
+    assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long))
+
+
+async def test_inline_index_requires_matching_resolve_key(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False)
+    ref = "data:audio/wav;base64,REVG"
+    audio.register(ref, 6)
+    enc = make_encoder(proc)
+
+    for _ in range(2):
+        await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    assert audio.resolve_calls == [ref, ref]
+
+
+async def test_inline_reference_shared_across_encoders(make_encoder, monkeypatch):
+    # Keep the AF_UNIX socket path independent of pytest's nested base path.
+    with TemporaryDirectory(prefix="moss-ref-") as shared_dir:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", shared_dir)
+        ref = "data:audio/wav;base64,R0hJ"
+        first_audio, second_audio = _DigestAudio(), _DigestAudio()
+        first_audio.register(ref, 9)
+        second_audio.register(ref, 9)
+        first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor())
+
+        a, key_a = await first.encode(
+            ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key
+        )
+        b, key_b = await second.encode(
+            ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key
+        )
+
+        assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == []
+        assert key_a == key_b and torch.equal(a, b)
+
+
+# --------------------------------------------------------------------------- #
+# Reference-encoder CUDA graphs: selection and processor-equivalent prep      #
+# --------------------------------------------------------------------------- #
+
+
+class _GraphsStub:
+    def __init__(self, result):
+        self.result = result
+        self.calls: list[list[torch.Tensor]] = []
+
+    def encode(self, wavs):
+        self.calls.append(wavs)
+        return self.result
+
+
+class _LoudnessProcessor(_FakeProcessor):
+    @staticmethod
+    def loudness_normalize(wav):
+        return wav * 2.0
+
+
+async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder):
+    proc = _LoudnessProcessor()
+    enc = make_encoder(proc)
+    graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)]
+    enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes)
+    prepared = [torch.ones(1, 100)]
+    assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == []
+    enc._graphs = _GraphsStub([None])  # no graph fits this clip
+    out = enc._encode_prepared(prepared)
+    assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3
+
+
+async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder):
+    enc = make_encoder(_LoudnessProcessor())
+    mono = enc._processor_prepare(torch.ones(100))
+    assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0))
+    many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4))
+    assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0)
+
+
+async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz, the Local-v1.5 reference rate
+    wav = torch.randn(1, 2400)
+    expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0
+    # The encoder resamples the mono clip before duplicating it. Some CPU conv
+    # kernels round a one-row batch differently from a two-row one, so the two
+    # orders agree to float32 precision, not bit for bit.
+    torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize("role", ["", "host", "client"])
+async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role):
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    enc = make_encoder(_FakeProcessor())
+    enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role)
+    captured, hosts, ready = [], [], []
+    monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True))
+    monkeypatch.setattr(enc, "_host_workers", lambda: [])
+    helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True))
+
+    def make_host(*args, **kwargs):
+        hosts.append(args)
+        return helper
+
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host)
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper)
+    enc.prepare()
+    enc.prepare()
+    # Every API process waits for the shared host before startup completes.
+    assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role]
+    assert len(ready) == (2 if role else 0)
+
+
+async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz
+    sent = []
+
+    def record_prepared(prepared):
+        sent.extend(prepared)
+        return [None] * len(prepared)
+
+    monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared)
+    wav = [float(v) for v in torch.randn(2400)]
+    enc._encode_batch_sync([(wav, _SR)])
+    # Resampled where the request arrived (a shared encoder receives it this way); still one channel.
+    expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000)
+    torch.testing.assert_close(sent[0], expected, rtol=0, atol=0)
+    # The eager path is told the clip is already at the tokenizer rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared(sent)
+    assert calls == [48000]
+
+
+async def test_batcher_runs_up_to_max_inflight_batches_at_once():
+    inside = threading.Barrier(2, timeout=5)
+    seen: list[int] = []
+
+    def encode_batch(payload):
+        inside.wait()  # two batches are encoding at the same time
+        seen.append(len(payload))
+        return [torch.zeros(1) for _ in payload]
+
+    batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2)
+    try:
+        first = asyncio.ensure_future(batcher.submit([1.0], _SR))
+        await asyncio.sleep(0.05)  # the first batch is in flight before the second clip arrives
+        await asyncio.gather(first, batcher.submit([2.0], _SR))
+    finally:
+        await batcher.aclose()
+    assert seen == [1, 1]
+
+
+def test_reference_stream_priority(monkeypatch):
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False)
+    assert re_mod.reference_stream_priority() == 0
+    monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high")
+    expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0
+    assert re_mod.reference_stream_priority() == expected
+
+
+async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False)  # CPU placement
+    enc = make_encoder(proc)
+    enc._gpu_prep = True
+    wav = torch.randn(1, 2400)
+    clip = enc._encoder_input(wav)
+    assert clip.shape == (1, 2400)  # left at the working rate for the encoder's device
+    expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0
+    torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6)
+    # The eager fallback is told the clip's own rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared([clip])
+    assert calls == [24000]
+    assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP"
@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode."""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    MossReferenceEncoderGraphs,
+    _cached_reference_rope,
+    _reference_rope_factors,
+    _validate_unmasked_reference_encoder,
+    _windowed_attention,
+    install_windowed_attention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+FRAME, N_VQ = 8, 3
+
+
+class _Stage(nn.Module):
+    """Causal per-frame stage: frames never mix, like a causal encoder layer."""
+
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x, lengths):
+        return torch.tanh(self.proj(x)), lengths
+
+
+class _Quantizer(nn.Module):
+    def forward(self, hidden, lengths, n_q):
+        # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices.
+        codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)])
+        return hidden, codes, lengths
+
+
+class _Tokenizer(nn.Module):
+    number_channels = 2
+    downsample_rate = FRAME
+    sampling_rate = 64
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)])
+        self.quantizer = _Quantizer()
+
+    def _flatten_channels_for_codec(self, audio, lengths):
+        batch, channels, samples = audio.shape
+        frames = samples // FRAME
+        values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1)
+        return values, lengths // FRAME
+
+    def _codec_inference_autocast(self):
+        return nullcontext()
+
+    def eager(self, wavs):
+        out = []
+        for wav in wavs:
+            values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]]))
+            for module in self.encoder:
+                values, lengths = module(values, lengths)
+            _, codes, _ = self.quantizer(values, lengths, N_VQ)
+            out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long())
+        return out
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda", torch.accelerator.current_device_index())
+
+
+def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda):
+    from types import SimpleNamespace
+
+    module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0))
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    torch.manual_seed(17)
+    # Prime once; all subsequent sequences and batches use an immutable prefix.
+    q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    _cached_reference_rope(module, q, k, cache)
+    stored = next(iter(cache.values()))
+    for batch, length in [(1, 17), (4, 64), (2, 128)]:
+        q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        before = (q.clone(), k.clone())
+        got = _cached_reference_rope(module, q, k, cache)
+        cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda)
+        torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0)
+        torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0)
+        for value, source, saved in zip(got, (q, k), before):
+            real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1)
+            expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype)
+            torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0)
+            torch.testing.assert_close(source, saved, rtol=0, atol=0)
+        assert len(cache) == 1
+        assert next(iter(cache.values())) is stored
+
+
+def test_cached_rope_distinguishes_period_and_head_width(cuda):
+    from types import SimpleNamespace
+
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]:
+        module = SimpleNamespace(rope=SimpleNamespace(max_period=period))
+        q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16)
+        _cached_reference_rope(module, q, q, cache)
+    assert len(cache) == 3
+
+
+@pytest.mark.parametrize("batched_transfer", [False, True])
+def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer):
+    torch.manual_seed(0)
+    tokenizer = _Tokenizer().to(cuda)
+    # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples.
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)]
+    assert graphs.max_batch == 4
+
+    for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]:
+        wavs = [torch.randn(2, n * FRAME) for n in frames]
+        codes = graphs.encode(wavs)
+        expected = tokenizer.eager([w.to(cuda) for w in wavs])
+        assert len(codes) == count
+        for got, want in zip(codes, expected):
+            assert got.device.type == "cpu" and got.dtype == torch.long
+            torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Longer than every bucket: left to the caller; the rest still use graphs.
+    long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME)
+    codes = graphs.encode([long_clip, short_clip])
+    assert codes[0] is None
+    torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0)
+
+    # Wider than every batch: split, longest first, and returned in input order.
+    wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Result ownership includes the degenerate single-frame shape. Reusing
+    # the pinned output slab must not overwrite a previous caller's codes.
+    previous = graphs.encode([torch.randn(2, FRAME)])[0]
+    snapshot = previous.clone()
+    graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)])
+    torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+
+
+def test_batch_audio_budget_limits_captured_graphs(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0
+    )
+    graphs.capture()
+    # 4 x 4 s exceeds the 8 s budget; a single clip is always captured.
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)]
+    wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_finer_singleton_buckets_preserve_multiclip_batching(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer,
+        n_vq=N_VQ,
+        batch_sizes=(1, 2, 4),
+        bucket_seconds=(2.0, 4.0),
+        singleton_bucket_seconds=(1.0, 3.0),
+    )
+    graphs.capture()
+    assert (1, 64) in graphs.captured and (1, 192) in graphs.captured
+    assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured
+    replay = graphs._replay
+    shapes = []
+
+    def record(entry, wavs):
+        shapes.append(tuple(entry.audio.shape))
+        return replay(entry, wavs)
+
+    graphs._replay = record
+    one = torch.randn(2, 20 * FRAME)
+    out = graphs.encode([one])
+    assert shapes == [(1, 2, 192)]
+    torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0)
+    shapes.clear()
+    wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]]
+    out = graphs.encode(wavs)
+    assert shapes == [(4, 2, 256)]
+    for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_compiled_capture_matches_eager(cuda):
+    torch.manual_seed(2)
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)]
+    wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+class _Attention(nn.Module):
+    """The attention surface the tokenizer's MHA exposes to the windowed kernel."""
+
+    def __init__(self, embed_dim=128, heads=2, context=5):
+        super().__init__()
+        self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True
+        self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+
+    def _project_qkv(self, x):
+        b, t, _ = x.shape
+        return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+
+    def _apply_dense_rope(self, q, k):
+        return q, k
+
+    def _forward_non_streaming_sdpa(self, x, input_lengths):
+        # Reference semantics: causal, distance < context, padded rows zeroed.
+        b, t, _ = x.shape
+        q, k, v = self._project_qkv(x)
+        pos = torch.arange(t, device=x.device)
+        delta = pos.view(-1, 1) - pos.view(1, -1)
+        mask = (delta >= 0) & (delta < self.context)
+        out = F.scaled_dot_product_attention(q, k, v, mask[None, None])
+        valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1)
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+        return out.transpose(1, 2).reshape(b, t, self.embed_dim)
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_windowed_attention_matches_masked_sdpa(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(1)
+    model = nn.Module()
+    model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)])
+    attention = model.encoder[0]
+    x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 3], device=cuda)
+    expected = attention._forward_non_streaming_sdpa(x, lengths)
+    assert install_windowed_attention(model, fa_version=fa_version) == 1
+    actual = attention._forward_non_streaming_sdpa(x, lengths)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0
+    # Non-bf16 inputs keep the tokenizer's own SDPA path.
+    x32 = x.float()
+    attention.float()
+    torch.testing.assert_close(
+        attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2
+    )
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(44)
+    attention = _Attention().to(cuda, torch.bfloat16)
+    x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 1, 0], device=cuda)
+    kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version)
+    with torch.no_grad():
+        masked = _windowed_attention(attention, x, lengths, **kwargs)
+        unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs)
+        poisoned = x.clone()
+        for row, length in enumerate([40, 17, 1, 0]):
+            poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30
+        changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs)
+        for row, length in enumerate([40, 17, 1, 0]):
+            torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0)
+            torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0)
+            assert masked[row, length:].count_nonzero() == 0
+        assert unmasked[1, 17:].count_nonzero() > 0
+        # The FP32 fallback retains its original masking contract.
+        attention.float()
+        got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs)
+        expected = attention._forward_non_streaming_sdpa(x.float(), lengths)
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+
+def test_unmasked_queries_reject_unknown_encoder_before_patching():
+    tokenizer = _Tokenizer()
+    with pytest.raises(ValueError, match="patch-downsample"):
+        install_windowed_attention(tokenizer, skip_padded_query_mask=True)
+
+
+@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"])
+def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part):
+    from types import SimpleNamespace
+
+    patch = type("MossAudioTokenizerPatchedPretransform", (), {})()
+    patch.is_downsample = bad_part != "upsample"
+    projected = type("MossAudioTokenizerProjectedTransformer", (), {})()
+    projected.transformer = SimpleNamespace(
+        layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))]
+    )
+    quantizer_cls = type(
+        "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {}
+    )
+    quantizer = quantizer_cls()
+    quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1)
+    tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer)
+    with pytest.raises(ValueError):
+        _validate_unmasked_reference_encoder(tokenizer)
@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""API processes share one reference encoder through a Unix socket."""
+
+import threading
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _codes(wav):
+    # Deterministic per clip: frames from the clip length, values from its first sample.
+    return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long)
+
+
+class _Recorder:
+    def __init__(self, gate=None):
+        self.batches = []
+        self.gate = gate
+
+    def __call__(self, wavs):
+        if self.gate is not None:
+            self.gate.wait(5)
+        self.batches.append(len(wavs))
+        return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs]
+
+
+@pytest.fixture
+def host_dir(monkeypatch):
+    monkeypatch.setattr(shared, "_host_lock_handle", None)
+    # AF_UNIX limits the full socket path, including pytest directory names.
+    with TemporaryDirectory(prefix="moss-ref-") as path:
+        yield path
+
+
+def test_only_one_process_per_directory_hosts(host_dir):
+    assert shared.elect_host(host_dir) is True
+    assert shared.elect_host(host_dir) is False  # a second open file description cannot take the lock
+
+
+def test_close_releases_queued_and_active_callers(host_dir):
+    entered, release = threading.Event(), threading.Event()
+
+    def encode(wavs):
+        entered.set()
+        assert release.wait(5)
+        return [_codes(wav) for wav in wavs]
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)])
+    results: dict[str, list[torch.Tensor | BaseException]] = {}
+    threads = [
+        threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True)
+        for key in ("active", "queued")
+    ]
+    try:
+        threads[0].start()
+        assert entered.wait(2)
+        threads[1].start()
+        with host._ready:
+            assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2)
+        host.close()
+        for thread in threads:
+            thread.join(1)
+            assert not thread.is_alive(), "closed encoder left a request waiting"
+        assert all(isinstance(results[key][0], RuntimeError) for key in results)
+    finally:
+        release.set()
+        host.close()
+        for thread in threads:
+            if thread.ident is not None:
+                thread.join(1)
+
+
+def test_worker_result_count_mismatch_fails_entire_batch(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)])
+    results = []
+    caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True)
+    try:
+        caller.start()
+        caller.join(2)
+        assert not caller.is_alive(), "missing encoder result left a request waiting"
+        assert len(results) == 2
+        assert all(isinstance(result, RuntimeError) for result in results)
+    finally:
+        host.close()
+        caller.join(1)
+
+
+def test_client_round_trip_and_per_clip_errors(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)]
+        out = client.encode(wavs)
+        assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0]))
+        assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1])
+        assert torch.equal(out[2], _codes(wavs[2]))
+        # The host process's own clips skip the socket.
+        assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2]))
+    finally:
+        host.close()
+
+
+def test_host_merges_clips_from_concurrent_clients(host_dir):
+    gate = threading.Event()
+    recorder = _Recorder(gate)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)])
+    try:
+        results = {}
+
+        def send(name, value):
+            results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2)
+
+        # Encoding is held until both clients' clips are queued.
+        warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)]))
+        warm.start()
+        threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))]
+        for thread in threads:
+            thread.start()
+        threading.Timer(0.5, gate.set).start()
+        for thread in (warm, *threads):
+            thread.join(10)
+        # Both clients' clips were encoded together, not one request at a time.
+        assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4
+        assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"])
+    finally:
+        host.close()
+
+
+def test_client_without_host_raises(host_dir, monkeypatch):
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    with pytest.raises((FileNotFoundError, ConnectionRefusedError)):
+        shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)])
+
+
+def test_host_window_collects_across_workers_without_reserving_singletons(host_dir):
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500)
+    try:
+        barrier = threading.Barrier(3)
+        results = []
+
+        def send(value):
+            barrier.wait()
+            results.extend(host.encode([torch.full((1, 20), value)]))
+
+        threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)]
+        for thread in threads:
+            thread.start()
+        barrier.wait()
+        for thread in threads:
+            thread.join(5)
+            assert not thread.is_alive()
+        assert recorder.batches == [2]
+        assert sorted(int(x[0, 0]) for x in results) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch):
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    proc = _FakeProcessor()
+    enc = re_mod.MossReferenceEncoder(
+        proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20)
+    )
+    placed = []
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True))
+    out = enc._encode_batch_sync([([7.0] * 50, _SR)])
+    assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7
+
+
+def test_workers_encode_concurrently_after_their_warmups(host_dir):
+    inside, release = threading.Barrier(3, timeout=5), threading.Event()
+    warmed = []
+
+    def worker(name):
+        def encode(wavs):
+            inside.wait()  # both workers are inside an encode at the same time
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, lambda: warmed.append(name)
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")])
+    try:
+        out = []
+
+        def submit(value):
+            out.append(host.encode([torch.full((1, 20), value)]))
+
+        first = threading.Thread(target=submit, args=(1.0,))
+        first.start()
+        # The first worker is busy by now, so the second clip goes to the other one.
+        threading.Event().wait(0.1)
+        second = threading.Thread(target=submit, args=(2.0,))
+        second.start()
+        inside.wait()
+        release.set()
+        for thread in (first, second):
+            thread.join(5)
+        assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_client_requests_in_flight_use_their_own_connections(host_dir):
+    entered, release = threading.Semaphore(0), threading.Event()
+
+    def worker():
+        def encode(wavs):
+            entered.release()
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, None
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        out = []
+        threads = [
+            threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0)
+        ]
+        for thread in threads:
+            thread.start()
+            # A worker holds this request, still unanswered, before the next is
+            # sent. Sent together, one free worker may take both as one batch.
+            assert entered.acquire(timeout=5)
+        release.set()  # both requests reached the host at once
+        for thread in threads:
+            thread.join(5)
+        assert sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+        assert len(client._idle) == 2  # both connections are kept for reuse
+    finally:
+        host.close()
+
+
+def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch):
+    entered, release = threading.Event(), threading.Event()
+    recorder = _Recorder()
+    warmed = []
+
+    def last_warmup():
+        entered.set()
+        assert release.wait(5)
+        warmed.append("last")
+
+    host = shared.SharedReferenceEncoderHost(
+        host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)]
+    )
+    client = shared.SharedReferenceEncoderClient(host_dir)
+    try:
+        assert entered.wait(5)
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03)
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            client.wait_until_ready()
+        assert warmed == ["first"] and recorder.batches == [] and not client._idle
+        release.set()
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5)
+        host.wait_until_ready()
+        client.wait_until_ready()
+        assert warmed == ["first", "last"] and recorder.batches == []
+        assert len(client._idle) == 1
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        host.close()
+
+
+def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch):
+    release = threading.Event()
+    monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05)
+    monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))])
+    timer = threading.Timer(0.2, release.set)
+    timer.start()
+    try:
+        # The first socket waits for readiness before starting its much shorter
+        # encode deadline. No API warmup call is required to protect this path.
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        timer.join(5)
+        host.close()
+
+
+def test_failed_host_warmup_is_not_ready(host_dir):
+    def fail():
+        raise RuntimeError("warmup failed")
+
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)])
+    try:
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            shared.SharedReferenceEncoderClient(host_dir).wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.encode([torch.ones(1, 20)])
+        assert recorder.batches == []
+    finally:
+        host.close()
+
+
+def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch):
+    from unittest.mock import Mock
+
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    enc = re_mod.MossReferenceEncoder(
+        _FakeProcessor(),
+        variant="local",
+        n_vq=_N_VQ,
+        sr_target=_SR,
+        speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20),
+    )
+    enc._shared_client = Mock()
+    enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling")
+    local = Mock()
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local)
+    with pytest.raises(shared.SharedReferenceEncoderStartupError):
+        enc._encode_shared_or_local([torch.ones(1, 20)])
+    local.assert_not_called()
@@ -0,0 +1,301 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerMultiheadAttention,
+    MossAudioTokenizerRotaryEmbedding,
+    RingKVCache,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)]
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("execution", ["graph", "compiled"])
+def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution):
+    torch.manual_seed(45)
+    batch, heads, states, dim = 3, 2, 5, 64
+    projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype)
+    expected_cache = cache.clone()
+    end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0]
+    expected_end = end.clone()
+    slots = torch.tensor([0, 4, 2], device="cuda")
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+
+    def candidate():
+        return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True)
+
+    graph = None
+    if execution == "compiled":
+        run = torch.compile(candidate, fullgraph=True)
+    else:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            candidate()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            actual = candidate()
+
+    for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]):
+        end.fill_(start)
+        expected_end.fill_(start)
+        cache.copy_(expected_cache)
+        projected.normal_()
+        slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17)
+        if graph is not None:
+            graph.replay()
+        else:
+            actual = run()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)])
+def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity):
+    q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)]
+    cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype)
+    end = torch.tensor([0, capacity + 3, 9, 0], device="cuda")
+    slots = torch.tensor([1, 3], device="cuda")
+    lengths = torch.zeros(2, device="cuda", dtype=torch.int32)
+    before_cache, before_end = cache.clone(), end.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    for active in [False, True, False, True]:
+        lengths.fill_(frames if active else 0)
+        expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity)
+        graph.replay()
+        torch.testing.assert_close(cache, before_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, before_end, rtol=0, atol=0)
+        torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025)
+        if not active:
+            assert out.count_nonzero() == 0
+
+
+def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity():
+    torch.manual_seed(19)
+    capacity, frames = 400, 480
+    state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32)
+    slots = torch.tensor([1, 0], device="cuda")
+    context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool))
+    for start in [0, frames]:
+        k = torch.randn(2, 2, frames, 64, device="cuda")
+        v = torch.randn_like(k)
+        state.complete(k, v, execution_context=context)
+        retained = torch.arange(frames - capacity, frames, device="cuda")
+        destinations = (start + retained) % capacity
+        torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0)
+
+
+def _ring_reference(q, k, v, cache, end, slots, lengths, context):
+    """Independent chronological ring oracle with FP32 attention arithmetic."""
+    capacity = cache.shape[-2]
+    output = torch.zeros_like(q)
+    for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())):
+        if length == 0:
+            continue
+        start = int(end[slot])
+        for token in range(max(0, length - capacity), length):
+            cache[0, slot, :, (start + token) % capacity] = k[row, :, token]
+            cache[1, slot, :, (start + token) % capacity] = v[row, :, token]
+        stop = start + length
+        physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long)
+        available = torch.arange(max(0, stop - capacity), stop, device=q.device)
+        physical_positions[available % capacity] = available
+        queries = torch.arange(start, stop, device=q.device)
+        distance = queries[:, None] - physical_positions[None]
+        mask = (physical_positions[None] >= 0) & (distance >= 0)
+        if context > 0:
+            mask &= distance < context
+        output[row, :, :length] = F.scaled_dot_product_attention(
+            q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None]
+        ).to(q.dtype)
+        end[slot] = stop
+    return output
+
+
+@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)])
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution):
+    torch.manual_seed(29)
+    batch, heads, state_capacity = 3, 2, 5
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    # Strided metadata is allowed by the kernel contract.
+    end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0]
+    slots = torch.arange(batch * 2, device="cuda")[::2]
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+    context = capacity
+    attention = slot_ring_attention
+    if execution == "compiled":
+        # Functionalization must preserve mutations of the persistent ring
+        # and its offsets between calls, including slot reuse.
+        attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+    end.zero_()
+    expected_cache, expected_end = cache.clone(), end.clone()
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        projected.normal_()
+        if step == 4:
+            end[1:3].zero_()
+            expected_end[1:3].zero_()
+        expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, end, slots, lengths, context)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025)
+        padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None]
+        assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)])
+def test_slot_attention_integrates_with_current_mha(frames, context):
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._slot_attention = slot_ring_attention
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
+
+
+@pytest.mark.parametrize(
+    "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)]
+)
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution):
+    torch.manual_seed(31)
+    batch, heads, state_capacity = 3, 2, 6
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    reference_cache = cache.clone()
+    end = torch.zeros(state_capacity, device="cuda", dtype=torch.long)
+    slots = torch.zeros(batch, device="cuda", dtype=torch.long)
+    lengths = torch.zeros(batch, device="cuda", dtype=torch.int32)
+    starts = torch.zeros(batch, device="cuda", dtype=torch.long)
+    attention = slot_ring_attention_rows
+    if execution == "compiled":
+        attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        cache.copy_(reference_cache)
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        starts.copy_(end[slots])
+        projected.normal_()
+        expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)])
+def test_fused_rows_attention_integrates_with_mha_offsets(frames, context):
+    torch.manual_seed(11)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    reference._slot_attention = slot_ring_attention
+    candidate._slot_attention_rows = slot_ring_attention_rows
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""A second streaming-decoder graph set can replay concurrently with the first."""
+
+import threading
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+CAP, DIM = 8, 64
+
+
+class _Codec(nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM)
+        self.state = torch.zeros(CAP + 16, DIM, device=device)
+
+    def decode_streaming_tensors(self, codes, lengths, slots, valid_rows):
+        x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None]
+        y = x @ self.weight
+        self.state.index_copy_(0, slots, y[:, -1])
+        return y.transpose(1, 2).contiguous(), lengths
+
+    def reset_decoder_state_slots(self, slots):
+        self.state.index_fill_(0, slots.to(self.state.device), 0)
+
+
+def _wrapper(codec, **kwargs):
+    return CUDAGraphStreamingDecoderWrapper(
+        codec,
+        state_capacity=CAP,
+        num_quantizers=2,
+        vllm_config=None,
+        compiled_decode=codec.decode_streaming_tensors,
+        **kwargs,
+    )
+
+
+def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    codec = _Codec(device)
+    main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3])
+    fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True)
+    main.warmup(device)
+    fast.warmup(device)
+    assert main.is_ready and fast.is_ready
+    assert main._capture_stream is None and fast._capture_stream is not None
+
+    codes_main = torch.randint(0, 5, (2, 4, 3), device=device)
+    codes_fast = torch.randint(0, 5, (2, 2, 1), device=device)
+    errors = []
+
+    def run(wrapper, codes, slots, iters):
+        stream = torch.cuda.Stream(device=device)
+        slot_ids = torch.tensor(slots, device=device)
+        try:
+            with torch.cuda.stream(stream):
+                for _ in range(iters):
+                    codec.reset_decoder_state_slots(slot_ids)
+                    audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone()
+                    expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2)
+                    stream.synchronize()
+                    torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3)
+        except Exception as error:  # surfaced below
+            errors.append(error)
+
+    threads = [
+        threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)),
+        threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=60)
+    assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish"
+    assert not errors, errors
@@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
    assert [resumed[i].item() for i in (0, 1)] == [2, 36]


-@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
-def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls):
+@pytest.mark.parametrize(
+    "graph_capacity,expected_groups",
+    [
+        # No graph, or graphs that do not cover the state capacity: exact tails.
+        (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        # Full graph coverage: finished tails (3, 7) join the padded 8-frame
+        # bucket with the live 8-frame row; the finished 1-frame row keeps the
+        # exact 1-frame graph and the live 15-frame row stays alone.
+        (8, {8: [0, 1, 2], 1: [3], 15: [4]}),
+    ],
+)
+def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups):
    s = session()
    s._cudagraph_wrapper = (
        None
@@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa
            for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
        ]
    )
-    assert step.call_count == expected_calls
+    # Slots are leased in request order, so slot == row here.
+    groups = {}
+    for call in step.call_args_list:
+        plan = call.args[0]
+        groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan)
+        terminal = call.kwargs["terminal_slots"] or set()
+        assert terminal == {slot for slot in plan if finished[slot]}
+    assert groups == expected_groups
+    assert step.call_count == len(expected_groups)
+    # Output is always cropped to the exact frame count, padded or not.
    assert [result[i].shape[-1] for i in range(5)] == lengths
-    if graph_capacity == 8:
-        assert list(step.call_args_list[0].args[0]) == [0, 1, 2]
-        assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1}
-        assert list(step.call_args_list[1].args[0]) == [3]  # Preserve the exact one-frame graph.
+    # Finished rows released their slots; live rows keep theirs.
+    assert set(decoder._stream_req_slots) == {"2", "4"}
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rows", [16, 19])
+@torch.inference_mode()
+def test_first_frame_graph_outputs_own_storage(dtype, rows):
+    class Decoder(torch.nn.Linear):
+        config = SimpleNamespace(num_quantizers=1)
+
+        def forward(self, codes):
+            return codes.to(dtype) * self.weight
+
+    model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
+    model.weight.fill_(1)
+    decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
+    decoder.capture()
+    codes = torch.arange(rows, device="cuda").reshape(-1, 1)
+    first = decoder.decode(codes)
+    second = decoder.decode(codes + 32)
+    assert first.dtype == second.dtype == torch.float32
+    torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
+    torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _codec(enabled):
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=8,
+        decoder_dim=32,
+        upsample_rates=[2, 2],
+        upsampling_ratios=[2],
+        codebook_size=16,
+        num_quantizers=2,
+    )
+    config._attn_implementation = "sdpa"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
+        )
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
+    model.precompute_snake_caches()
+    return model
+
+
+def _flags():
+    c = torch.backends.cudnn
+    return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
+
+
+@torch.inference_mode()
+def test_autotuned_codec_graph_replay_restores_flags():
+    torch.manual_seed(21)
+    model = _codec(True)
+    codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
+    original = _flags()
+    for _ in range(3):
+        model(codes)
+    assert _flags() == original
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = model(codes)
+    assert _flags() == original
+    codes.random_(0, 16)
+    graph.replay()
+    expected = model(codes)
+    torch.testing.assert_close(captured, expected)
+    assert _flags() == original
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+@torch.inference_mode()
+def test_codec_failure_restores_cudnn_settings(mocker, enabled):
+    model = _codec(enabled)
+    codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
+    original = _flags()
+
+    def fail(hidden):
+        if enabled:
+            assert torch.backends.cudnn.benchmark
+            assert torch.backends.cudnn.benchmark_limit == 10
+        else:
+            assert _flags() == original
+        raise RuntimeError("decode failed")
+
+    mocker.patch.object(model, "_decode_waveform", side_effect=fail)
+    with pytest.raises(RuntimeError, match="decode failed"):
+        model(codes)
+    assert _flags() == original
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeSnakeBeta,
+)
+
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _blocks():
+    config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
+    with torch.no_grad():
+        for module in blocks.modules():
+            if isinstance(module, Qwen3OmniMoeSnakeBeta):
+                module.alpha.uniform_(-0.5, 0.5)
+                module.beta.uniform_(-0.5, 0.5)
+    return blocks
+
+
+def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
+    blocks = _blocks()
+    hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    keys = set(blocks.state_dict())
+
+    assert use_fused_snake(blocks) == hf_count > 0
+    assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
+    assert set(blocks.state_dict()) == keys  # checkpoint weights load unchanged
+
+
+def test_fused_snake_blocks_match_hf_blocks():
+    reference = _blocks()
+    fused = copy.deepcopy(reference)
+    use_fused_snake(fused)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 32, 20)
+    with torch.no_grad():
+        expected = x
+        actual = x
+        for ref_block, fused_block in zip(reference, fused):
+            expected = ref_block(expected)
+            actual = fused_block(actual)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("option", "device", "cuda", "fused"),
+    [
+        (False, "cuda", True, False),
+        (True, "cuda", True, True),
+        (True, "cpu", True, False),
+        (True, "cuda", False, False),
+    ],
+)
+def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+    )
+    model = Qwen3OmniMoeCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
+            ),
+            device_config=SimpleNamespace(device=device),
+        )
+    )
+    assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+    plan_decode_groups,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 4
+_CODEBOOK = 16
+
+
+def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
+    torch.manual_seed(0)
+    config = Qwen3OmniMoeCode2WavConfig(
+        codebook_size=_CODEBOOK,
+        num_quantizers=quantizers,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+        sliding_window=8,
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
+    return model.eval()
+
+
+def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (3, _Q))
+    with torch.inference_mode():
+        for row in range(codes.shape[0]):
+            pcm = decoder.decode(codes[row : row + 1])
+            chunk0 = code2wav.chunked_decode_streaming(
+                codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
+            )[0]
+            assert pcm.dtype == torch.float32
+            assert torch.equal(pcm[0], chunk0.reshape(-1).float())
+            assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
+        batched = decoder.decode(codes)
+    assert batched.shape == (3, pcm.shape[-1])
+
+
+def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (_Q, 4))
+    with torch.inference_mode():
+        first = decoder.decode(codes[:, :1].T)[0]
+        chunk0 = code2wav.chunked_decode_streaming(
+            codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
+        )[0]
+    torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
+
+
+def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
+    import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
+
+    real = wrapper_module.CUDAGraphDecoderWrapper
+    wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
+    monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
+    extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
+    code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
+    code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
+        code2wav, device=torch.device("cuda"), **kw
+    )
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.code2wav = code2wav
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
+    )
+    for batches in ([], [2, 4, 8]):
+        extra["codec_graph_batch_sizes"] = batches
+        Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+        args = wrapper.call_args.kwargs
+        sizes = args["capture_sizes"]
+        assert {1, 3, 7, 15, 31, 50} <= set(sizes)
+        assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
+        assert set(args["extra_capture_shapes"]) == {
+            (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
+        }
+    assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
+    extra.pop("codec_chunk_ramp")
+    extra.pop("codec_graph_batch_sizes")
+    Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+    assert wrapper.call_args.kwargs["capture_sizes"] is None
+
+
+_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
+
+
+def _bucket(length):
+    return next((size for size in _SIZES if size >= length), None)
+
+
+def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
+    assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
+    calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
+    assert [len(rows) for rows, _length in calls] == [16, 4]
+    calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
+    assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
+    assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
+
+
+@pytest.mark.parametrize("frames", [65, 129])
+def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.enabled = wrapper._warmed_up = True
+    wrapper._bucket_sizes = [8, 32, 64]
+    wrapper._compiled_shapes = {(1, 64), (8, 64)}
+    wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
+    wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
+    assert groups == [([0], frames), ([1], frames)]
+    codes = torch.arange(2 * frames).reshape(2, 1, frames)
+    for rows, width in groups:
+        actual = wrapper.decode(codes[rows, :, :width])
+        torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
+    assert wrapper.decoder.call_count == 2
+    for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
+        graph.replay.assert_not_called()
+
+
+class _PaddingGraphs:
+    def __init__(self, model, batch_sizes):
+        self.model = model
+        self.batch_sizes = sorted(batch_sizes)
+        self.calls = []
+
+    def _get_padded_size(self, length):
+        return _bucket(length)
+
+    def decode(self, codes):
+        rows, _q, frames = codes.shape
+        size = _bucket(frames)
+        batch = next(b for b in self.batch_sizes if b >= rows)
+        self.calls.append((rows, frames))
+        padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
+        padded[:rows, :, :frames] = codes
+        out = self.model(padded)
+        return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
+
+
+def test_grouped_streaming_decode_matches_one_padded_batch():
+    code2wav = _code2wav()
+    lengths = [1, 3, 50, 7, 31]
+    lefts = [0, 1, 25, 3, 15]
+    codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
+    for row, length in enumerate(lengths):
+        codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
+    counts = [length * _Q for length in lengths]
+    with torch.inference_mode():
+        reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+        graphs = _PaddingGraphs(code2wav, [1, 2, 4])
+        code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
+        code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
+        grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+
+    for want, got in zip(reference, grouped, strict=True):
+        assert want.shape == got.shape
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
+def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+    codec = _code2wav(quantizers)
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
+    model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
+        codes, left_context_size=left, seq_token_counts=counts
+    )
+    codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
+    with torch.inference_mode():
+        reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
+        first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
+        out = model.forward(
+            input_ids=codes.flatten(),
+            positions=None,
+            seq_token_counts=[quantizers * frames] * 2,
+            runtime_additional_information=[
+                {"meta": {"left_context_size": 0, "first_audio": True}},
+                {"meta": {"left_context_size": 0, "first_audio": False}},
+            ],
+        ).multimodal_outputs
+    torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
+    torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
+    assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]
Author	SHA1	Message	Date
tzhouam	bedc42a94d	[Bugfix] Preserve session IDs when upgrading Omni engine requests Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-10-04 03:19:39 +08:00
andyluo7	ee8fdab1dc	[CI/Build] Fix MOSS CPU tests on ROCm workers (#8447 ) Signed-off-by: andyluo7 <andy.luo@amd.com>	2026-10-03 09:13:51 +08:00
Sy03andYueqian Lin	3bc3f1a7d0	[Model] Optimize MOSS Local 1.5 MRV2 serving and progressive chunks (#8213 ) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>	2026-10-02 18:03:14 +00:00
Zheng Wengangandamy-why-3459	527982d886	[Bugfix] Skip async-chunk SHM put for stage-0-final requests (#7245 ) Signed-off-by: ZhengWG <zwg0606@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-10-02 09:14:39 -07:00
Sy03	2ba5074ca0	[Model] Optimize CosyVoice3 packed streaming with bounded HiFT graphs (#8420 ) Signed-off-by: Sy03 <1370724210@qq.com>	2026-10-02 08:49:50 -07:00
Yueqian Lin	4025bb30f6	[Perf][AuK] Graph the encoder prefill, cache reference latents and fuse the codec activation (#8305 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-10-02 15:43:55 +00:00
Yueqian Lin	13c66da752	[Perf][AuK] Precompute adaLN modulations per request and run the conv position embedding as one GEMM (#8328 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-10-02 08:12:29 -07:00
Nick CaoandCodex	0e6791e9c4	[Bugfix] Initialize eager state in Higgs test fixture (#8426 ) Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Codex <noreply@openai.com>	2026-10-02 10:22:52 -04:00
Zheng Wengangandamy-why-3459	eb1699bfca	[FEAT][Fronted] Support incremental prefill for realtime video streaming (#5894 ) Signed-off-by: ZhengWG <zwg0606@gmail.com> Signed-off-by: Zheng Wengang <zwg0606@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-10-02 22:18:51 +08:00
Sy03	e372aa3a41	[Bugfix] Restore Fish Speech and CosyVoice3 audio quality at high concurrency (#8422 ) Signed-off-by: Sy03 <1370724210@qq.com>	2026-10-02 07:11:04 -07:00
Sy03 princepride Claude Opus 5.5	ae58c2d8b9	[Model] YuE2: async single-stage synthesis and CUDA graph acceleration (#8407 ) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: Wang Zhipeng <wangzhipeng628@gmail.com> Co-authored-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>	2026-10-02 05:17:06 -07:00
linzhenpl07 Claude Opus 5 Zhou Taichang	b10270a01c	[AR-Diffusion] Page in kernel-legal blocks so the default resolution can run (#6481 ) Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>	2026-10-02 20:05:04 +08:00
Sy03	e84014dffb	[Model] Add single-stage Qwen3-TTS with in-Talker streaming codec decode (#8259 ) Signed-off-by: Sy03 <1370724210@qq.com>	2026-10-02 21:34:30 +10:00
Rahul Steiger	596487fe62	[Core] Prototype diffusion attention execution contracts with FA4 (#7379 )	2026-10-02 17:58:47 +08:00
MaciejBalaNVand汪志鹏	97d6c693a6	[feature] Add HSDP pre-sharded option (#7948 ) Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>	2026-10-02 09:35:09 +00:00
汪志鹏	b742f86136	[CI/Build] Restore streaming executor fixture contract (#8418 ) Signed-off-by: princepride <wangzhipeng628@gmail.com>	2026-10-02 01:45:56 -07:00
andyluo7	53c7340668	test: measure offload memory inside diffusion workers (#8189 ) Signed-off-by: andyluo7 <andy.luo@amd.com>	2026-10-02 14:07:20 +08:00
Sy03	d93e1a8e09	[Model] Optimize Qwen3-Omni MRv2 handoff, code prediction and audio decoding (#8223 ) Signed-off-by: Sy03 <1370724210@qq.com>	2026-10-01 22:22:54 -07:00
wtz2333andYuanheng Zhao	68003cf6ad	[Bugfix] Fix washed-out LingBot Video outputs (#7037 ) Signed-off-by: wtz2333 <2955110911@qq.com> Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>	2026-10-02 03:57:50 +00:00