fix: preserve direct execution of the review verifier

fix: align migrated review helpers with repository lint policy
docs: distinguish generic review selection from session eligibility
2026-10-01 17:21:00 -07:00 · 2026-10-01 17:18:54 -07:00 · 2026-10-01 17:11:12 -07:00 · 2026-10-01 17:08:01 -07:00
360 changed files with 2990 additions and 27893 deletions
@@ -16,9 +16,6 @@ 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
@@ -88,16 +85,11 @@ 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
@@ -102,7 +102,6 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe
 SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
 CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
 CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
-CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"

 # Bootstrap Buildkite ``if`` expressions.
 # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
@@ -621,11 +620,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
    if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
        commands = merged.get("commands")
        if isinstance(commands, list):
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
        elif commands is None:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
        else:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands]
    # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
    if "retry" in step:
        merged["retry"] = step["retry"]
@@ -1,38 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
-
-from __future__ import annotations
-
-import importlib
-import re
-from pathlib import Path
-
-
-def main() -> None:
-    dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
-    match = re.search(
-        r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$",
-        dockerfile.read_text(),
-        re.MULTILINE,
-    )
-    if match is None:
-        raise RuntimeError(f"Missing vLLM release version in {dockerfile}")
-    expected_version = match.group(1)
-
-    import torch
-    import vllm
-
-    assert vllm.__version__ == expected_version, vllm.__version__
-    assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
-    for module in ("flashinfer", "vllm._C_stable_libtorch"):
-        importlib.import_module(module)
-    print(
-        f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
-        f"expected={expected_version} torch={torch.__version__} "
-        f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -124,30 +124,6 @@ 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
@@ -48,8 +48,6 @@ steps:
      - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
        source_file_dependencies: omni_qwen3_omni_doc
        timeout_in_minutes: 90
-        artifact_paths:
-          - audio_*.wav
        commands:
          - pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model"

@@ -158,11 +156,6 @@ steps:
      - label: ":full_moon: TTS · Function Test with L4"
        source_file_dependencies: tts_qwen3_tts_function
        timeout_in_minutes: 120
-        env:
-          VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
-        artifact_paths:
-          - failed_speech_*.wav
-          - failed_speech_*.json
        commands:
          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
          - export VLLM_USE_DEEP_GEMM="0"
@@ -25,16 +25,6 @@ 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,6 +9,15 @@ 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:
@@ -103,6 +112,12 @@ 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 |
@@ -110,6 +125,15 @@ If docs and live code disagree, verify the code/tests and report the drift.
 | [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,11 +20,15 @@ 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;
- run-level and domain markers place hardware/model tests in the intended CI
-  lane without silently skipping the contract.
+- 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.

-Map source symbols to tests with bounded `rg` searches; do not assume the test
-tree mirrors production paths.
+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.

 ## Runtime check

@@ -0,0 +1,50 @@
+# 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.
@@ -0,0 +1,80 @@
+# 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.
@@ -0,0 +1,87 @@
+#!/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
@@ -0,0 +1,73 @@
+#!/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}]"
@@ -0,0 +1,129 @@
+#!/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())
@@ -0,0 +1,8 @@
+#!/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" "$@"
@@ -0,0 +1,227 @@
+# 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()
@@ -1,7 +1,10 @@
-# Use the official release image with its matching CUDA/PyTorch stack.
-ARG VLLM_VERSION=0.31.0
+# Target vLLM version: v0.30.0 (tag @ ced6857afa0e).
+# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
+# runtime. Do not reinstall a commit wheel or repair dependencies on top of
+# an older base image at a release boundary. For a future unreleased target,
+# restore an explicit wheel pin and validate its complete runtime stack.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v${VLLM_VERSION}
+ARG VLLM_BASE_TAG=v0.30.0

 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents
@@ -9,7 +12,6 @@ ARG VLLM_BASE_TAG=v${VLLM_VERSION}
 # everything except the actual project dir, since that should not
 # change much across branches.
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps
-ARG VLLM_VERSION
 # Keep APP_DIR outside /workspace so k8s Buildkite workspace mounts cannot
 # shadow the editable install target (#5364). Docker WORKDIR creates /opt if
 # needed.
@@ -34,17 +36,12 @@ RUN apt-get update && \
 # for CI
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-# Scope the override to Omni; dependency source builds need their own versions.
-ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
 # Editable install: deps + project metadata are cached in this layer, while
 # the package itself is linked to APP_DIR (populated later by COPY --link).
 # Do NOT use a non-editable install here — that copies an empty stub into
 # site-packages and breaks ``import vllm_omni.entrypoints``.
 RUN uv pip install --system -e ".[dev]"
-# Fail the image build if dependency resolution replaced the selected vLLM.
-# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start
-# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py.
-RUN python3 -c "import vllm; assert vllm.__version__ == '${VLLM_VERSION}', vllm.__version__"

 # The final stage just copies the working dir on top of the
 # dependencies image. This uses --link so that we can skip
@@ -1,6 +1,6 @@
 # NOTE: this is largely taken from vLLM

-# Exclude .git for CI builds — version is set via VLLM_OMNI_VERSION_OVERRIDE
+# Exclude .git for CI builds — version is set via SETUPTOOLS_SCM_PRETEND_VERSION
 .git
 .venv
 /build
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.31.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.30.0
 FROM ${BASE_IMAGE}

 ARG COMMON_WORKDIR=/app
@@ -1,7 +1,7 @@
 # Keep the ROCm runtime aligned with the vLLM release targeted by
 # docker/Dockerfile.ci. A stale base image can build successfully while failing
 # at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0

 # Keep the dependency-heavy portion independent from the source tree. The
 # registry cache is keyed on these inputs, while the final source copy can be
@@ -12,11 +12,11 @@ FROM ${BASE_IMAGE} AS deps
 # BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
 # VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when
 # USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's
-# release version. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
+# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
 # intentional custom source pin must update that regression in the same change,
 # and should return to the stable release as soon as one is available.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.31.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}

@@ -75,8 +75,7 @@ WORKDIR ${COMMON_WORKDIR}/vllm-omni
 RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-# Scope the override to Omni; dependency source builds need their own versions.
-ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0

 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
        --no-cache-dir -e ".[dev]" --no-build-isolation && \
@@ -2,7 +2,7 @@

 This section lists the most common options for running vLLM-Omni.

-For options within a vLLM Engine, please refer to the [vLLM 0.31 configuration guide](https://docs.vllm.ai/en/v0.31.0/configuration/index.html).
+For options within a vLLM Engine, please refer to the [vLLM 0.30 configuration guide](https://docs.vllm.ai/en/v0.30.0/configuration/index.html).

 Each model defines fixed topology in a registered `PipelineConfig` and runtime
 overrides in a deploy YAML.
@@ -186,7 +186,7 @@ their keys only.
 ## Inherited vLLM variables

 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.31 environment-variable reference](https://docs.vllm.ai/en/v0.31.0/configuration/env_vars.html)
+the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.0/configuration/env_vars.html)
 for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
 XPU, ModelScope, and FlashInfer workspace settings.

@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.

 ## Model-specific variables

-The inventory includes 93 variables read by a single model or pipeline family. They are
+The inventory includes 66 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 | 64 | Move a stable setting into typed stage or model configuration. |
+| Promote | 38 | 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 | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Internalize | 16 | 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
@@ -1,22 +0,0 @@
-# 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,9 +319,7 @@ 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.
@@ -352,73 +350,12 @@ 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
-two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile.
-This changes the default CUDA MRV2 streaming path:
+single-GPU 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.
@@ -514,50 +451,3 @@ 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,9 +4,6 @@ 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
@@ -1,133 +0,0 @@
-# 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,7 +68,6 @@ 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
@@ -117,5 +116,4 @@ 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)
@@ -1,22 +0,0 @@
-# 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.
@@ -1,7 +1,7 @@
 # Installation

 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. This source checkout targets vLLM 0.31.0. The published vLLM-Omni 0.30.0 wheel instructions use vLLM 0.30.0; choose the installation path that matches your Omni revision.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0.

 vLLM-Omni supports the following hardware platforms:

@@ -1,13 +1,11 @@
 # --8<-- [start:requirements]

- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, etc.)
+- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)

 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]

-vLLM-Omni depends on the matching major/minor release of vLLM. This source
-checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
-the published 0.30.0 release with its matching vLLM dependency.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.

 !!! note
    PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -43,17 +41,13 @@ uv pip install 'vllm-omni[demo]==0.30.0'
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, install the final 0.31.0
-release wheel from its commit-pinned index:
+If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library

 ```bash
-uv pip install vllm==0.31.0 --torch-backend=auto \
-  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
+uv pip install vllm==0.30.0 --torch-backend=auto
 ```

-The default 0.31.0 wheel contains CUDA 12.9 binaries. If you need a different
-CUDA variant or want to reuse an existing PyTorch installation, build vLLM
-from source instead.
+The 0.30.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -74,16 +68,16 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.31.0
+git checkout v0.30.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-# Default CUDA 12.9 wheel for v0.31.0 on x86_64
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.30.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.30.0/vllm-0.30.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
-VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto
+uv pip install --editable .
 ```
 Install vllm with command below (If you already have PyTorch).
 ```bash
@@ -134,13 +128,12 @@ docker run --runtime nvidia --gpus 2 \
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda .
 ```

-The source build defaults to `vllm/vllm-openai:v0.31.0`. To select that
-upstream release image explicitly:
+If you want to specify the base vLLM version:

 ```bash
 DOCKER_BUILDKIT=1 docker build \
  -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \
  -t vllm-omni-cuda .
 ```

@@ -10,9 +10,8 @@

 For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images.

-vLLM-Omni depends on the matching major/minor release of vLLM. This source
-checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
-the published 0.30.0 release with its matching vLLM dependency.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
+0.30.x release line uses vLLM 0.30.x.

 The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The
 `vllm-omni` package does not install vLLM as a dependency, so non-Docker source
@@ -50,14 +49,13 @@ uv pip install onnxruntime-rocm

 #### Installation of vLLM

-If you do not need to modify source code of vLLM, install the final 0.31.0
-release wheel for ROCm 7.2.3:
+If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library

 ```bash
-uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
+uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
 ```

-The pre-built 0.31.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.30.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni

@@ -77,7 +75,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.31.0
+git checkout v0.30.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -89,7 +87,7 @@ python3 setup.py develop
 #### Build docker image

 The source build defaults to the published upstream base image
-`vllm/vllm-openai-rocm:v0.31.0`, aligned with this source checkout.
+`vllm/vllm-openai-rocm:v0.30.0`, aligned with the vLLM release used by CI.
 Older or custom bases must provide the vLLM APIs checked by the Dockerfile's
 image-build canary. This upstream base is distinct from the prebuilt
 `vllm/vllm-omni-rocm` images discussed below; their published-tag availability
@@ -104,7 +102,7 @@ To select the upstream ROCm base explicitly, pass `BASE_IMAGE`:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
  -f docker/Dockerfile.rocm \
-  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 \
  -t vllm-omni-rocm .
 ```

@@ -20,16 +20,12 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
 #### Build docker image

 ```bash
-export VLLM_VERSION=v0.31.0
-DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
-  --build-arg VLLM_VERSION="$VLLM_VERSION" .
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
 ```

 This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
-base image, which Docker pulls automatically. The Dockerfile still defaults to
-v0.30.0; the command above explicitly selects v0.31.0 to match this source
-checkout. To target a different vLLM release, pass
-`--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
+base image (v0.30.0 by default), which Docker pulls automatically. To target a different vLLM release,
+pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
 the base from upstream's own Dockerfile first and point the build at it:

 ```bash
@@ -71,14 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886

 # --8<-- [start:installation-main]

-You can also install vLLM-Omni from source if you want to use the latest features
-or bug fixes. This source checkout targets vLLM 0.31.0 and needs a matching
-vLLM-Ascend runtime. The v0.30.0 images above are for the published Omni v0.30.0
-release. Check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886)
-for aligned versions and NPU validation before selecting a development image.
+You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.

 ```bash
-# Inside a container with matching vLLM and vLLM-Ascend, install Omni from source
+# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
 cd /vllm-workspace
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -19,20 +19,17 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.31.0 --torch-backend=auto \
-  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
+uv pip install vllm==0.30.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
+uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
 ```

-These commands install vLLM 0.31.0 for this source checkout. The CUDA index is
-pinned to the final release commit. For published vLLM-Omni wheels and additional
-installation methods, see the [installation guide](installation/README.md).
+For additional installation methods — please see the [installation guide](installation/README.md).

 !!! note
    It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
@@ -16,10 +16,6 @@ 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
@@ -58,7 +54,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. 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.
+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.

 **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`.

@@ -68,10 +64,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. 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. |
+| `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. |
 | `system_prompt` | string or null | null | Optional custom system prompt. |
-| `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. |
+| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
 | `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. |
@@ -111,42 +107,9 @@ 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

- 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).
+- 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.
 - 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.
@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes

 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
+- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.

 ---
@@ -155,22 +155,16 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```

-Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
+Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.

-The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
-
-The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
-
-Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
-
-Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
+Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:

 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```

-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.

 ### CLI client

@@ -200,7 +194,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes

 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).

 ---
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
-import copy
 import functools
 import json
 import os
@@ -11,11 +10,13 @@ from typing import Any

 import numpy as np
 import soundfile as sf
+from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio

 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config

 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -58,7 +59,8 @@ def run_e2e():
        "--deploy-config",
        type=str,
        default=None,
-        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
+        help="Override the deploy config path. If unset, auto-loads "
+        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
    )
    parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
    parser.add_argument(
@@ -103,6 +105,8 @@ def run_e2e():
        profiler_config=args.profiler_config,
    )

+    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
+
    print("Model initialized. Preparing inputs...")
    ref_audio_path = args.ref_audio or _default_ref_audio()
    if not os.path.exists(ref_audio_path):
@@ -134,7 +138,7 @@ def run_e2e():

    print(f"Generating for prompt: {args.text}")

-    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
+    config = CosyVoice3Config()
    tokenizer = get_qwen_tokenizer(
        token_path=args.tokenizer,
        skip_special_tokens=config.skip_special_tokens,
@@ -145,10 +149,29 @@ def run_e2e():
    min_len = int(base_len * config.min_token_text_ratio)
    max_len = int(base_len * config.max_token_text_ratio)

-    # Keep the deploy profile's sampling policy, penalties and required stops.
-    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
-    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
-    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
+    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
+    gpt_sampling = SamplingParams(
+        temperature=1.0,
+        top_p=sampling_cfg["top_p"],
+        top_k=sampling_cfg["top_k"],
+        repetition_penalty=2.0,
+        min_tokens=min_len,
+        max_tokens=max_len,
+        stop_token_ids=[sampling_cfg["eos_token_id"]],
+        # allowed_token_ids=list(range(6561+3)),
+        detokenize=False,
+    )
+    # Not used
+    s2mel_sampling = SamplingParams(
+        temperature=1.0,
+        top_p=1.0,
+        top_k=-1,
+        repetition_penalty=2.0,
+        max_tokens=256,
+        detokenize=False,
+    )
+
+    sampling_params_list = [gpt_sampling, s2mel_sampling]

    profiler_enabled = args.profiler_config is not None
    if profiler_enabled:
@@ -28,9 +28,7 @@ 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,
@@ -43,12 +41,10 @@ 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
@@ -75,13 +71,7 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
        # conditioning prefix from the scheduled tokens alone.
        KEY_PREFIX_IDS: list(prompt_ids),
    }
-    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.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
    params.stop_token_ids = list(STOP_TOKEN_IDS)
    params.detokenize = False
    return params
@@ -154,9 +144,6 @@ 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,
@@ -1,15 +1,11 @@
 #!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
-# other devices use vllm_omni/deploy/cosyvoice3.yaml.
+# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.

 set -e
@@ -35,10 +35,6 @@ 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/)
@@ -156,296 +152,3 @@ 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,11 +40,9 @@ 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, 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.
+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).

 ## Running

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

 ## Notes

-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.
+- **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).
 - 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
@@ -2,7 +2,7 @@
 # kernels==0.16.1 requires the LayerRepository version=/revision= integration
 # introduced in transformers 5.13. Keep the tested upper bound: 5.15 has
 # a reported model-construction regression; later versions need validation.
-# This range also satisfies vLLM 0.31.0's transformers>=5.10.4,<5.18.0 requirement.
+# This range also satisfies vLLM 0.30.0's transformers>=5.10.4 requirement.
 # MOSS-TTS-Nano's list/set weight-loading compatibility is handled by
 # transformers_keys_to_ignore_compat in model_executor/models/utils.py.
 transformers >= 5.13.0, < 5.15
@@ -10,8 +10,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 REPO_ROOT = Path(__file__).resolve().parents[2]
 AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2"
 AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh"
-CUDA_RELEASE_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.cuda"
-CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
+CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
 ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm"
 ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore"

@@ -31,18 +30,14 @@ def _line_index(lines: list[str], prefix: str) -> int:
    return matches[0]


-def test_rocm_base_tracks_cuda_vllm_release() -> None:
-    cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE")
-    cuda_image, _, cuda_release = cuda_base.rpartition(":")
+def test_rocm_base_tracks_ci_vllm_release() -> None:
+    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
    rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE")

    image_ref, separator, image_tag = rocm_base.rpartition(":")
    assert separator, f"expected a tagged ROCm base image, got {rocm_base}"
-    assert cuda_image == "vllm/vllm-openai"
-    assert image_ref == "vllm/vllm-openai-rocm"
-    assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}"
-    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image
-    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}"
+    assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref
+    assert image_tag == ci_release


 def test_rocm_defaults_to_prebuilt_base_image() -> None:
@@ -108,11 +103,11 @@ def test_rocm_build_context_excludes_git_history() -> None:
    assert ".git" in patterns


-def test_rocm_source_ref_tracks_cuda_vllm_release() -> None:
-    cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION")
+def test_rocm_source_ref_tracks_ci_vllm_release() -> None:
+    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
    rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH")

-    assert rocm_source_ref == f"v{cuda_release}"
+    assert rocm_source_ref == ci_release


 def test_rocm_dockerfile_contains_vllm_api_canary() -> None:
@@ -172,11 +172,7 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
    env_names = {item["name"] for item in container["env"]}
    assert "VLLM_CI_HF_TOKEN" in env_names
    assert "HF_TOKEN" not in env_names
-    assert step["commands"] == [
-        CUDA_HF_TOKEN_EXPORT,
-        "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
-        "pytest -sv tests/example",
-    ]
+    assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]


 def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
@@ -243,10 +239,6 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
        assert "VLLM_CI_HF_TOKEN" in env_names
        assert "HF_TOKEN" not in env_names
        assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
-        assert step["commands"][1:] == [
-            "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
-            "pytest -sv tests/example",
-        ]


 def _gpu_limit(step: dict) -> int:
@@ -571,56 +571,6 @@ class TestPipelineDiscovery:
        assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
        assert "qwen3_tts" in OMNI_PIPELINES

-    @pytest.mark.parametrize(
-        "cuda,major,memory_gib,mps,optimized",
-        [
-            (False, None, 0, True, False),
-            (True, None, 141, True, False),
-            (True, 8, 141, True, False),
-            (True, 9, 80, True, False),
-            (True, 10, 180, True, False),
-            (True, 9, 141, False, False),
-            (True, 9, 141, True, True),
-        ],
-    )
-    def test_cosyvoice3_device_default_and_explicit_overrides(
-        self, monkeypatch, cuda, major, memory_gib, mps, optimized
-    ):
-        from vllm.platforms import current_platform
-        from vllm.platforms.interface import DeviceCapability
-
-        from vllm_omni.platforms import current_omni_platform
-
-        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
-        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
-        capability = DeviceCapability(major, 0) if major is not None else None
-        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
-        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
-        monkeypatch.setattr(
-            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
-            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
-        )
-        pipeline = resolve_pipeline_config("cosyvoice3")
-        assert pipeline is not None
-        config = VllmOmniConfig.from_pipeline_config(pipeline)
-        talker, codec = config.stage_configs
-        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
-        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
-        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
-        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
-        if optimized:
-            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
-            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
-            override = VllmOmniConfig.from_pipeline_config(
-                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
-            )
-            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
-        explicit = VllmOmniConfig.from_pipeline_config(
-            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
-        )
-        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
-        assert explicit.stage_by_id(0).model_config.hf_overrides is None
-
    def test_registry_resolver_qwen3_omni_all_stages(self):
        """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
        pipeline = resolve_pipeline_config(
@@ -2978,17 +2928,9 @@ class TestPlatformOverrides:

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

-    @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):
+    @pytest.mark.parametrize("platform", ["cuda", "npu"])
+    def test_recommended_native_runner_platform_defaults(self, platform):
+        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
        deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
        deploy = _apply_platform_overrides(deploy, platform=platform)
        expect_v2 = platform == "cuda"
@@ -3004,60 +2946,6 @@ 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,161 +35,3 @@ 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: 95,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
        EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 7,
+        EnvironmentVariableCategory.INTERNAL: 6,
    }

    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: 66,
+        ModelEnvironmentVariableDisposition.PROMOTE: 46,
        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
        ModelEnvironmentVariableDisposition.EXTERNAL: 0,
        ModelEnvironmentVariableDisposition.INTERNALIZE: 17,
@@ -1046,7 +1046,6 @@ 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",
@@ -1,258 +0,0 @@
-# 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.kv_holding_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()
-    waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args
-    assert list(waiting) == []
-    assert running is s.running
-    assert ready == {"0"} and finished == 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
@@ -1,107 +0,0 @@
-# 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))
@@ -11,9 +11,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm import SamplingParams
-from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 from vllm.v1.core.sched.interface import PauseState
-from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus

@@ -55,14 +53,12 @@ class FakeAdapter:
 def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
    scheduler.max_num_scheduled_tokens = 8
-    scheduler.max_num_active_reqs = 1
    scheduler.max_num_running_reqs = 1
    scheduler._pause_state = PauseState.UNPAUSED
    scheduler.running = []
    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
    scheduler.waiting.add_request(waiting_request)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
    scheduler.requests = {waiting_request.request_id: waiting_request}
    scheduler.policy = SchedulingPolicy.FCFS
    scheduler.chunk_transfer_adapter = FakeAdapter()
@@ -90,7 +86,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
        get_manager_metadata=lambda: None,
    )
    scheduler.connector = None
-    scheduler.aux_output_connector = None
    scheduler.ec_connector = None
    scheduler.prev_step_scheduled_req_ids = set()
    scheduler._pending_finish_reqs = []
@@ -110,45 +105,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
    return scheduler


-def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None:
-    request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
-    request.block_hashes.append(b"block-hash")
-    scheduler = _make_generation_scheduler(request, use_v2_model_runner=True)
-    scheduler.aux_output_connector = AuxOutputSchedulerConnector()
-
-    output = scheduler._build_generation_scheduler_output(
-        new_reqs_data=[],
-        cached_reqs_data=CachedRequestData.make_empty(),
-        num_scheduled_tokens={request.request_id: 1},
-        scheduled_spec_decode_tokens={},
-        scheduled_encoder_inputs={},
-        num_common_prefix_blocks=[0],
-    )
-
-    metadata = output.aux_output_connector_metadata
-    assert metadata is not None
-    assert metadata.requests == {request.request_id: 0}
-    assert list(metadata.block_hashes[request.request_id]) == request.block_hashes
-
-
-@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)])
-def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None:
-    fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
-    holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None)
-    holding.num_computed_tokens = 1
-    scheduler = _make_generation_scheduler(fresh)
-    scheduler.max_num_running_reqs = runner_slots
-    scheduler.max_num_active_reqs = active_limit
-    scheduler.requests[holding.request_id] = holding
-    scheduler.kv_holding_waiting.add_request(holding)
-
-    output = scheduler.schedule()
-
-    assert set(output.num_scheduled_tokens) == {holding.request_id}
-    assert list(scheduler.waiting) == [fresh]
-    assert list(scheduler.kv_holding_waiting) == []
-
-
 def _chunk_request(request_id, **kwargs):
    defaults = dict(
        request_id=request_id,
@@ -167,13 +123,7 @@ def _chunk_request(request_id, **kwargs):
        record_event=lambda *args, **kwargs: None,
    )
    defaults.update(kwargs)
-    return _HashableChunkRequest(**defaults)
-
-
-class _HashableChunkRequest(SimpleNamespace):
-    # Deferred queues require identity hashing; SimpleNamespace types it as None.
-    __hash__ = object.__hash__  # type: ignore[assignment]
-    __eq__ = object.__eq__
+    return SimpleNamespace(**defaults)


 def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
@@ -188,8 +138,7 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
    )
    scheduler.running = []
    scheduler.waiting = create_request_queue(scheduler.policy)
-    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)

    completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
    completed.status = RequestStatus.RUNNING
@@ -343,7 +292,6 @@ def _express_scheduler():
    scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
    scheduler.waiting.add_request(first)
    scheduler.requests[first.request_id] = first
-    scheduler.max_num_active_reqs = 4
    scheduler.max_num_running_reqs = 4
    scheduler._native_data_plane = True
    scheduler.chunk_transfer_adapter = None
@@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises
 admits in a finished state -- which kills the stage's engine core, not just
 the request.

-Regression for the sweep missing ``kv_holding_waiting``, the third queue
+Regression for the sweep missing ``skipped_waiting``, the third queue
 upstream admits from. An aborted duplex session parked there was re-selected
 by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
 stage.
@@ -49,30 +49,28 @@ class _StubRequest:
        return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)


-def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()):
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
    """Bare scheduler with just the surface the sweep reads."""
    scheduler = OmniARScheduler.__new__(OmniARScheduler)
    scheduler.waiting = create_request_queue(policy)
    for req in waiting:
        scheduler.waiting.add_request(req)
-    scheduler.kv_holding_waiting = create_request_queue(policy)
-    scheduler.deferred_waiting = set()
-    for req in kv_holding_waiting:
-        scheduler.kv_holding_waiting.add_request(req)
+    scheduler.skipped_waiting = create_request_queue(policy)
+    for req in skipped_waiting:
+        scheduler.skipped_waiting.add_request(req)
    scheduler.running = list(running)
    return scheduler


@pytest.mark.parametrize("policy", _POLICIES)
-def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
+def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
    monkeypatch: pytest.MonkeyPatch,
    policy: SchedulingPolicy,
 ) -> None:
    """The public schedule path must remove aborted skipped requests first."""
    aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
-    scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted])
+    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
    scheduler.chunk_transfer_adapter = None
-    scheduler.max_num_active_reqs = 8
    scheduler.max_num_running_reqs = 8
    scheduler.input_coordinator = None
    scheduler.num_waiting_for_streaming_input = 0
@@ -87,7 +85,7 @@ def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
    sentinel = object()

    def upstream_schedule(_self, _throttle_prefills: bool = False):
-        assert list(scheduler.kv_holding_waiting) == []
+        assert list(scheduler.skipped_waiting) == []
        return sentinel

    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
@@ -1,16 +1,10 @@
-# 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

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


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


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

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


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

-    assert not omits_kv
-    assert omits_chunk
+    assert not omits_transfer
    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.deferred_waiting = set()
-    scheduler.kv_holding_waiting = []
-    scheduler.log_stats = False
-    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
-
-    class Session(SimpleNamespace):
-        __hash__ = object.__hash__
-
-    session = Session(
-        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)
@@ -11,14 +11,12 @@ from types import MethodType, SimpleNamespace

 import numpy as np
 import pytest
-import torch
 from vllm.v1.engine import FinishReason
 from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus

 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.outputs import OmniModelRunnerOutput

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

@@ -134,13 +132,11 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
    scheduler = SimpleNamespace(
        perf_metrics=None,
        connector=None,
-        aux_output_connector=None,
        chunk_transfer_adapter=None,
        requests={request.request_id: request for request in requests},
        running=list(requests),
        waiting=_RequestQueue(),
-        kv_holding_waiting=_RequestQueue(),
-        deferred_waiting=set(),
+        skipped_waiting=_RequestQueue(),
        structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True),
        transfer_triggered_requests=set(),
        active_kv_transfers=set(),
@@ -296,36 +292,6 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
    np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])


-@pytest.mark.parametrize("scored", [False, True])
-def test_fixed_token_scores_reach_engine_output(scored) -> None:
-    request = _Request("req")
-    request.sampling_params.num_logprobs = None
-    scheduler = _make_scheduler_stub([request])
-    _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
-    scheduler_output = SimpleNamespace(
-        num_scheduled_tokens={"req": 3},
-        scheduled_spec_decode_tokens={},
-        num_invalid_spec_tokens=0,
-    )
-    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
-    model_runner_output = OmniModelRunnerOutput(
-        req_ids=["req"],
-        req_id_to_index={"req": 0},
-        sampled_token_ids=[[7]],
-        prompt_logprobs_dict={},
-        prompt_token_id_logprobs_dict={"req": scores} if scored else {},
-    )
-
-    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
-    (output,) = outputs[0].outputs
-
-    assert output.new_token_ids == [7]
-    if scored:
-        torch.testing.assert_close(output.prompt_token_id_logprobs, scores)
-    else:
-        assert output.prompt_token_id_logprobs is None
-
-
 def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
    return SimpleNamespace(
        sampled_token_ids=[[]],
@@ -40,7 +40,6 @@ 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

@@ -73,8 +72,7 @@ def _replace_streaming_session(session: Request) -> None:
    sched.num_waiting_for_streaming_input = 0
    sched.log_stats = False
    sched.chunk_transfer_adapter = None
-    sched.kv_holding_waiting = set()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = set()
    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
    update = StreamingUpdate(
        mm_features=None,
@@ -105,11 +103,10 @@ 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, spec=None) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
    """Feed one model-runner frame through update_from_output.

    Returns True when the frame's tokens were delivered (the append ran),
@@ -117,7 +114,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 = {session.request_id: spec} if spec else {}
+    scheduler_output.scheduled_spec_decode_tokens = {}
    scheduler_output.num_invalid_spec_tokens = 0

    model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -137,27 +134,6 @@ 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."""
@@ -249,80 +225,3 @@ 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,7 +26,6 @@ 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

@@ -42,8 +41,7 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS
    sched.num_waiting_for_streaming_input = 0
    sched.log_stats = False
    sched.chunk_transfer_adapter = None
-    sched.kv_holding_waiting = set()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = set()
    sched._free_request_blocks = MagicMock()
    sched.encoder_cache_manager = MagicMock()
    sched._inflight_prefills = set()
@@ -177,7 +175,6 @@ 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}
@@ -315,7 +312,6 @@ 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}
@@ -404,7 +400,6 @@ 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}
@@ -1170,8 +1165,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
    sched.running = []
    sched.waiting = queue()
-    sched.kv_holding_waiting = queue(session)
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = queue(session)
    sched.num_waiting_for_streaming_input = 1
    sched.chunk_transfer_adapter = SimpleNamespace(
        receives_chunks=True,
@@ -1183,7 +1177,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:

    assert session.status == RequestStatus.WAITING
    assert session in sched.waiting.requests
-    assert session not in sched.kv_holding_waiting.requests
+    assert session not in sched.skipped_waiting.requests
    assert sched.num_waiting_for_streaming_input == 0
    assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests

@@ -1210,8 +1204,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
        receives_chunks=receives_chunks,
        segment_finished_requests={session.request_id},
    )
-    sched.kv_holding_waiting = MagicMock()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = MagicMock()
    sched._enqueue_waiting_request = MagicMock()

    sched._resume_downstream_chunk_receiver(session)
@@ -1219,7 +1212,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
    assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
    assert sched.num_waiting_for_streaming_input == 1
    assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
-    sched.kv_holding_waiting.remove_requests.assert_not_called()
+    sched.skipped_waiting.remove_requests.assert_not_called()
    sched._enqueue_waiting_request.assert_not_called()


@@ -1275,8 +1268,7 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None:
    session.num_computed_tokens = 6
    session.num_output_placeholders = 0
    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    sched.kv_holding_waiting.add_request(session)
-    sched.deferred_waiting.add(session)
+    sched.skipped_waiting.add_request(session)
    sched.num_waiting_for_streaming_input = 1


@@ -1299,7 +1291,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi
    assert session.status == RequestStatus.FINISHED_ERROR
    assert session.request_id not in sched.requests
    assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
    assert sched.num_waiting_for_streaming_input == 0
    sched._free_request_blocks.assert_called_once_with(session)
    assert sched.finished_req_ids == {session.request_id}
@@ -1407,8 +1399,7 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
    sched.max_model_len = max_model_len
    sched.policy = SchedulingPolicy.FCFS
    sched.waiting = create_request_queue(sched.policy)
-    sched.kv_holding_waiting = create_request_queue(sched.policy)
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = create_request_queue(sched.policy)
    sched.running = []
    sched.requests = {}
    sched._free_request = MagicMock()  # needs a KV manager; not what this covers
@@ -1437,7 +1428,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -

    assert finished is True
    assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
    assert session.status == RequestStatus.FINISHED_ERROR
    assert session.resumable is False
    assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
@@ -1456,8 +1447,7 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
    finished = sched._handle_stopped_request(session)

    assert finished is False
-    assert list(sched.kv_holding_waiting) == [session]
-    assert session not in sched.deferred_waiting
+    assert len(sched.waiting) == 1
    assert session.status == RequestStatus.WAITING
    assert not getattr(sched, "_streaming_context_overflow", {})
    sched._free_request.assert_not_called()
@@ -1479,7 +1469,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:

    assert finished is True
    assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
    assert session.status == RequestStatus.FINISHED_ERROR
    assert session.resumable is False
    sched._free_request.assert_not_called()
@@ -1502,7 +1492,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc
    finished = sched._handle_stopped_request(session)

    assert finished is True
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
    assert sched.num_waiting_for_streaming_input == 0
    assert session.status == RequestStatus.FINISHED_ERROR
    sched._free_request.assert_not_called()
@@ -1535,7 +1525,6 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
    sched.running = []
    sched._native_data_plane = native
    sched.use_v2_model_runner = native
-    sched.max_num_active_reqs = 8
    sched.max_num_running_reqs = 8
    sched.input_coordinator = (
        SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
@@ -1562,7 +1551,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
    observed_limits: list[int] = []

    def fake_schedule(self, _throttle_prefills=False):
-        observed_limits.append(self.max_num_active_reqs)
+        observed_limits.append(self.max_num_running_reqs)
        return SimpleNamespace(scheduled_new_reqs=[])

    monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
@@ -1570,4 +1559,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
    sched.schedule()

    assert observed_limits == [7 if native else 8]
-    assert sched.max_num_active_reqs == 8
+    assert sched.max_num_running_reqs == 8
@@ -79,19 +79,16 @@ def _make_scheduler(
        else:
            hf_config.voxcpm2_runtime_config = runtime_config
    sched.vllm_config = SimpleNamespace(model_config=model_config)
-    # Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the
+    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
    # abort sweep at the top of ``schedule()`` reads it.
-    sched.kv_holding_waiting = _MockQueue()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = _MockQueue()
    return sched


-@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
-def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
    scheduler = _make_scheduler()
    scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue()
-    setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])

    assert scheduler._should_defer_waiting_for_unified_decode_graph()

@@ -147,17 +144,12 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa

 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
    scheduler = _make_scheduler()
-    scheduler.max_num_active_reqs = 8
    scheduler.max_num_running_reqs = 8
    scheduler.running = [_MockRequest("decode")]
    original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
    deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
    original_waiting = _MockQueue([original_waiting_req])
    scheduler.waiting = original_waiting
-    kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING)
-    original_kv_holding = _MockQueue([kv_holder])
-    scheduler.kv_holding_waiting = original_kv_holding
-    preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED)
    scheduler.policy = "fcfs"
    scheduler.chunk_transfer_adapter = None
    scheduler.input_coordinator = None
@@ -168,10 +160,7 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
    def fake_upstream_schedule(self, throttle_prefills: bool = False):
        assert self.waiting is not original_waiting
        assert not self.waiting
-        assert self.kv_holding_waiting is not original_kv_holding
-        assert not self.kv_holding_waiting
        self.waiting.add_request(deferred_by_upstream)
-        self.kv_holding_waiting.add_request(preempted_by_upstream)
        raise RuntimeError("stop before output wrapping")

    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
@@ -180,6 +169,3 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
        scheduler.schedule()

    assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]
-    assert scheduler.waiting is original_waiting
-    assert scheduler.kv_holding_waiting is original_kv_holding
-    assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder]
@@ -23,12 +23,10 @@ def _scheduler(monkeypatch, waiting, running=()):
    scheduler.waiting = create_request_queue(scheduler.policy)
    for request in waiting:
        scheduler.waiting.add_request(request)
-    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
    scheduler.running = list(running)
    scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
    scheduler.max_model_len = 128
-    scheduler.max_num_active_reqs = 2
    scheduler.max_num_running_reqs = 2
    scheduler.max_num_scheduled_tokens = 256
    scheduler.block_size = 16
@@ -52,11 +52,10 @@ class _SchedulerStub(OmniGenerationScheduler):
        self.num_waiting_for_streaming_input = 0
        self.log_stats = log_stats
        self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
-        self.kv_holding_waiting = _SkippedWaitingStub()
-        self.deferred_waiting: set[Request] = set()
+        self.skipped_waiting = _SkippedWaitingStub()

    def _enqueue_waiting_request(self, session: Request) -> None:
-        raise AssertionError("unexpected enqueue for kv_holding_waiting miss")
+        raise AssertionError("unexpected enqueue for skipped_waiting miss")


 class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
@@ -135,8 +134,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
    sched._handle_stopped_request.side_effect = rearm_resumable_request
    sched.running = [session]
    sched.waiting = MagicMock()
-    sched.kv_holding_waiting = MagicMock()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = MagicMock()
    sched.structured_output_manager.accept_tokens.return_value = True
    # Async scheduling can observe the segment boundary in the next schedule()
    # before this model output is applied and defer the same request for finish.
@@ -208,8 +206,7 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
        OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted)
    )
    sched.waiting = MagicMock()
-    sched.kv_holding_waiting = MagicMock()
-    sched.deferred_waiting = set()
+    sched.skipped_waiting = MagicMock()
    sched.structured_output_manager.accept_tokens.return_value = False
    sched._pending_finish_reqs = []
    sched.recompute_kv_load_failures = False
@@ -353,8 +350,6 @@ class _RealignSchedulerStub(OmniSchedulerMixin):
        self.requests = requests
        self.running = running
        self.waiting = waiting
-        self.kv_holding_waiting: list[Request] = []
-        self.deferred_waiting: set[Request] = set()


 class TestRealignRequestStatusToQueues:
@@ -123,21 +123,13 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request

    sched = MagicMock()
    sched.max_num_scheduled_tokens = 8
-    sched.max_num_active_reqs = 8
    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]
    sched.waiting = []  # falsy: the waiting loop never runs
-    sched.kv_holding_waiting = []
-    sched.deferred_waiting = set()
    sched.policy = SchedulingPolicy.FCFS
    sched.chunk_transfer_adapter = None
    sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 """Integration tests for the post-finish defensive purge of
 ``self.running`` placed in
 ``OmniARScheduler.finish_requests`` /
@@ -55,8 +52,6 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting):
    scheduler.requests = requests
    scheduler.running = running
    scheduler.waiting = waiting
-    scheduler.kv_holding_waiting = []
-    scheduler.deferred_waiting = set()
    return scheduler


@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 from __future__ import annotations

 from types import SimpleNamespace
@@ -52,8 +49,6 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
    scheduler.requests = {}
    scheduler.running = []
    scheduler.waiting = []
-    scheduler.kv_holding_waiting = []
-    scheduler.deferred_waiting = set()

    def fake_finish_requests(self, request_ids, finished_status):
        assert request_ids == ["req-a", "req-b"]
@@ -81,7 +76,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": (1, True, False)}
+    scheduler._omits_kv_transfer_cache = {"req-free": True}
    scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
    scheduler.finished_req_ids = set()
    scheduler.finished_req_ids_dict = None
@@ -7,10 +7,7 @@ from typing import Any

 import pytest
 from vllm.config import SchedulerConfig, VllmConfig
-from vllm.sampling_params import SamplingParams
-from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import FinishReason
-from vllm.v1.request import Request

 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
@@ -21,7 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


 class _Scheduler(OmniSchedulerMixin):
-    _is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status)
+    pass


 def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
@@ -123,11 +120,7 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
    calls: list[tuple[Any, ...]] = []
    scheduler = _Scheduler()
-    waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None)
-    holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None)
-    holding.num_computed_tokens = 1
-    scheduler.waiting = [waiting]
-    scheduler.kv_holding_waiting = [holding]
+    scheduler.waiting = ["waiting"]
    scheduler.running = ["running"]
    scheduler.requests = {"request": object()}
    scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
@@ -144,16 +137,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
    scheduler.chunk_transfer_adapter = SimpleNamespace(
        receives_chunks=True,
        process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
-            ("process", list(waiting), running, scheduler_requests)
+            ("process", waiting, running, scheduler_requests)
        ),
        restore_queues=lambda waiting, running, scheduler_requests: calls.append(
-            ("restore-chunks", list(waiting), running, scheduler_requests)
+            ("restore-chunks", waiting, running, scheduler_requests)
        ),
        collect_timed_out_request_ids=_collect_timed_out,
        collect_failed_send_request_ids=_collect_failed_sends,
    )
    scheduler.input_coordinator = SimpleNamespace(
-        restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running))
+        restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
    )

    scheduler._process_pending_omni_inputs("ar")
@@ -162,13 +155,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
    assert calls == [
        ("consume", "ar"),
        ("timeouts",),
-        ("process", [holding, waiting], scheduler.running, scheduler.requests),
+        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
        # The chunk deadline runs after chunks are applied, so a chunk that
        # arrived this cycle resets the clock before it is measured (R1.1).
        ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
        ("failed-sends",),
-        ("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests),
-        ("restore-full", [holding, waiting], scheduler.running),
+        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("restore-full", scheduler.waiting, scheduler.running),
    ]


@@ -5,13 +5,8 @@
 from types import SimpleNamespace

 import pytest
-from vllm import SamplingParams
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
-from vllm.v1.request import Request, RequestStatus

-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
 from vllm_omni.outputs import OmniConnectorOutput

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -20,8 +15,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
    scheduler = OmniSchedulerMixin()
    scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
-    scheduler.kv_holding_waiting = []
-    scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status
    coordinator = SimpleNamespace(
        _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
    )
@@ -39,92 +32,7 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
    coordinator.update_request_metadata.assert_called_once_with(
        scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
    )
-    coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"})
-    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
+    coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
    scheduler._consume_pending_connector_output(model_mode="ar")
    assert coordinator.update_request_metadata.call_count == 1
-    coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set())
-    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
-
-
-@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY])
-@pytest.mark.parametrize("async_chunk", [False, True])
-def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk):
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(policy)
-    scheduler.kv_holding_waiting = create_request_queue(policy)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = OmniSchedulingCoordinator(
-        stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk
-    )
-    scheduler._init_omni_connector_output_inbox()
-    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    holder.num_computed_tokens = 2
-    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
-    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
-    scheduler._enqueue_waiting_request(fresh)
-    scheduler._enqueue_waiting_request(holder)
-
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert not scheduler.waiting and not scheduler.kv_holding_waiting
-    expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT
-    assert holder.status == fresh.status == expected_status
-    coordinator = scheduler.input_coordinator
-    registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations
-    assert [entry.request_id for entry in registrations] == ["holder", "fresh"]
-
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-    ready_output = (
-        OmniConnectorOutput(chunk_ready_req_ids={"holder"})
-        if async_chunk
-        else OmniConnectorOutput(stage_recv_req_ids={"holder"})
-    )
-    scheduler.enqueue_omni_connector_output(ready_output)
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert holder.status == RequestStatus.WAITING
-    assert fresh.status == expected_status
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert not scheduler.waiting
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-
-@pytest.mark.parametrize("async_chunk", [False, True])
-@pytest.mark.parametrize(
-    "status",
-    [
-        RequestStatus.WAITING_FOR_STREAMING_REQ,
-        RequestStatus.WAITING_FOR_REMOTE_KVS,
-        RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
-    ],
-)
-def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status):
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    request.status = status
-    request.num_computed_tokens = 2
-    scheduler.requests = {request.request_id: request}
-    scheduler._enqueue_waiting_request(request)
-    scheduler.input_coordinator = OmniSchedulingCoordinator(
-        stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk
-    )
-    scheduler._init_omni_connector_output_inbox()
-
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert request.status == status
-    assert list(scheduler.kv_holding_waiting) == [request]
-    assert scheduler.deferred_waiting == {request}
-    assert not scheduler.input_coordinator.pending_chunk_registrations
-    assert not scheduler.input_coordinator.pending_input_registrations
+    coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
@@ -59,8 +59,7 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp
    scheduler.requests = requests
    scheduler.running = running
    scheduler.waiting = waiting
-    scheduler.kv_holding_waiting = skipped if skipped is not None else []
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = skipped if skipped is not None else []
    scheduler.num_waiting_for_streaming_input = counter
    return scheduler

@@ -98,8 +97,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status(
    assert scheduler.num_waiting_for_streaming_input == 0


-def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None:
-    """``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
+def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
    so the resync has to look at both or it under-counts and reintroduces the
    spin the counter exists to prevent."""
    parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
@@ -129,11 +128,9 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
    scheduler.requests = {parked.request_id: parked}
    scheduler.running = [parked]
    scheduler.waiting = []
-    scheduler.kv_holding_waiting = []
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = []
    scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
    scheduler.policy = "fcfs"
-    scheduler.max_num_active_reqs = 8
    scheduler.max_num_running_reqs = 8

    seen: dict[str, int] = {}
@@ -164,13 +161,12 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -
    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
    scheduler.waiting.add_request(aborted)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
-    scheduler.kv_holding_waiting.add_request(aborted)
+    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.skipped_waiting.add_request(aborted)
    scheduler.running = [aborted]

    scheduler._drop_aborted_queued_requests()

    assert list(scheduler.waiting) == []
-    assert list(scheduler.kv_holding_waiting) == []
+    assert list(scheduler.skipped_waiting) == []
    assert scheduler.running == []
@@ -216,40 +216,44 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────

+# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
+_REALTIME_WS_MODEL_MISMATCH = object()
+_REALTIME_WS_INVALID_AUDIO_APPEND = object()
+

@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize(
    "send_frames_spec, ws_error_code, err_message",
    [
-        pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
+        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
        pytest.param(
            json.dumps({"type": "session.update"}),
            "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_missing_session",
+            ("model", "Missing required field"),
+            id="session_update_missing_model",
        ),
        pytest.param(
-            json.dumps({"type": "session.update", "session": []}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_invalid_session",
+            _REALTIME_WS_MODEL_MISMATCH,
+            "model_not_found",
+            ("model", "does not exist"),
+            id="session_update_model_mismatch",
        ),
        pytest.param(
-            json.dumps({"type": "input_audio_buffer.commit"}),
-            "invalid_request_error",
-            "Input audio buffer is empty",
-            id="commit_without_audio",
+            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
+            "model_not_validated",
+            ("session.update", "validate the model"),
+            id="commit_without_model_validation",
        ),
        pytest.param(
            json.dumps({"type": "unknown.event"}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
+            "unknown_event",
+            "Unknown event type",
            id="unknown_event_type",
        ),
        pytest.param(
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-            "invalid_request_error",
-            "Invalid base64 audio data",
+            _REALTIME_WS_INVALID_AUDIO_APPEND,
+            "invalid_audio",
+            "Invalid audio data",
            id="invalid_audio_append",
        ),
    ],
@@ -262,12 +266,27 @@ def test_realtime_invalid_requests(
    ws_error_code: str,
    err_message: str | tuple[str, ...],
 ) -> None:
+    send_frames: Any
+    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
+        send_frames = json.dumps(
+            {
+                "type": "session.update",
+                "model": "this-model-is-not-served-on-this-instance",
+            }
+        )
+    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
+        send_frames = [
+            json.dumps({"type": "session.update", "model": omni_server.model}),
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+        ]
+    else:
+        send_frames = send_frames_spec
    online_client.send_realtime_ws_request(
        {
-            "send_frames": send_frames_spec,
+            "send_frames": send_frames,
            "timeout": 120,
            "ws_max_size": None,
-            "ws_skip_types": ["session.created", "conversation.created"],
+            "ws_skip_types": ["session.created"],
            "ws_json_type": "error",
            "ws_error_code": ws_error_code,
            "err_message": err_message,
@@ -376,17 +376,10 @@ 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 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.
+    # One tiny LingBot-like session requires 128 bytes; two require 192.
    runner = make_runner(
        pipeline,
-        available_bytes=2048,
+        available_bytes=128,
        gpu_memory_fraction=1.0,
    )
    kv = runner.kv_cache
@@ -672,106 +665,3 @@ 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,74 +625,6 @@ 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,7 +23,6 @@ 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]
@@ -37,26 +36,10 @@ NEG = "negative"


 def make_state(
-    *,
-    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,
+    *, num_layers=1, window_chunks=2, 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=chunk_size,
-        window_chunks=window_chunks,
-        sink_chunks=sink_chunks,
-        reset_at_boundary=reset_at_boundary,
-        reuse_history_staging=reuse_history_staging,
+        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
    )
    kv = ARDiffusionKVCache(
        cfg,
@@ -126,13 +109,11 @@ def _commit_video_span(
    n_chunks: int,
    dtype: torch.dtype,
    device: torch.device,
-    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    span = n_chunks * chunk_size
-    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
+    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
    ctx.ensure_video_slots(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)
+    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)
    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)
@@ -566,644 +547,6 @@ 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])
@@ -1,359 +0,0 @@
-# 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"}
@@ -1,118 +0,0 @@
-# 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,7 +21,6 @@ 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]
@@ -433,35 +432,6 @@ 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 = []
@@ -580,199 +550,6 @@ 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) -------------


@@ -1,308 +0,0 @@
-# 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),
-    )
@@ -329,9 +329,7 @@ def _uneven_cfg_sharded_worker(
        init_method=init_method,
        rank=rank,
        world_size=4,
-        # Spawned workers import torch/vLLM before joining; allow startup
-        # skew without changing the collective deadlock checks below.
-        timeout=datetime.timedelta(seconds=60),
+        timeout=datetime.timedelta(seconds=10),
    )
    try:
        fs_groups = [
@@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
    # Independent expected native extras: a crossing item has a negative
    # relative offset; a block can include the end of one item and the next.
    first, second = (feature.identifier for feature in row.mm_features)
-    extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
+    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
    parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
    expected = []
    for start, keys in zip(range(0, 12, 4), extras, strict=True):
@@ -31,10 +31,6 @@ def producer(monkeypatch):
    # vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34,
    # #56033), the MLA/Mamba flags and the head-resharding validator.
    worker.pp_size = 1
-    worker.opaque_packed_storages = set()
-    worker._prepared_transfer_regions = {}
-    worker.region_shared_groups = []
-    worker.region_row_offsets = []
    worker.use_mla = False
    worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False)
    worker._validate_head_resharding_layout = Mock(return_value=None)
@@ -1253,15 +1253,14 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
        speculative_config=None,
    )

-    def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
-        # Exercise the real upstream selector, replacing only platform resolution.
+    def select_backend(*, backend, attn_selector_config, num_heads):
+        # Exercise the real 0.29 selector, replacing only platform resolution.
        # Ulysses has already gathered tokens and sharded heads before the
        # native kernel; the MoE PCP mapping must not request PCP attention.
        assert not attn_selector_config.use_pcp
        assert not attn_selector_config.use_dcp
        assert attn_selector_config.use_non_causal
        assert num_heads == 4
-        assert _run_kv_cache_dtype == config.cache_config.cache_dtype
        selected_backends.append((backend, config.attention_config.backend_per_kind))
        if selection_fails:
            raise ValueError("test backend unavailable")
@@ -14,7 +14,6 @@ from typing import Literal
 import pytest
 import torch

-from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -57,8 +56,6 @@ def _init_worker(
            "MASTER_PORT": str(master_port),
        }
    )
-    if device_kind == "cpu":
-        configure_cpu_collective_worker()
    backend = "gloo" if device_kind == "cpu" else None
    init_distributed_environment(backend=backend)

@@ -230,10 +227,6 @@ def _run_ring_p2p(
            atol=atol,
            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
        )
-        if device_kind == "cpu":
-            # P2P completion only synchronizes neighbors. Keep fast CPU ranks
-            # alive until every rank has finished constructing and using groups.
-            torch.distributed.barrier()
    finally:
        destroy_distributed_env()

@@ -250,7 +243,7 @@ def _spawn_4d_identity(
    device_kind: DeviceKind,
    master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        _run_4d_identity,
        args=(
            world_size,
@@ -264,7 +257,6 @@ def _spawn_4d_identity(
            master_port,
        ),
        nprocs=world_size,
-        device_kind=device_kind,
    )


@@ -280,7 +272,7 @@ def _spawn_5d_identity(
    device_kind: DeviceKind,
    master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        _run_5d_identity,
        args=(
            world_size,
@@ -294,7 +286,6 @@ def _spawn_5d_identity(
            master_port,
        ),
        nprocs=world_size,
-        device_kind=device_kind,
    )


@@ -308,11 +299,10 @@ def _spawn_ring_p2p(
    device_kind: DeviceKind,
    master_port: int,
 ) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        _run_ring_p2p,
        args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
        nprocs=world_size,
-        device_kind=device_kind,
    )


@@ -509,9 +499,8 @@ def _run_fused_qkv_matches_three_4d(
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
 def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        _run_fused_qkv_matches_three_4d,
        args=(world_size, "cpu", 29613),
        nprocs=world_size,
-        device_kind="cpu",
    )
@@ -1,200 +0,0 @@
-# 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,13 +3,11 @@
 """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

@@ -20,7 +18,6 @@ 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]

@@ -341,92 +338,3 @@ 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]))
@@ -19,7 +19,6 @@ from vllm.utils.network_utils import get_file_store_init_method
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors

 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
-from tests.helpers.distributed import collective_context, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -40,25 +39,6 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


-@pytest.fixture(autouse=True)
-def _restore_collective_platforms():
-    """CPU baselines run in the parent process; retain its original platforms."""
-    import vllm.platforms
-
-    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
-
-    original_platform = vllm.platforms.current_platform
-    original_omni_platforms = {
-        module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
-    }
-    try:
-        yield
-    finally:
-        vllm.platforms.current_platform = original_platform
-        for module, platform in original_omni_platforms.items():
-            module.current_omni_platform = platform
-
-
 def _cleanup_distributed() -> None:
    """Destroy omni distributed state left dangling after a test.

@@ -429,10 +409,6 @@ def init_dist(
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)
-    else:
-        from tests.helpers.distributed import configure_cpu_collective_worker
-
-        configure_cpu_collective_worker()
    backend = "gloo" if device_kind == "cpu" else None
    init_distributed_environment(
        world_size=world_size,
@@ -524,14 +500,13 @@ def isend_irecv_worker(


 def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    q = manager.Queue()
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        isend_irecv_worker,
        args=(pp_size, init_method, device_kind, q),
        nprocs=pp_size,
-        device_kind=device_kind,
    )
    results = {label: tensor for label, tensor in [q.get(), q.get()]}
    torch.testing.assert_close(
@@ -681,16 +656,15 @@ def _run_predict_noise(

    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)

-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    pp_q = manager.Queue()

    world_size = pp_size * cfg_size
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        predict_noise_worker,
        args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
        nprocs=world_size,
-        device_kind=device_kind,
    )

    pp_out = pp_q.get()
@@ -926,16 +900,15 @@ def _run_scheduler_step(

    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)

-    mp_context = collective_context(device_kind)
+    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    q = manager.Queue()

    world_size = pp_size * cfg_size
-    start_collective_workers(
+    torch.multiprocessing.spawn(
        scheduler_step_worker,
        args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
        nprocs=world_size,
-        device_kind=device_kind,
    )

    result = q.get()
@@ -27,7 +27,6 @@ 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
@@ -2090,206 +2089,3 @@ 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,10 +61,7 @@ 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
-    # _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 prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
    assert step_spy.call_count == 2
    capture_spy.assert_not_called()
    assert not wrapper._cache
@@ -87,8 +84,7 @@ 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, 0)
-    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
    assert wrapper.max_graphs == 32


@@ -157,21 +153,18 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
            inputs["ref"],
            inputs["ref_mask"],
        )
-        # 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)
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
        assert key in wrapper._cache

        # The static context was refreshed for this request's conditioning.
        entry = wrapper._cache[key]
-        grid = torch.tensor([0.0, 0.4], device="cuda")
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
        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

@@ -1,139 +0,0 @@
-# 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,9 +127,7 @@ 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

-    # 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)
+    torch.testing.assert_close(got, x, rtol=0, atol=0)
    assert dit.text_cond is None and dit.text_uncond is None


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

 from tests.helpers.mark import hardware_test
-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.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
 from vllm_omni.platforms import current_omni_platform

@@ -247,72 +246,6 @@ 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,18 +291,6 @@ 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()

@@ -504,40 +492,6 @@ 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)
@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
 )
 def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
    from transformers import AutoModel
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym

    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config

    od_config = OmniDiffusionConfig(
        model="dummy-boogu",
@@ -287,14 +287,15 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
        mock_dependencies["mllm_config_loader"].assert_not_called()
        mock_dependencies["mllm_builder"].assert_not_called()
        assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
-    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)


 def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
-    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config

-    quant_config = DiffusionFp8Config(
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    quant_config = Fp8Config(
        ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
    )
    od_config = OmniDiffusionConfig(
@@ -868,93 +868,6 @@ 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

@@ -16,13 +16,17 @@ class _RecordingBlock(nn.Module):


 def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
    from vllm_omni.diffusion.models.flux2 import flux2_transformer
-    from vllm_omni.quantization import build_quant_config

    monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
    monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)

-    quant_config = build_quant_config("fp8")
+    quant_config = Fp8Config(
+        is_checkpoint_fp8_serialized=False,
+        activation_scheme="dynamic",
+    )
    model = flux2_transformer.Flux2Transformer2DModel(
        in_channels=4,
        out_channels=4,
@@ -148,8 +148,7 @@ 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 isinstance(video["video"], np.ndarray)
-    np.testing.assert_array_equal(video["video"], frames.numpy())
+    assert video["video"] is frames
    image = post(
        {"image": torch.ones(4, 4, 3)},
        SimpleNamespace(output_type="pt"),
@@ -163,40 +162,6 @@ 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
@@ -216,9 +181,7 @@ def test_postprocess_preserves_image_contract_for_serving():
        sampling_params=sampling,
        request_id="lingbot-image-contract-test",
    )
-    image = torch.zeros(8, 12, 3)
-    image[4:] = 1.0
-    raw_output = {"image": image}
+    raw_output = {"image": torch.ones(8, 12, 3)}
    processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
    normalized = normalize_diffusion_postprocess_output(processed)

@@ -236,8 +199,6 @@ 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():
@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):

 def test_fp8_scope_and_prefix_propagation(monkeypatch):
    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
    from vllm.model_executor.layers.quantization.utils.quant_utils import (
        is_layer_skipped,
    )

    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config

    monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
    monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
        "blocks.0.adaln_proj.linear",
        "final_layer.adaln_proj.linear",
    }
-    fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
+    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
    model = h3.MiniMaxH3DiTModel(
        _small_od_config(),
        quant_config=fp8_config,
@@ -13,6 +13,7 @@ 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
@@ -54,7 +55,6 @@ 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 offloader and allocator state."""
+    """Test-only worker extension exposing the configured offloader state."""

    model_runner: Any

@@ -86,21 +86,6 @@ 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):
@@ -111,6 +96,13 @@ 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
@@ -125,27 +117,15 @@ def _extract_audio(output: Any) -> np.ndarray | None:
    return np.asarray(audio)


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


 def _measurement_order() -> tuple[bool, bool]:
@@ -164,6 +144,8 @@ 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
@@ -184,45 +166,52 @@ def run_inference(
        if not offload_states:
            raise AssertionError("The offload-state worker probe returned no diffusion worker results")

-        # 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")
+        # 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()

-        # 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")
+        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()

        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 {
-        "peak_allocated_mb": max(float(result["peak_allocated_mb"]) for result in memory_results),
-        "memory_results": memory_results,
+        "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,
        "audio": audio,
        "offload_states": offload_states,
    }
@@ -270,8 +259,12 @@ def _assert_audio_outputs(


 def _print_measurement(label: str, measurement: dict[str, Any]) -> None:
-    print(f"{label}: worker_peak_allocated={measurement['peak_allocated_mb']:.1f} MB")
-    print(f"{label} worker memory: {measurement['memory_results']}")
+    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} offload state: {measurement['offload_states']}")


@@ -305,33 +298,6 @@ 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
@@ -400,8 +366,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["peak_allocated_mb"]
-    layerwise_offload_peak_memory = layerwise_offload["peak_allocated_mb"]
+    no_offload_peak_memory = no_offload["incremental_peak_mb"]
+    layerwise_offload_peak_memory = layerwise_offload["incremental_peak_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 "
@@ -15,74 +15,6 @@ def test_build_quant_config_fp8():
    assert config is not None
    assert config.get_name() == "fp8"
    assert config.activation_scheme == "dynamic"
-    assert not config.is_checkpoint_fp8_serialized
-
-
-def test_build_quant_config_fp8_uses_native_online_quantization():
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config("fp8", ignored_layers=["proj_out"])
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args.linear.weight == kFp8StaticTensorSym
-    assert config.args.moe.weight == kFp8StaticTensorSym
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_build_quant_config_upstream_online_fp8_shorthand():
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args.linear.weight == kFp8StaticTensorSym
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_online_fp8_rejects_static_activation():
-    from vllm_omni.quantization import build_quant_config
-
-    with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
-        build_quant_config("fp8", activation_scheme="static")
-
-
-def test_build_online_quant_config_preserves_explicit_native_args():
-    from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
-    config = build_quant_config("online", args=args)
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args is args
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_build_online_quant_config_requires_args():
-    from vllm_omni.quantization import build_quant_config
-
-    with pytest.raises(ValueError, match="requires quantization config arguments"):
-        build_quant_config("online")
-
-
-def test_serialized_checkpoint_replaces_online_fp8_config():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
-
-    config = resolve_quant_config_from_disk(
-        build_quant_config("fp8"),
-        {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
-    )
-    assert isinstance(config, Fp8Config)
-    assert config.is_checkpoint_fp8_serialized
-    assert config.activation_scheme == "static"


 def test_build_quant_config_none():
@@ -107,7 +39,7 @@ def test_build_quant_config_invalid():
 def test_build_quant_config_dict():
    from vllm_omni.quantization import build_quant_config

-    config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
+    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
    assert config is not None
    assert config.get_name() == "fp8"
    assert config.activation_scheme == "static"
@@ -116,7 +48,7 @@ def test_build_quant_config_dict():
 def test_build_quant_config_dict_not_mutated():
    from vllm_omni.quantization import build_quant_config

-    original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
+    original = {"method": "fp8", "activation_scheme": "static"}
    copy = original.copy()
    build_quant_config(original)
    assert original == copy
@@ -167,7 +99,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated():
    """Inner component dicts should not be mutated by build_quant_config."""
    from vllm_omni.quantization import build_quant_config

-    inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
+    inner = {"method": "fp8", "activation_scheme": "static"}
    original = inner.copy()
    build_quant_config({"transformer": inner, "vae": None})
    assert inner == original
@@ -183,18 +115,20 @@ def test_flat_dict_not_misdetected_as_per_component():


 def test_build_quant_config_passthrough():
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config

-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import build_quant_config
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
    assert build_quant_config(fp8) is fp8


 def test_component_config_routing():
-    from vllm_omni.quantization import ComponentQuantizationConfig
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config

-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
    config = ComponentQuantizationConfig(
        component_configs={"transformer": fp8, "vae": None},
    )
@@ -206,10 +140,11 @@ def test_component_config_routing():


 def test_component_config_with_default():
-    from vllm_omni.quantization import ComponentQuantizationConfig
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config

-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
    config = ComponentQuantizationConfig(
        component_configs={"vae": None},
        default_config=fp8,
@@ -232,7 +167,7 @@ def test_integration_dict():

    config = OmniDiffusionConfig(
        model="test",
-        quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
+        quantization_config={"method": "fp8", "activation_scheme": "static"},
    )
    assert config.quantization_config is not None
    assert config.quantization_config.get_name() == "fp8"
@@ -1,72 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Exercise online FP8 loading and inference through the native CUDA path."""
-
-import pytest
-import torch
-from torch.nn import functional as F
-from vllm.config import set_current_vllm_config
-from vllm.distributed import (
-    destroy_distributed_environment,
-    destroy_model_parallel,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
-from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
-from vllm.utils.network_utils import get_open_port
-from vllm.utils.torch_utils import set_default_torch_dtype
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
-from vllm_omni.quantization import build_quant_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
-    ignored_key = "ignored_layers" if method == "fp8" else "ignore"
-    config = build_quant_config(method, **{ignored_key: ["skip"]})
-    od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
-    vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
-
-    try:
-        init_distributed_environment(
-            world_size=1,
-            rank=0,
-            local_rank=0,
-            distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
-            backend="nccl",
-        )
-        with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
-            initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
-            layer = ReplicatedLinear(
-                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
-            )
-            assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
-            assert layer.weight.is_meta
-
-            generator = torch.Generator(device="cuda").manual_seed(42)
-            weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
-            inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
-            expected = F.linear(inputs, weights).float()
-
-            # This is the loader installed by the native online quantizer.
-            # It materializes and quantizes the BF16 checkpoint in one pass.
-            layer.weight.weight_loader(layer.weight, weights)
-            assert layer.weight.dtype == torch.float8_e4m3fn
-            actual = layer(inputs)[0].float()
-            relative_mse = (actual - expected).square().mean() / expected.square().mean()
-            assert relative_mse.item() < 0.002
-
-            ignored = ReplicatedLinear(
-                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
-            )
-            assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
-            assert ignored.weight.dtype == torch.bfloat16
-    finally:
-        destroy_model_parallel()
-        destroy_distributed_environment()
@@ -77,7 +77,6 @@ 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(),
        )
@@ -128,9 +127,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")
-    metadata = envelope.get("metadata")
-    return payload if isinstance(payload, dict) else {}, metadata if isinstance(metadata, 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


 class TestPipelineStreamingOutputToStageDiffusionClient:
@@ -231,7 +230,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
        client._response_socket.bind(response_address)
        client._encoder = OmniMsgpackEncoder()
        client._decoder = OmniMsgpackDecoder()
-        client._output_queue = asyncio.Queue[OmniRequestOutput]()
+        client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
        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-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 from types import SimpleNamespace

@@ -14,7 +14,6 @@ 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):
@@ -39,7 +38,6 @@ 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):
@@ -62,10 +60,8 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
    assert calls == ["impl"]


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

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

-    assert calls == (["boundary"] if initialized_with_hsdp else ["impl"])
+    assert calls == ["impl"]
@@ -11,12 +11,10 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 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
@@ -97,30 +95,6 @@ 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
@@ -636,45 +610,6 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
    return _build


-@pytest.mark.parametrize("active_stream_window", [0, 2])
-def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window):
-    adapter, _ = build_adapter(active_stream_window=active_stream_window)
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    scheduler.chunk_transfer_adapter = adapter
-    scheduler.input_coordinator = None
-    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    holder.num_computed_tokens = 2
-    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
-    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
-    scheduler._enqueue_waiting_request(fresh)
-    scheduler._enqueue_waiting_request(holder)
-
-    scheduler._process_pending_omni_inputs(model_mode="ar")
-
-    assert not scheduler.waiting and not scheduler.kv_holding_waiting
-    assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh]
-    assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-    adapter._finished_load_reqs.add(holder.request_id)
-    scheduler._process_pending_omni_inputs(model_mode="ar")
-
-    assert holder.status == RequestStatus.WAITING
-    assert fresh.status == RequestStatus.WAITING_FOR_CHUNK
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert not scheduler.waiting
-    assert adapter.requests_with_ready_chunks == {holder.request_id}
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-
 def _dequeue_load_entry(adapter, request):
    """Register and take the entry as the background recv loop would."""
    adapter.load_async(request)
@@ -2485,8 +2420,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
    scheduler.requests = {request.request_id: request}
    scheduler.running = [request] if placement == "running" else []
    scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
-    scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
    scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)

    freed = []
@@ -2510,7 +2444,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
    assert not adapter.waiting_for_chunk_running_requests
    assert scheduler.running == []
    assert list(scheduler.waiting) == []
-    assert list(scheduler.kv_holding_waiting) == []
+    assert list(scheduler.skipped_waiting) == []
    assert scheduler.num_waiting_for_streaming_input == 0

    fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
@@ -2533,8 +2467,7 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
    scheduler.requests = {request.request_id: request}
    scheduler.running = []
    scheduler.waiting = DummyWaitingQueue()
-    scheduler.kv_holding_waiting = DummyWaitingQueue()
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = DummyWaitingQueue()
    scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")

    assert (
@@ -2608,7 +2541,6 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc

    scheduler = mocker.MagicMock()
    scheduler._first_chunk_express = False
-    scheduler.aux_output_connector = None
    scheduler._express_min_slack_s = 0.0
    scheduler.chunk_transfer_adapter = adapter_mock
    scheduler.connector = None
@@ -2688,10 +2620,10 @@ 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.aux_output_connector = None
    scheduler.chunk_transfer_adapter = adapter_mock
-    bind_omits_transfer_helpers(scheduler)
    scheduler.connector = None
    scheduler.perf_metrics = None
    scheduler.log_stats = False
@@ -2788,8 +2720,6 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
    sched.requests = {}
    sched.running = []
    sched.waiting = []
-    sched.kv_holding_waiting = []
-    sched.deferred_waiting = set()

    with patch.object(VLLMScheduler, "finish_requests", _super_finish):
        OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -2805,8 +2735,7 @@ def _parked_sender_scheduler(adapter, session):
    scheduler.requests = {session.request_id: session}
    scheduler.running = []
    scheduler.waiting = DummyWaitingQueue()
-    scheduler.kv_holding_waiting = DummyWaitingQueue([session])
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = DummyWaitingQueue([session])
    scheduler.num_waiting_for_streaming_input = 1
    scheduler.finished_req_ids_dict = {}

@@ -213,7 +213,7 @@ class TestCleanup:
        connector.put("s0", "s1", "cleanup_req_42", data)
        assert "cleanup_req_42" in connector._pending_keys

-        assert connector.cleanup("cleanup_req_42") is True
+        connector.cleanup("cleanup_req_42")
        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)

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

    def test_close_cleans_all_pending(self, connector):
@@ -296,24 +296,6 @@ 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 ───────────────────────────────────────────────────


@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-
-
-@pytest.fixture
-def sana_model_parallel(tmp_path):
-    """Construct native parallel layers with real single-rank GPU groups."""
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_distributed_env,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
-    )
-    try:
-        initialize_model_parallel()
-        yield
-    finally:
-        destroy_distributed_env()
Author	SHA1	Message	Date
Hongsheng Liu	d29a35bb65	fix: preserve direct execution of the review verifier	2026-10-01 17:21:00 -07:00
Hongsheng Liu	f2441ddd2a	fix: align migrated review helpers with repository lint policy	2026-10-01 17:18:54 -07:00
Hongsheng Liu	9454e7987c	docs: distinguish generic review selection from session eligibility	2026-10-01 17:11:12 -07:00
Hongsheng Liu	2c43541fa7	docs: consolidate repository-owned review helpers and backends	2026-10-01 17:08:01 -07:00