Merge branch 'main' into dev/vllm-align

2026-09-10 11:23:08 +08:00
commit 0ae9bdd8c2
@@ -17,6 +17,8 @@
 #   a2b3_npu_1  — ascend-a2b3, resource_class npu-1 (A2/B3 CI image)
 #   a2b3_npu_4  — ascend-a2b3, resource_class npu-4
 #   a3_npu_2    — ascend-a3, resource_class npu-2 (A3 CI image tag …:$BUILDKITE_COMMIT-a3)
+#   a5_npu_2    — ascend-a5, resource_class npu-2
+#   310p_npu_1  — ascend-310p, resource_class npu-1
 #
 # B200 (b200-k8s): matches vLLM device ``b200-k8s`` from ci-infra
 # (AgentQueue.B200_K8S + b200_plugin_template). Queue routes the K8s pool;
@@ -649,3 +651,52 @@ mirror_hardwares:
      resource_class: "npu-16"
    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
    plugins: *npu_k8s_plugins
+
+  a5_npu_2:
+    agents:
+      queue: "ascend-a5"
+      resource_class: "npu-2"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5"
+    plugins: *npu_k8s_plugins
+
+  a5_npu_4:
+    agents:
+      queue: "ascend-a5"
+      resource_class: "npu-4"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5"
+    plugins: *npu_k8s_plugins
+
+  a5_npu_8:
+    agents:
+      queue: "ascend-a5"
+      resource_class: "npu-8"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5"
+    plugins: *npu_k8s_plugins
+
+  a5_npu_16:
+    agents:
+      queue: "ascend-a5"
+      resource_class: "npu-16"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5"
+    plugins: *npu_k8s_plugins
+
+  310p_npu_1:
+    agents:
+      queue: "ascend-310p"
+      resource_class: "npu-1"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p"
+    plugins: *npu_k8s_plugins
+
+  310p_npu_2:
+    agents:
+      queue: "ascend-310p"
+      resource_class: "npu-2"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p"
+    plugins: *npu_k8s_plugins
+
+  310p_npu_4:
+    agents:
+      queue: "ascend-310p"
+      resource_class: "npu-4"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p"
+    plugins: *npu_k8s_plugins
@@ -70,7 +70,15 @@ from tests.helpers.mark import (  # noqa: E402

 LOG = "upload_pipeline"
 BOOTSTRAP_STEPS_FILENAME = "bootstrap-upload-steps.yml"
-BOOTSTRAP_IMAGE_BUILD_KEYS = frozenset({"image-build", "image-build-a2", "image-build-a3"})
+BOOTSTRAP_IMAGE_BUILD_KEYS = frozenset(
+    {
+        "image-build",
+        "image-build-a2",
+        "image-build-a3",
+        "image-build-a5",
+        "image-build-310p",
+    }
+)
 BOOTSTRAP_UPLOAD_IF_KEYS = {
    "upload-ready-pipeline": "ready",
    "upload-merge-pipeline": "merge",
@@ -93,7 +93,7 @@ steps:
    commands:
      - |
        timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
+                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
        "
    mirror_hardwares: l4_1

@@ -3,10 +3,10 @@
 steps:
  - label: ":buildkit: Build and Push NPU Test Image - A2"
    key: image-build-a2
-    agents:
+    agents: &npu-image-build-agents
      queue: "ascend-a2b3"
      resource_class: "npu-1"
-    plugins:
+    plugins: &npu-image-build-plugins
      - kubernetes:
          metadata:
            annotations:
@@ -46,123 +46,47 @@ steps:
      IMAGE_NAME: "vllm-omni-ci-npu"
      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
-
-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu.ci \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+    command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}"

  - label: ":buildkit: Build and Push NPU Test Image - A3"
    key: image-build-a3
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    plugins:
-      - kubernetes:
-          metadata:
-            annotations:
-              vault.hashicorp.com/agent-init-first: "true"
-              vault.hashicorp.com/agent-inject: "true"
-              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-              vault.hashicorp.com/agent-pre-populate-only: "true"
-              vault.hashicorp.com/agent-run-as-group: "1000"
-              vault.hashicorp.com/agent-run-as-user: "1000"
-              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-              vault.hashicorp.com/role: ascend-gha-runners
-              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-              vault.hashicorp.com/tls-skip-verify: "true"
-          podSpecPatch:
-            volumes:
-            - name: token-vol
-              projected:
-                defaultMode: 420
-                sources:
-                - serviceAccountToken:
-                    audience: api
-                    expirationSeconds: 600
-                    path: token
+    agents: *npu-image-build-agents
+    plugins: *npu-image-build-plugins
    env:
      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a3"
      IMAGE_NAME: "vllm-omni-ci-npu"
      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
+    command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -a3

-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
+  - label: ":buildkit: Build and Push NPU Test Image - A5"
+    key: image-build-a5
+    agents: *npu-image-build-agents
+    plugins: *npu-image-build-plugins
+    env:
+      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a5"
+      IMAGE_NAME: "vllm-omni-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -a5

-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu.ci.a3 \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+  - label: ":buildkit: Build and Push NPU Test Image - 310P"
+    key: image-build-310p
+    agents: *npu-image-build-agents
+    plugins: *npu-image-build-plugins
+    env:
+      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-310p"
+      IMAGE_NAME: "vllm-omni-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -310p

  - label: "Upload Ready Pipeline"
    depends_on:
      - image-build-a2
      - image-build-a3
+      - image-build-a5
+      - image-build-310p
    key: upload-ready-pipeline
    agents:
      queue: "cpu_queue_premerge"
@@ -174,6 +98,8 @@ steps:
    depends_on:
      - image-build-a2
      - image-build-a3
+      - image-build-a5
+      - image-build-310p
    key: upload-nightly-pipeline
    commands:
      - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-nightly.yml
@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+set -ex
+
+dockerfile=${1:?Dockerfile is required}
+image_tag=${2:?Image tag is required}
+device_suffix=${3:-}
+
+image_name=${IMAGE_NAME:-vllm-omni-ci-npu}
+image_registry=${IMAGE_REGISTRY:-swr.cn-southwest-2.myhuaweicloud.com/modelfoundry}
+buildkitd_addr=${BUILDKITD_ADDR:-tcp://buildkitd-service.buildkitd:1234}
+
+echo "--- Building and pushing NPU Test Image"
+echo "Image: ${image_registry}/${image_name}:${image_tag}"
+echo "Dockerfile: ${dockerfile}"
+
+if ! command -v buildctl &> /dev/null; then
+  mkdir -p /tmp/buildkit
+  buildkit_version=v0.29.0
+  wget -q \
+    "https://github.com/moby/buildkit/releases/download/${buildkit_version}/buildkit-${buildkit_version}.linux-arm64.tar.gz" \
+    -O /tmp/buildkit.tar.gz
+  tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+  cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+fi
+
+build_args=()
+if [[ -n "${device_suffix}" ]]; then
+  build_args+=(--opt "build-arg:VLLM_ASCEND_DEVICE_SUFFIX=${device_suffix}")
+fi
+
+export DOCKER_CONFIG=/home/user/.docker
+buildctl \
+  --addr="${buildkitd_addr}" \
+  --tlscacert=/home/user/.docker/ca.pem \
+  --tlscert=/home/user/.docker/cert.pem \
+  --tlskey=/home/user/.docker/key.pem \
+  build \
+  --frontend dockerfile.v0 \
+  --local context=. \
+  --local dockerfile=./docker \
+  --opt "filename=${dockerfile}" \
+  "${build_args[@]}" \
+  --secret id=dockerconfig,src=/home/user/.docker/config.json \
+  --output "type=image,name=${image_registry}/${image_name}:${image_tag},push=true" \
+  --progress=plain
+
+echo "--- Image pushed successfully"
+echo "${image_registry}/${image_name}:${image_tag}"
@@ -1,28 +1,59 @@
 steps:
-  - label: "Platforms Test"
-    key: ready-platforms-test
+  - group: ":card_index_dividers: Simple Test"
    depends_on: upload-ready-pipeline
-    mirror_hardwares: a3_npu_2
-    commands:
-      - >-
-        pytest -s -v
-        tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
-        -m 'core_model and npu and A3 and omni'
-        --run-level 'core_model'
+    steps:
+      - label: "Platforms Test"
+        key: ready-platforms-test
+        mirror_hardwares: a3_npu_2
+        commands:
+          - >-
+            pytest -s -v
+            tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
+            -m 'core_model and npu and A3 and omni'
+            --run-level 'core_model'

-  - label: "Diffusion NPU layer tests"
-    key: ready-diffusion-npu-layer-tests
-    depends_on: upload-ready-pipeline
-    mirror_hardwares: a3_npu_2
-    commands:
-      - >-
-        pytest -s -v
-        tests/diffusion/layers/test_activation.py
-        tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
-        tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
-        tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py
-        -m 'core_model and npu and A3'
-        --run-level 'core_model'
+      - label: "Diffusion NPU layer tests - A2"
+        key: ready-diffusion-npu-layer-tests-a2
+        mirror_hardwares: a2b3_npu_1
+        commands:
+          - >-
+            pytest -s -v
+            tests/diffusion/layers/test_activation.py
+            -m 'core_model and npu and A2'
+            --run-level 'core_model'
+
+      - label: "Diffusion NPU layer tests - A3"
+        key: ready-diffusion-npu-layer-tests
+        mirror_hardwares: a3_npu_2
+        commands:
+          - >-
+            pytest -s -v
+            tests/diffusion/layers/test_activation.py
+            tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
+            tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
+            tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py
+            -m 'core_model and npu and A3'
+            --run-level 'core_model'
+
+      - label: "Diffusion NPU layer tests - A5"
+        key: ready-diffusion-npu-layer-tests-a5
+        mirror_hardwares: a5_npu_2
+        commands:
+          - >-
+            pytest -s -v
+            tests/diffusion/layers/test_activation.py
+            -m 'core_model and npu and A5'
+            --run-level 'core_model'
+
+      - label: "Diffusion NPU layer tests - 310P"
+        key: ready-diffusion-npu-layer-tests-310p
+        mirror_hardwares: 310p_npu_1
+        commands:
+          - >-
+            pytest -s -v
+            tests/diffusion/layers/test_activation.py
+            -m 'core_model and npu and 310P'
+            --run-level 'core_model'

  - group: ":card_index_dividers: Omni Model Test"
    key: ready-omni-group
@@ -123,10 +123,11 @@ it changes again, report the churn and wait for a stable target.

 For a trusted PR head, materialize the pinned head in an isolated detached
 worktree. A worktree freezes identity but is not a security sandbox. Treat fork
-heads as untrusted unless the user and environment policy explicitly establish
-otherwise: execute them only in a disposable, secret-free sandbox with restricted
-filesystem, network, and resources; without one, use static SHA-addressed reads
-and CI evidence only. For a local review, freeze the committed, index, worktree,
+heads as untrusted until the user and environment policy explicitly establish
+trust: execute their code only after that trust is recorded for this host,
+with the trusted SHA noted in the review state and secrets kept out of the
+execution scope; without recorded trust, use static SHA-addressed reads and CI
+evidence only. For a local review, freeze the committed, index, worktree,
 and NUL-safe in-scope untracked contents. Follow
 [review-execution.md](references/process/review-execution.md) for trust gates,
 state fingerprints, and byte-for-byte staleness checks.
@@ -196,7 +197,7 @@ style preferences as simplification findings.

 Before each validation group, verify the frozen SHA plus the tracked, index,
 untracked, and ignored-file fingerprint, or recreate a pristine snapshot. On a
-trusted head or inside the required sandbox, run an import/version preflight,
+head the user explicitly trusted for this host, run an import/version preflight,
 then the narrowest relevant tests and low-cost static checks. Bind every result
 to the head SHA, snapshot fingerprint, and environment fingerprint. Never run
 imports, tests, builds, hooks, or repo-configurable tooling from an untrusted
@@ -22,6 +22,19 @@ Keep exact commands and report variability rather than a cherry-picked run.
 Use the repository's benchmark when it covers the claim; otherwise use the
 smallest reproducible workload that reaches the changed production path.

+## Isolation per claim
+
+Isolate one claim per comparison. When a PR stacks several changes (a fused or
+restructured load path, a distilled step schedule, a new cache), request one
+column per claim at equal steps/workload instead of one blended base-vs-head
+delta, and name the unchanged stage or metric that pins equivalence.
+
+When the changed path is runtime-switchable (feature flag or environment
+variable), run every column on the same head and toggle the switch rather than
+diffing against a base commit that cannot separate the stacked changes:
+base, head with the switch off, head with the switch on, then head with the
+switch on plus each further claim.
+
 ## Evidence table

 | Dimension | Typical evidence |
@@ -34,6 +47,100 @@ smallest reproducible workload that reaches the changed production path.
 Every optimization needs correctness or quality evidence. A faster result on a
 different workload, precision, seed, or topology is not a valid comparison.

+## E2e stage-attribution table
+
+For an end-to-end latency claim, request the stage-attribution table the
+repository's perf PRs use:
+
+- Rows carry the benchmark's own field names: total e2e, per-stage execution,
+  postprocess, response encoding, and a request-to-output roll-up. Add a stage
+  handoff/idle row when the pipeline spans stages — connector transfer and
+  scheduling gaps are invisible in per-stage latency deltas.
+- Sanity rows pin comparability: inference steps, output frames/pixels/samples,
+  and generated tokens held identical across columns.
+- Report P50 and P100 from the stated warmup and measured counts. One
+  population per number: never mix cold-start and warm samples into one value.
+- Name the baseline and candidate commit SHAs, and close with an isolation
+  statement naming the metric that stayed unchanged ("X remained effectively
+  unchanged, isolating the gain to Y").
+
+When the serving path caches, captures, or regenerates at runtime (CUDA-graph
+bucket growth, cache rebuilds, recompilation), medians hide the resulting tail.
+Also request the event count and the maximum observed stall.
+
+Component-level tables pair with, not replace, the e2e table: per-op reference
+vs candidate timings (with the numerical-equality result per op) plus one
+complete-pipeline check under the production shapes.
+
+## Proportionality
+
+When a PR adds a model and explicitly labels its numbers as smoke observations
+rather than claims, do not demand the attribution machinery. Ask only for
+evidence protecting the model's stated operating contract — for a
+realtime/duplex model that is sustained cadence over a realistic session
+(enough ticks to leave warmup), with any mid-session capture/regeneration
+event and its stall reported, since a cadence that collapses mid-session is a
+functional failure, not a perf claim.
+
+## Realtime/duplex world models
+
+Request the operating-contract evidence as serving metrics, not claims:
+
+- Sustained per-tick latency over one realistic session: median, max, and the
+  argmax tick.
+- Stall events with their causes.
+- Peak memory at first versus last tick — session-owned state must not grow
+  per tick.
+- Control latency from a submitted action to the tick that applies it.
+- Realtime-versus-offline parity under the same controls, using the
+  repository's video-similarity helpers. Parity plotted over tick index is the
+  drift-over-horizon curve and the honest measure of session lifetime;
+  distribution-level suites (FVD, VBench) belong to model evaluation, not the
+  serving integration.
+
+Watch the contract boundaries the model documents: if a control input is
+tick-API-only, the parity pair must use a horizon both paths can run.
+
+## Speech-generation (TTS) models
+
+The streaming contract outranks the average — a synthesis that pauses
+mid-stream is broken for realtime use regardless of its averages. Request:
+
+- The benchmark's own speech metrics over one long utterance (enough audio to
+  leave warmup, tens of seconds): `audio_ttfp` (time to first packet),
+  `audio_rtf` (real-time factor; above 1.0 cannot keep up), `e2el`, and
+  `ttft`/`tpot` where tokens precede audio.
+- Mid-utterance stalls as the benchmark's continuity fields, not prose:
+  `audio_underrun` percentiles and the `audio_continuity_ok` rate
+  (`vllm_omni/benchmarks/metrics/metrics.py`), computed by
+  `vllm_omni/benchmarks/audio_continuity.py` — the companion to `audio_ttfp`
+  and `audio_rtf`. A synthesis that bursts chunks fast enough on average but
+  underruns mid-stream fails the streaming contract even when the averages
+  pass.
+- Quality as a paired comparison against the reference implementation under
+  the same text, speaker prompt, and seed; report sample rate and frame
+  count, and assert output completeness — truncated or dropped trailing audio
+  is a known failure class, not a quality nuance.
+- For realtime or duplex TTS: interruption (barge-in) latency and the
+  per-session memory trend across turns.
+
+## Omni models
+
+Apply the e2e stage-attribution table directly to the voice-to-voice path of
+omni (speech-to-speech and multimodal AR) models:
+
+- TTFT for text output, `audio_ttfp` for audio output, and per-stage
+  ITL/TPOT cadence across the perception-understanding-generation chain.
+- Cross-modal parity — the omni analogue of realtime-versus-offline parity:
+  the same semantic request through text and audio inputs, compared on output
+  content or timing.
+- For duplex omni sessions: turn-taking latency, interruption handling, the
+  concurrent session count actually validated, and a per-turn memory trend —
+  session-owned KV state must not grow per turn. When either turn produces
+  audio, continuity is judged with the same `audio_underrun` and
+  `audio_continuity_ok` fields as TTS, since turn-taking latency alone does
+  not expose mid-utterance underruns.
+
 ## Reviewer verification ladder

 1. Run base/head A/B on suitable hardware when affordable.
@@ -121,6 +121,19 @@ Search bounded callers and sibling implementations before reporting dead code,
 duplication, or a missed consumer. Keep new abstractions only when they have a
 distinct live caller, invariant, or compatibility need.

+Verify absence and equality claims with whole-file evidence. A line window, a
+grep for one symbol, or a commit message proves presence, never absence: to
+assert that a file matches another version or lacks a change, diff or hash the
+full files (for example `git diff <sha1>:<path> <sha2>:<path>` or the contents
+API's size and SHA fields) before anchoring the finding. A verification
+shortcut that would itself be wrong to run is not evidence.
+
+For a value written or finalized in more than one place, establish which site
+executes first on each entry path — and whether the earlier site makes the
+later ones unreachable — before naming where a fix belongs. Enumerating all
+call sites correctly still misses the bug when a path the enumeration ranked
+last actually finalizes first and clears the state the proposed fix checks.
+
 ## Finding bar

 Anchor each finding to a changed `path:line`; name the trigger or call path,
@@ -39,13 +39,13 @@ isolated detached worktree. Run every source read, `rg` search, import, and test
 from that snapshot, not the caller's checkout. A detached worktree provides
 snapshot isolation, not a security boundary.

-Treat a fork head as untrusted unless the user and environment policy explicitly
-establish trust. Never run its imports, tests, builds, hooks, package setup, or
-repo-configurable linters/plugins on the reviewer host. Execute them only in a
-disposable sandbox or VM with no credentials or inherited secrets, no host agent
-or service sockets, minimal read-only host mounts, disabled or explicitly
-allowlisted network, resource/time limits, and destruction after the run. If
-that boundary is unavailable, limit the review to the remote diff, SHA-addressed
+Treat a fork head as untrusted until the user and environment policy
+explicitly establish trust. Without recorded trust, never run its imports,
+tests, builds, hooks, package setup, or repo-configurable linters/plugins on
+the reviewer host. Once the user explicitly trusts this pinned commit for
+this host, record the trusted SHA with the review state and execute with
+credentials, agent sockets, and other secrets out of scope. Without recorded
+trust, limit the review to the remote diff, SHA-addressed
 reads such as `git show <head_sha>:<path>`, and existing CI evidence; report all
 executable validation as a gap.

@@ -1,20 +1,7 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.28.0
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
-
-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+ARG VLLM_ASCEND_VERSION=v0.28.0
+ARG VLLM_ASCEND_DEVICE_SUFFIX=
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -24,6 +11,7 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v "mindiesd==3.1.0" && \
    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation

 RUN python3 -m pip install -v --no-deps \
@@ -1,36 +0,0 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.28.0-a3
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
-
-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
-
-ARG APP_DIR=/vllm-workspace/vllm-omni
-WORKDIR ${APP_DIR}
-COPY . .
-
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
-
-RUN python3 -m pip install -v --no-deps \
-      "stepaudio2-minicpmo>=0.1.1" \
-      "step-audio2>=1.0.0"
-
-ENV VLLM_TARGET_DEVICE=npu
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-ENTRYPOINT []
@@ -1,10 +1,10 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.28.0
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+ARG VLLM_ASCEND_VERSION=v0.28.0
+ARG VLLM_ASCEND_DEVICE_SUFFIX=
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
-COPY . .

 RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list

@@ -22,11 +22,9 @@ RUN apt-get update && \
    rm -rf /var/cache/apt/* && \
    rm -rf /var/lib/apt/lists/*

-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
+RUN python3 -m pip install -U pip && \
    python3 -m pip install -v \
+      "mindiesd==3.1.0" \
      "pytest>=7.0.0" \
      "pytest-asyncio>=0.21.0" \
      "pytest-cov>=4.0.0" \
@@ -48,7 +46,13 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
      "datasets>=3.0.0" \
      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60 && \
+      --default-timeout=60
+
+COPY . .
+
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
@@ -1,72 +0,0 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.28.0-a3
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-ARG APP_DIR=/vllm-workspace/vllm-omni
-WORKDIR ${APP_DIR}
-COPY . .
-
-RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list
-
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends \
-        ffmpeg \
-        libgl1 \
-        libglib2.0-0 \
-        libgomp1 \
-        libsm6 \
-        libxext6 \
-        libxrender1 \
-        libxcb1 \
-        espeak-ng && \
-    rm -rf /var/cache/apt/* && \
-    rm -rf /var/lib/apt/lists/*
-
-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v \
-      "pytest>=7.0.0" \
-      "pytest-asyncio>=0.21.0" \
-      "pytest-cov>=4.0.0" \
-      "pytest-mock>=3.10.0" \
-      "openpyxl>=3.0.0" \
-      "psutil>=7.2.0" \
-      "soundfile>=0.13.1" \
-      "pyttsx3>=2.99" \
-      "opencc==1.4.1" \
-      "opencv-python>=4.12.0.88" \
-      "imageio[ffmpeg]>=0.6.0" \
-      "jiwer>=3.0.0" \
-      "zhon>=2.0.0" \
-      "zhconv>=1.4.2" \
-      "funasr>=1.0.0" \
-      "hyperpyyaml>=1.2.0" \
-      "onnx>=1.16.0" \
-      "librosa>=0.10.0" \
-      "torchmetrics>=1.9.0" \
-      "datasets>=3.0.0" \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60 && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
-
-# MiniCPM-o 4.5 Stage-2 vocoder assets: s3tokenizer plus flashcosyvoice, which
-# is shipped by step-audio2. --no-deps keeps pip from replacing the Ascend torch
-# build and from shadowing onnxruntime-cann with the vanilla onnxruntime wheel;
-# the deps these two actually need are installed above.
-RUN python3 -m pip install -v --no-deps \
-      "s3tokenizer>=0.3.0" \
-      "step-audio2>=1.0.0" \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
-
-ENV VLLM_TARGET_DEVICE=npu
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-ENTRYPOINT []
@@ -129,7 +129,7 @@ Canonical layout (prefer these paths for new changes):

 === "NPU"

-    **Bootstrap:** [`npu/pipeline-npu.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/pipeline-npu.yml) + [`npu/bootstrap-upload-steps.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/bootstrap-upload-steps.yml)—same split as CUDA; builds A2/B3 and A3 CI images, then uploads child test pipelines.
+    **Bootstrap:** [`npu/pipeline-npu.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/pipeline-npu.yml) + [`npu/bootstrap-upload-steps.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/bootstrap-upload-steps.yml)—same split as CUDA; builds A2/B3, A3, A5, and 310P CI images, then uploads child test pipelines.

    **Test YAML:** `npu/test-npu-ready.yml` (L2), `test-npu-nightly.yml` (L4).

@@ -137,7 +137,12 @@ Canonical layout (prefer these paths for new changes):

    **Upload:** same as CUDA—`upload_pipeline.py --upload`.

-    **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. Presets: `a2b3_npu_1`, `a2b3_npu_4`, `a3_npu_2` in `common/ci_mirror_hardwares.yml`.
+    **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. NPU presets in `common/ci_mirror_hardwares.yml` are:
+
+    - A2: `a2b3_npu_1`, `a2b3_npu_4`, `a2b3_npu_8`
+    - A3: `a3_npu_2`, `a3_npu_4`, `a3_npu_8`, `a3_npu_16`
+    - A5: `a5_npu_2`, `a5_npu_4`, `a5_npu_8`, `a5_npu_16`
+    - 310P: `310p_npu_1`, `310p_npu_2`, `310p_npu_4`

    **Conventions**

@@ -48,6 +48,8 @@ Defined in `pyproject.toml`:
 | `MI325`            | Tests that require MI325 GPU (AMD/ROCm) *                                                               |
 | `A2`               | Tests that require A2 NPU *                                                                             |
 | `A3`               | Tests that require A3 NPU *                                                                             |
+| `A5`               | Tests that require A5 NPU *                                                                             |
+| `310P`             | Tests that require 310P NPU *                                                                           |
 | `cards_{n}`        | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` *         |
 | `slow`             | Slow tests (may skip in quick CI)                                                                       |
 | `benchmark`        | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
@@ -76,7 +78,7 @@ def test_video_to_audio()
 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:

 1. **Applies platform and resource markers**  
-   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`).
+   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`, `A5`, `310P`).
   ```python
   @pytest.mark.cuda
   @pytest.mark.L4
@@ -7,13 +7,34 @@ For detailed hardware and software requirements, please refer to the [vLLM-Ascen

 The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built Docker images:

+=== "A2"
+
+    ```bash
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0
+    ```
+
+=== "A3"
+
+    ```bash
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a3
+    ```
+
+=== "A5"
+
+    ```bash
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a5
+    ```
+
+=== "310P (Experimental)"
+
+    ```bash
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-310p
+    ```
+
 ```bash
-# vLLM-Ascend image
-# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3
 docker run --rm \
    --name vllm-omni-npu \
-    --shm-size=1g \
+    --shm-size=64g \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
@@ -27,8 +48,8 @@ docker run --rm \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /root/.cache:/root/.cache \
-    -p 8000:8000 \
-    -it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash
+    -p 8091:8091 \
+    -it "$IMAGE" bash

 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
@@ -72,53 +93,146 @@ Supported images as following.

 | image name | Hardware | OS |
 |-|-|-|
-| image-tag | Atlas A2 | Ubuntu |
-| image-tag-a3 | Atlas A3 | Ubuntu |
+| image-tag | A2 | Ubuntu |
+| image-tag-a3 | A3 | Ubuntu |
+| image-tag-a5 | A5 | Ubuntu |
+| image-tag-310p | 310P (experimental) | Ubuntu |

 Here's an example deployment command that has been verified on 4 x NPUs:

-```bash
-# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3
-docker run --rm \
-    --name vllm-omni-a3 \
-    --shm-size=4g \
-    --device /dev/davinci0 \
-    --device /dev/davinci1 \
-    --device /dev/davinci2 \
-    --device /dev/davinci3 \
-    --device /dev/davinci_manager \
-    --device /dev/devmm_svm \
-    --device /dev/hisi_hdc \
-    -v /usr/local/dcmi:/usr/local/dcmi \
-    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-    --device /dev/davinci_manager \
-    --device /dev/devmm_svm \
-    --device /dev/hisi_hdc \
-    -v /usr/local/dcmi:/usr/local/dcmi \
-    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-    -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -v ~/.cache:/root/.cache \
-    -p 8091:8091 \
-    -it quay.io/ascend/vllm-omni:v0.28.0-a3 bash
-```
+=== "A2"
+
+    ```bash
+    export IMAGE=quay.io/ascend/vllm-omni:v0.28.0
+    docker run --rm \
+        --name vllm-omni-a2 \
+        --shm-size=64g \
+        --device /dev/davinci0 \
+        --device /dev/davinci1 \
+        --device /dev/davinci2 \
+        --device /dev/davinci3 \
+        --device /dev/davinci_manager \
+        --device /dev/devmm_svm \
+        --device /dev/hisi_hdc \
+        -v /usr/local/dcmi:/usr/local/dcmi \
+        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+        -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+        -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+        -v /etc/ascend_install.info:/etc/ascend_install.info \
+        -v ~/.cache:/root/.cache \
+        -p 8091:8091 \
+        -it "$IMAGE" bash
+    ```
+
+=== "A3"
+
+    ```bash
+    export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a3
+    docker run --rm \
+        --name vllm-omni-a3 \
+        --shm-size=64g \
+        --device /dev/davinci0 \
+        --device /dev/davinci1 \
+        --device /dev/davinci2 \
+        --device /dev/davinci3 \
+        --device /dev/davinci_manager \
+        --device /dev/devmm_svm \
+        --device /dev/hisi_hdc \
+        -v /usr/local/dcmi:/usr/local/dcmi \
+        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+        -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+        -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+        -v /etc/ascend_install.info:/etc/ascend_install.info \
+        -v ~/.cache:/root/.cache \
+        -p 8091:8091 \
+        -it "$IMAGE" bash
+    ```
+
+=== "A5"
+
+    ```bash
+    export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a5
+    docker run --rm \
+        --name vllm-omni-a5 \
+        --shm-size=64g \
+        --device /dev/davinci0 \
+        --device /dev/davinci1 \
+        --device /dev/davinci2 \
+        --device /dev/davinci3 \
+        --device /dev/davinci_manager \
+        --device /dev/devmm_svm \
+        --device /dev/hisi_hdc \
+        -v /usr/local/dcmi:/usr/local/dcmi \
+        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+        -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+        -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+        -v /etc/ascend_install.info:/etc/ascend_install.info \
+        -v ~/.cache:/root/.cache \
+        -p 8091:8091 \
+        -it "$IMAGE" bash
+    ```
+
+=== "310P (Experimental)"
+
+    ```bash
+    export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-310p
+    docker run --rm \
+        --name vllm-omni-310p \
+        --shm-size=64g \
+        --device /dev/davinci0 \
+        --device /dev/davinci1 \
+        --device /dev/davinci2 \
+        --device /dev/davinci3 \
+        --device /dev/davinci_manager \
+        --device /dev/devmm_svm \
+        --device /dev/hisi_hdc \
+        -v /usr/local/dcmi:/usr/local/dcmi \
+        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+        -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+        -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+        -v /etc/ascend_install.info:/etc/ascend_install.info \
+        -v ~/.cache:/root/.cache \
+        -p 8091:8091 \
+        -it "$IMAGE" bash
+    ```

 !!! tip
    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.

-Or build IMAGE from **source code**:
+Or build the image from **source code**:

 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
-# A2
-# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
-# A3
-docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
 ```

+All hardware variants use `docker/Dockerfile.npu`. Choose the device suffix for your hardware:
+
+=== "A2"
+
+    ```bash
+    docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu .
+    ```
+
+=== "A3"
+
+    ```bash
+    docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \
+        --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-a3 .
+    ```
+
+=== "A5"
+
+    ```bash
+    docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \
+        --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-a5 .
+    ```
+
+=== "310P (Experimental)"
+
+    ```bash
+    docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \
+        --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-310p .
+    ```
+
 # --8<-- [end:pre-built-images]
@@ -199,12 +199,21 @@ curl -X POST http://localhost:8091/v1/videos \
 | `boundary_ratio`      | float  | None    | Boundary split ratio for low/high DiT (Wan2.2)   |
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
+| `num_outputs_per_prompt` | int  | 1       | Number of videos to generate (1-10, MiniMax H3); only the first is returned, see below |
 | `lora`                | object | None    | LoRA configuration                               |
 | `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding |
 | `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x |
 | `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs |
 | `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` |

+!!! note "Only the first output is returned"
+    `num_outputs_per_prompt` is forwarded to the pipeline, so the model does
+    generate that many videos and you pay the generation cost for all of them.
+    `/v1/videos` and `/v1/videos/sync` then return only the **first** one,
+    because an async video job stores a single file per video id. The server
+    logs a warning when it discards the extra outputs. Keep this at `1`
+    unless you have a specific reason to generate videos you will not receive.
+
 ## Frame Interpolation

 Frame interpolation is an optional post-processing step for `/v1/videos` and
@@ -259,21 +259,25 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.

 ### Prerequisites
+
 None beyond the base install: the neural audio codec is implemented in tree and
 its weights (`codec.pth`) ship with the checkpoint.

 ### Launch
+
 ```bash
 vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
 # or:
 ./audio8_tts/run_server.sh
 ```
+
 The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
 `model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
 registers its own `arktts` config, and transformers would otherwise prefer the
 checkpoint's remote code and bypass it.

 ### Text-only synthesis
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -284,6 +288,7 @@ curl -X POST http://localhost:8092/v1/audio/speech \
 ```

 ### Voice cloning
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -293,18 +298,21 @@ curl -X POST http://localhost:8092/v1/audio/speech \
        "ref_text": "The exact transcript of the reference recording."
    }' --output cloned.wav
 ```
+
 `ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
 mandatory whenever `ref_audio` is present and must match the recording.
 Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
 reference codes across requests (`voice: "<name>"`).

 ### Streaming
+
 ```bash
 python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
 ffplay -f s16le -ar 44100 -ac 1 out.pcm
 ```

 ### Gradio demo
+
 ```bash
 ./audio8_tts/run_gradio_demo.sh          # server + demo
 python audio8_tts/gradio_demo.py --api-base http://localhost:8092   # demo only
@@ -328,7 +336,7 @@ these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
 deploy defaults:

 | | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
-|---|---|---|---|---|
+| --- | --- | --- | --- | --- |
 | RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
 | Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
 | E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
@@ -347,7 +355,7 @@ job holding ~80% SM, no repeats), so read them as directional, not as measured
 speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):

 | stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
 | 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
 | 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
 | 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
@@ -376,6 +384,7 @@ underrun. Raise it to 8 if first-packet latency under load matters more than a
 rare sub-300 ms gap.

 ### Notes
+
 - Output: 44.1 kHz mono; ~21.5 codec frames per second.
 - No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
 - `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
@@ -553,6 +562,46 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.

 ---

+## MOSS-TTS Local Transformer v1.5
+
+For a single H200, the optional
+[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml)
+deployment places both the talker and codec on logical GPU 0:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+    vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+    --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \
+    --disable-log-stats
+```
+
+Run from the repository root and select an available physical GPU with
+`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and
+disables detailed per-request statistics logging to reduce CPU overhead at
+high concurrency. Remove `--disable-log-stats` when those statistics are needed;
+keep these settings identical when comparing performance.
+This preset configures 256 request slots per stage,
+a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens,
+and codec batch buckets through 256. It requires more than 80 GiB of GPU memory;
+request capacity also depends on input and generated lengths. The default
+[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains
+available for smaller deployments.
+
+The preset enables the optional codec backend with
+`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the
+streaming ring-cache mask and uses BF16 attention with 64-dimensional heads;
+other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the
+default attention implementation. Codec terminal tails share execution only
+when they already map to the same padded CUDA Graph; returned audio retains
+each request's actual length.
+
+Voice cloning requests use `ref_audio` and `ref_text`. For streaming output,
+set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`;
+the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute
+audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare
+the same requests, concurrency, and reference-cache state.
+
 ## MOSS-TTS-Nano

 Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets.
@@ -283,6 +283,8 @@ markers = [
    "S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)",
    "A2: [hardware-resource] [npu] Tests that require A2 NPU",
    "A3: [hardware-resource] [npu] Tests that require A3 NPU",
+    "A5: [hardware-resource] [npu] Tests that require A5 NPU",
+    "310P: [hardware-resource] [npu] Tests that require 310P NPU",
    # Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts().
    "cards_1: [hardware-cards] Tests that require 1 accelerator card",
    "cards_2: [hardware-cards] Tests that require 2 accelerator cards",
@@ -7,13 +7,14 @@ recipes it uses **no offload of any kind** — see the capacity note below.
 Validated on:

 - Host: DGX Spark (GB10), aarch64
- GPUs: 1 (unified memory)
- Driver: <FILL_DRIVER_VERSION>
+- GPUs: 1, or 2 hosts with one unified-memory GPU each
+- Driver: 580.173.02 for the two-host qualification
 - vLLM: 0.26.0
 - vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
-  `seed=1101`
-
+- One-host workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps
+- Two-host workload: T2VA, 1344x768, `duration=5.0` and `duration=10.0`,
+  50 steps, text-encoder TP2, and DiT USP2
+- Common sampling settings: `flow_shift=12`, `seed=1101`

 ## Capacity requirements

@@ -24,6 +25,7 @@ Validated on:
 | Checkpoint storage | 135 GiB per partition |
 | Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
 | Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
+| Measured peak per rank, two-host T2VA 5 s / 10 s (FP8) | 75.45 / 87.17 GiB |

 `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
 at a time.
@@ -164,6 +166,41 @@ curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
 The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
 32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.

+## Two GB10 hosts: T2VA at 1344x768
+
+The two-host qualification uses one GB10 GPU per DGX Spark, connected through a
+single ConnectX-7 RoCE interface. The FL2VA checkpoint and software revision are
+identical on both hosts. The distributed configuration is:
+
+- world size 2;
+- text-encoder tensor parallel size 2;
+- DiT Ulysses sequence parallel size 2 and ring degree 1;
+- VAE patch parallel size 1;
+- online FP8, eager execution, cuDNN attention, and tiled VAE decode.
+
+The current diffusion multiprocess executor starts all ranks on one host and sets
+`MASTER_ADDR=localhost`. These measurements therefore used an external `env://`
+launcher with one rank on each host and a benchmark-only adaptation that maps
+both hosts' local worker to `cuda:0`. This qualifies the two-rank MiniMax-H3
+pipeline and its output, but it is not yet a supported cross-host `vllm serve`
+command. Do not infer that the single-host launch command above works unchanged
+across two machines.
+
+One 5-second, 2-step request warmed both workers before the measured runs. The
+reported pipeline elapsed time covers distributed model execution through GPU
+completion; MP4 encoding happened afterward and is not included.
+
+| Duration | Steps | Text encode | Denoise | Per step | VAE decode | Pipeline elapsed | Max peak per rank |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 5 s | 50 | 0.21 s | 1570.19 s | 31.40 s | 79.45 s | 1649.87 s (27 min 30 s) | 75.45 GiB |
+| 10 s | 50 | 0.15 s | 4448.12 s | 88.96 s | 160.71 s | 4609.00 s (76 min 49 s) | 87.17 GiB |
+
+Both ranks reported stage times and allocator peaks within measurement noise.
+The generated outputs were H.264 at 1344x768 and 24 FPS with 32 kHz stereo AAC:
+124 frames (5.167 s) and 243 frames (10.125 s). Both files decoded successfully,
+and neither run encountered an OOM or NCCL error. These are single qualification
+runs (`n=1`), not throughput guarantees.
+
 ## Measured latency

 Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
@@ -238,9 +275,11 @@ shorter than a 50-step Ref2VA run.

 ## Known limitations

- One GPU means no Ulysses, ring, TP, or VAE patch parallelism. Pass
+- With one GPU there is no Ulysses, ring, TP, or VAE patch parallelism. Pass
  `--usp 1 --ring 1 --vae-patch-parallel-size 1` explicitly if a profile or
  script would otherwise supply higher degrees.
+- The two-host results require the external launch adaptation described above;
+  cross-host diffusion worker launch is not exposed by the documented server CLI.
 - `--enforce-eager` is used because regional compile adds startup time without a
  parallelism win on a single rank.
 - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
@@ -255,8 +294,7 @@ shorter than a 50-step Ref2VA run.
  adding references or raising the output shape.
 - Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
  under **Validated on**: FP8 weight loading and image-only Ref2VA are both
-  broken on `release/v0.26.0`.suggestion is  `v0.26.1`.
+  broken on `release/v0.26.0`. The suggested version is `v0.26.1`.
 - Online FP8 is incompatible with layerwise offload — the offload path produces a
  weight stride the Cutlass FP8 kernel rejects. This is not a practical
  restriction here since offload is unusable on GB10 anyway.
-```
@@ -16,7 +16,7 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc
 Chinese and English). Choose a variant based on your latency and quality requirements:

 | Model | Params | Use case |
-|---|---|---|
+| --- | --- | --- |
 | MOSS-TTS | 8B | General TTS, highest quality |
 | MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
 | MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
@@ -24,9 +24,17 @@ Chinese and English). Choose a variant based on your latency and quality require
 | MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
 | MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |

-All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
+The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
 output 24 kHz mono audio.

+MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz
+stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an
+accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the
+reference transcript and audio as a continuation prefix, then generates only
+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.
+
 ## References

 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -55,6 +55,7 @@ recipes/
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
 | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
 | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
+| [`MiniMaxAI/MiniMax-H3-Spark-GB10.md`](./MiniMaxAI/MiniMax-H3-Spark-GB10.md) | MiniMax H3 T2VA and Ref2VA on DGX Spark | 1x or 2x DGX Spark (GB10) |
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
@@ -26,6 +26,7 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.seed_tts_dataset"

 if _MODULE_NAME not in sys.modules:
    _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    assert _spec is not None and _spec.loader is not None
    _mod = importlib.util.module_from_spec(_spec)
    sys.modules[_MODULE_NAME] = _mod
    _spec.loader.exec_module(_mod)
@@ -237,7 +238,8 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
        assert req.seed_tts_utterance_id == "ok"


-def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
+@pytest.mark.parametrize("audio_seconds", [0.1, 61])
+def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds):
    from vllm_omni.benchmarks.data_modules import seed_tts_eval

    calls = {}
@@ -252,8 +254,9 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
            return self

    class FakeProcessor:
-        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
+        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs):
            calls["return_attention_mask"] = return_attention_mask
+            calls["processor_kwargs"] = kwargs
            assert sampling_rate == 16000
            assert return_tensors == "pt"
            assert len(wav) > 0
@@ -262,9 +265,10 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
                attention_mask=FakeTensor("mask") if return_attention_mask else None,
            )

-        def get_decoder_prompt_ids(self, *, language, task):
+        def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True):
            assert language == "english"
            assert task == "transcribe"
+            calls["no_timestamps"] = no_timestamps
            return [(1, 2)]

        def batch_decode(self, predicted_ids, *, skip_special_tokens):
@@ -283,10 +287,18 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
    monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
    monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")

-    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
+    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32))

    assert text == "hello"
    assert calls["return_attention_mask"] is True
    assert calls["input_features"].device == "cuda:1"
    assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
    assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]
+    if audio_seconds > 30:
+        assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"}
+        assert calls["generate_kwargs"]["return_timestamps"] is True
+        assert calls["no_timestamps"] is False
+    else:
+        assert calls["processor_kwargs"] == {}
+        assert "return_timestamps" not in calls["generate_kwargs"]
+        assert calls["no_timestamps"] is True
@@ -5,7 +5,7 @@ import pytest
 import torch
 import torch.nn.functional as F

-from tests.helpers.mark import hardware_test
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.layers.activation import SiluAndMul

 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -19,9 +19,17 @@ def test_silu_and_mul_native_matches_packed_reference() -> None:
    torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)


-@hardware_test(res={"npu": "A3"}, num_cards=1)
-def test_silu_and_mul_npu_matches_packed_reference() -> None:
-    packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A2"}, num_cards=1), id="A2-bf16"),
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A3"}, num_cards=1), id="A3-bf16"),
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A5"}, num_cards=1), id="A5-bf16"),
+        pytest.param(torch.float16, marks=hardware_marks(res={"npu": "310P"}, num_cards=1), id="310P-fp16"),
+    ],
+)
+def test_silu_and_mul_npu_matches_packed_reference(dtype: torch.dtype) -> None:
+    packed = torch.randn(257, 256, device="npu", dtype=dtype)
    gate, up = packed.chunk(2, dim=-1)

    torch.testing.assert_close(
@@ -37,8 +37,10 @@ def clear_platform_cache():
        (90, True, True, True, True),
        (90, True, False, True, False),
        (90, False, True, False, False),
-        (100, True, True, False, False),
-        (103, True, True, False, False),
+        (100, True, True, True, False),
+        (100, True, False, True, False),
+        (103, True, True, True, False),
+        (103, True, False, True, False),
        (110, True, True, False, False),
        (120, True, True, False, False),
        (121, True, True, False, False),
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

-"""CUDA regression tests for LTX-2 BWE vocoder precision."""
+"""GPU regression tests for LTX-2 BWE vocoder precision."""

 import pytest
 import torch
@@ -14,10 +14,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 class _BWEConvVocoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
+        # Match the production LTX vocoder's conv_pre geometry. The previous
+        # single-channel 1x1 probe intermittently had no MIOpen algorithm.
        self.conv = torch.nn.Conv1d(
-            1,
-            1,
-            kernel_size=1,
+            128,
+            1024,
+            kernel_size=7,
+            padding=3,
            bias=False,
            device="cuda",
            dtype=torch.bfloat16,
@@ -33,12 +36,12 @@ class _BWEConvVocoder(torch.nn.Module):
        return output


-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
    from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder

    vocoder = _BWEConvVocoder()
-    generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
+    generated_mel = torch.ones((1, 128, 64), device="cuda", dtype=torch.bfloat16)

    output = _run_ltx_vocoder(vocoder, generated_mel)

@@ -300,7 +300,8 @@ def test_single_stage_zimage_fp8_uses_less_memory():
    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"


-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "xpu": "B60"})
+@pytest.mark.advanced_model
 def test_single_stage_qwen_image_fp8():
    """Qwen-Image (random weights) with FP8 generates valid images."""
    model = "riverclouds/qwen_image_random"
@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Check replica startup ordering without model-loading timing assumptions."""
+
+import contextlib
+import fcntl
+import os
+import threading
+from collections.abc import Iterator
+from pathlib import Path
+
+import pytest
+from pytest_mock import MockerFixture
+from vllm.config import VllmConfig
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.utils import CoreEngineProcManager, EngineZmqAddresses
+
+from vllm_omni.config.stage_config import StageConfig
+from vllm_omni.engine import stage_init_utils, stage_runtime
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan, StageMetadata
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def launch_runtime(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mocker: MockerFixture
+) -> Iterator[stage_runtime.StageRuntime]:
+    runtime = stage_runtime.StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=5,
+        async_chunk=False,
+    )
+    # These are virtual device IDs: only the real per-device file locking and
+    # scheduling run. No model, CUDA context, or engine process is created.
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda _stage, cfg: cfg["devices"])
+    monkeypatch.setattr(
+        stage_init_utils, "device_init_lock_path", lambda device: str(tmp_path / f"device-{device}.lock")
+    )
+    mocker.patch.object(
+        stage_runtime.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        return_value=mocker.Mock(spec=stage_runtime.StageClient),
+    )
+    monkeypatch.setenv("VLLM_OMNI_TEST_LAUNCH_ENV", "parent")
+    yield runtime
+    if runtime._stage_init_executor is not None:
+        runtime._stage_init_executor.shutdown(wait=True)
+
+
+def _plan(stage_id: int, device: str) -> LogicalStageInitPlan:
+    metadata = StageMetadata(
+        stage_id=stage_id,
+        stage_type="llm",
+        engine_output_type="text",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source=[],
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=SamplingParams(),
+        custom_process_input_func=None,
+        model_stage="dummy",
+        runtime_cfg={"devices": device, "env": {"VLLM_OMNI_TEST_LAUNCH_ENV": str(stage_id)}},
+    )
+    replica = ReplicaInitPlan(
+        replica_id=0,
+        num_replicas=1,
+        launch_mode="local",
+        stage_cfg=StageConfig(stage_id=stage_id, model_stage="dummy", yaml_runtime=metadata.runtime_cfg),
+        metadata=metadata,
+        stage_connector_spec={},
+        omni_kv_connector=(None, None, None),
+        stage_vllm_config=VllmConfig(),
+        executor_class=object,
+        engine_args_dict={},
+    )
+    return LogicalStageInitPlan(stage_idx=stage_id, stage_id=stage_id, replicas=[replica])
+
+
+def _resources(mocker: MockerFixture) -> StageReplicaResources:
+    return StageReplicaResources(
+        manager=mocker.Mock(spec=CoreEngineProcManager),
+        addresses=EngineZmqAddresses(inputs=["unused-input"], outputs=["unused-output"]),
+    )
+
+
+def test_distinct_devices_spawn_before_either_is_ready(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    spawned = [threading.Event(), threading.Event()]
+
+    @contextlib.contextmanager
+    def launch(*, stage_id: int, stage_visible_devices: str, **_kwargs: object) -> Iterator[StageReplicaResources]:
+        assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == str(stage_id)
+        spawned[stage_id].set()
+        yield _resources(mocker)
+        # The device remains exclusively locked throughout loading/profiling,
+        # even though another device's replica must be able to start.
+        with open(stage_init_utils.device_init_lock_path(int(stage_visible_devices)), "r+") as lock:
+            with pytest.raises(BlockingIOError):
+                fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        assert spawned[1 - stage_id].wait(5), "A different device's spawn was blocked by the READY wait"
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    clients = launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "1")], 5)
+    assert all(clients[stage][0] is not None for stage in (0, 1))
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+
+
+def test_same_device_waits_for_ready_before_next_spawn(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    events: list[tuple[int, str]] = []
+
+    @contextlib.contextmanager
+    def launch(*, stage_id: int, **_kwargs: object) -> Iterator[StageReplicaResources]:
+        events.append((stage_id, "spawn"))
+        yield _resources(mocker)
+        events.append((stage_id, "ready"))
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "0")], 5)
+    assert events == [(0, "spawn"), (0, "ready"), (1, "spawn"), (1, "ready")]
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+
+
+def test_ready_failure_cleans_up_resources_and_releases_device(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    manager = mocker.Mock(spec=CoreEngineProcManager)
+    resources = StageReplicaResources(manager=manager, addresses=EngineZmqAddresses(inputs=[], outputs=[]))
+
+    @contextlib.contextmanager
+    def launch(**_kwargs: object) -> Iterator[StageReplicaResources]:
+        yield resources
+        raise RuntimeError("READY handshake failed")
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    with pytest.raises(RuntimeError, match="READY handshake failed"):
+        launch_runtime._initialize_local_llm_replica(_plan(0, "0").replicas[0], 5)
+
+    manager.shutdown.assert_called_once_with()
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+    with open(stage_init_utils.device_init_lock_path(0), "r+") as lock:
+        fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
@@ -166,6 +166,56 @@ def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker):
    assert prepared.tts_params["seed"] == [42]


+@pytest.mark.parametrize(
+    "variant,ref_text,mode",
+    [
+        ("local", " Reference. ", "continuation"),
+        ("local", None, "generation"),
+        ("local", "  ", "generation"),
+        ("tts", "Reference.", "generation"),
+    ],
+)
+def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker):
+    import torch
+
+    reference = [torch.ones((3, 12), dtype=torch.int64)]
+    unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13)
+    processor = mocker.Mock(return_value={"input_ids": unified})
+    server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}})
+    server._voice_created_at.return_value = 123
+    request = OpenAICreateSpeechRequest(
+        input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048
+    )
+    adapter = MossTTSAdapter(SpeechServingContext(server=server))
+
+    adapter._moss_variant = variant
+    mocker.patch.object(adapter, "_get_moss_processor", return_value=processor)
+    mocker.patch.object(
+        adapter, "_encode_moss_references", new=mocker.AsyncMock(return_value=(reference, {0: "reference-key"}))
+    )
+
+    prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=False))
+
+    if mode == "continuation":
+        processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English")
+        processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference)
+        conversation = [processor.build_user_message.return_value, processor.build_assistant_message.return_value]
+    else:
+        processor.build_user_message.assert_called_once_with(text="Target.", language="English", reference=reference)
+        processor.build_assistant_message.assert_not_called()
+        conversation = [processor.build_user_message.return_value]
+    processor.assert_called_once_with(conversations=[conversation], mode=mode)
+    assert prepared.prompt["prompt_token_ids"] == unified[0, :, 0].tolist()
+    assert torch.equal(prepared.tts_params["codes"]["ref"], unified[0, :, 1:])
+    assert prepared.tts_params["max_new_frames"] == [2048]
+    assert prepared.tts_params["ref_audio_cache_key"] == "reference-key"
+    assert prepared.tts_params["voice_name"] == ["speaker"]
+    assert prepared.tts_params["voice_created_at"] == [123]
+    assert prepared.tts_params["seed"] == [0]
+    assert prepared.prompt["cache_salt"]
+    assert request.input == "Target."
+
+
 def test_qwen3_tts_metadata():
    assert Qwen3TTSAdapter.backend == "ar"
    assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations

 import queue
@@ -97,9 +100,11 @@ class FakeAsyncOmniEngine:
    ) -> None:
        self.model = model
        self.config_path = None
-        self.stage_configs: list[Any] = []
        self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]]
        self.num_stages = len(self.stage_metadata)
+        self.stage_configs = [
+            StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages)
+        ]
        self.default_sampling_params_list = default_sampling_params_list or [
            SamplingParams(max_tokens=8) for _ in range(self.num_stages)
        ]
@@ -222,6 +227,25 @@ def test_resolve_sampling_params_list_preserves_stage_constraints():
    assert caller_params.stop_token_ids == [7]


+@pytest.mark.parametrize("use_defaults", [False, True])
+def test_moss_local_output_policy_preserves_codec_streaming(use_defaults):
+    from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE
+
+    base = _make_base()
+    base.engine.num_stages = 2
+    base.sampling_constraints_list = [stage.sampling_constraints for stage in MOSS_TTS_LOCAL_PIPELINE.stages]
+    base.default_sampling_params_list = [
+        base._apply_sampling_constraints(SamplingParams(), constraints)
+        for constraints in base.sampling_constraints_list
+    ]
+    caller = [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(2)]
+    result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True)
+
+    assert [params.output_kind for params in result] == [RequestOutputKind.FINAL_ONLY, RequestOutputKind.DELTA]
+    assert all(params.output_kind == RequestOutputKind.DELTA for params in caller)
+    assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY
+
+
 def _stage_spec(
    stage_id: int,
    *,
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CPU parity tests for MOSS-TTS Local Depth RoPE caching."""

 import pytest
@@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None:
    # Smaller slices retain the fixed allocation used by CUDA Graph.
    attn.prepare_rope_cache(6, device, torch.float32)
    assert attn._rope_cos_cache.data_ptr() == cache_ptr
+
+
+def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32)
+    cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8))
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(2, 12, 32)
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = model._forward_prefix(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerMultiheadAttention,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "batch,heads,queries,keys", [(2, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)]
+)
+def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys):
+    torch.manual_seed(7)
+    q = torch.randn(batch, queries, 3, heads, 64, device="cuda", dtype=torch.bfloat16)[:, :, 0].permute(0, 2, 1, 3)
+    k = torch.randn(batch, heads, keys, 64, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    positions = torch.arange(keys, device="cuda").roll(17)
+    query_positions = torch.arange(queries, device="cuda") + keys - queries
+    mask = (positions[None, :] <= query_positions[:, None])[None, None].expand(batch, 1, -1, -1).clone()
+    mask[0, :, 0] = False  # Cover an all-masked row alongside valid rows.
+    actual = masked_attention(q, k, v, mask)
+    expected = F.scaled_dot_product_attention(q, k, v, mask)
+    assert bool(actual.isfinite().all())
+    torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01)
+    assert actual[0, :, 0].count_nonzero() == 0
+
+
+def test_streaming_attention_preserves_ring_wrap_reordering_and_slot_reset():
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128, 2, causal=True, context=9, device="cuda", dtype=torch.bfloat16
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._streaming_attention = masked_attention
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, 3, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=context)
+            actual = candidate(x, x, x, execution_context=context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        torch.testing.assert_close(candidate._streaming_state.offset, reference._streaming_state.offset, rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import (
+    MossTTSCodecDecoder,
+    _MossCodecStreamSession,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch)
+
+    def reset_decoder_state_slots(self, slots):
+        self.state[slots] = 0
+
+    def decode_streaming_batch(self, codes, lengths, slots, valid_rows):
+        audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None]
+        self.state[slots] = audio[:, -1]
+        return SimpleNamespace(audio=audio[:, None], audio_lengths=lengths)
+
+
+def session():
+    return _MossCodecStreamSession(CausalCodec(), state_capacity=8, n_vq=2, vllm_config=None)
+
+
+def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
+    s = session()
+    assert [s.acquire(), s.acquire()] == [0, 1]
+    plan = {0: torch.ones(2, 3), 1: torch.ones(2, 15)}
+    with pytest.raises(ValueError, match="Only terminal"):
+        s.step(plan)
+    assert s._codec.state.count_nonzero() == 0
+    s.step({0: torch.ones(2, 1), 1: torch.full((2, 1), 2)})
+    actual = s.step(plan, terminal_slots={0})
+    torch.testing.assert_close(actual[0], torch.tensor([[4.0, 6.0, 8.0]]))
+    torch.testing.assert_close(actual[1], torch.arange(6.0, 36.0, 2.0)[None])
+    s.release(0, state_already_reset=True)
+    assert s.acquire() == 0
+    resumed = s.step({0: torch.ones(2, 1), 1: torch.ones(2, 1)})
+    assert [resumed[i].item() for i in (0, 1)] == [2, 36]
+
+
+@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
+def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls):
+    s = session()
+    s._cudagraph_wrapper = (
+        None
+        if graph_capacity is None
+        else SimpleNamespace(
+            batch_sizes=[1, graph_capacity],
+            _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None),
+            decode=lambda *args, **kwargs: None,
+        )
+    )
+    step = mocker.spy(s, "step")
+    decoder = object.__new__(MossTTSCodecDecoder)
+    nn.Module.__init__(decoder)
+    decoder._stream_max_step_frames = 15
+    decoder._stream_state_capacity = 8
+    decoder._stream_req_slots = {}
+    decoder._ensure_stream_session = lambda: s
+    lengths = [3, 7, 8, 1, 15]
+    finished = [True, True, False, True, False]
+    result = decoder._decode_streaming_batch(
+        [
+            (i, str(i), torch.ones(2, length), done)
+            for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
+        ]
+    )
+    assert step.call_count == expected_calls
+    assert [result[i].shape[-1] for i in range(5)] == lengths
+    if graph_capacity == 8:
+        assert list(step.call_args_list[0].args[0]) == [0, 1, 2]
+        assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1}
+        assert list(step.call_args_list[1].args[0]) == [3]  # Preserve the exact one-frame graph.
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+def test_async_packed_snapshot_survives_source_reuse():
+    values = [
+        torch.arange(12, device="cuda").view(1, 12),
+        torch.empty((0, 12), device="cuda", dtype=torch.long),
+        torch.arange(24, device="cuda").view(2, 12),
+    ]
+    expected = [v.cpu() for v in values]
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        {"list": values, "tuple": tuple(values)}, copy_stream=torch.cuda.Stream(), pin_memory=True
+    )
+    for v in values:
+        v.fill_(-99)
+    snapshot.wait()
+    torch.testing.assert_close(snapshot.payload, {"list": expected, "tuple": tuple(expected)}, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("offsets", [[0, 1], [1, 4]])
+def test_mtp_preserves_rows_and_snapshot_without_cuda_sync(monkeypatch, offsets):
+    import vllm_omni.worker.gpu_model_runner as mod
+    from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.model.gpu_resident_buffer_keys = {("codes", "audio")}
+    for name in ("talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"):
+        buffer = getattr(runner, name)
+        buffer.gpu = buffer.gpu.to("cuda")
+    embeds = torch.ones((2, 4), device="cuda")
+    codes = torch.arange(2, device="cuda").view(2, 1)
+    runner.talker_mtp = lambda *args, **kwargs: (embeds, codes)
+    output = torch.full((5, 4), -10.0, device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        runner._talker_mtp_forward(["r1", "r2"], output, offsets)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    expected = torch.full((5, 4), -10.0)
+    expected[offsets] = 1.0
+    torch.testing.assert_close(output.cpu(), expected, rtol=0, atol=0)
+    codes.fill_(999)
+    assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0
+    assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1
+
+
+def test_mixed_stop_flags_do_not_synchronize_cuda():
+    model = object.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.n_vq = 12
+    model.audio_pad_token_id = 1024
+    valid = torch.arange(12, device="cuda").view(1, 12)
+    stopped = torch.full_like(valid, 1024)
+    infos = [{"audio_codes": {"current": valid}}, {}, {"audio_codes": {"current": stopped}}]
+    hidden = torch.ones((3, 4), device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        output = model.make_omni_output(hidden, runtime_additional_information=infos)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    assert model._batch_should_continue.tolist() == [True, True, False]
+    torch.testing.assert_close(output.multimodal_outputs["codes"]["audio"], [valid, valid[:0], stopped])
+    assert all("current" not in info.get("audio_codes", {}) for info in infos)
+    prefill = model.make_omni_output(hidden[:1], runtime_additional_information=[{}])
+    assert model._batch_should_continue is None and prefill.multimodal_outputs == {}
@@ -464,25 +464,33 @@ def _transcribe_en_f32_16k(wav_f32: np.ndarray) -> str:
        return ""
    with _lock:
        assert _en_processor is not None and _en_model is not None and _device is not None
+        # Whisper's default feature extraction truncates at 30 seconds. Keep
+        # abnormal tails so WER evaluates the complete generated response.
+        long_audio = len(wav_f32) > 30 * 16000
+        processor_kwargs = {"truncation": False, "padding": "longest"} if long_audio else {}
        try:
            inputs = _en_processor(
                wav_f32,
                sampling_rate=16000,
                return_tensors="pt",
                return_attention_mask=True,
+                **processor_kwargs,
            )
        except TypeError:
-            inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt")
+            inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt", **processor_kwargs)
        input_features = inputs.input_features.to(_device)
        attention_mask = getattr(inputs, "attention_mask", None)
        if attention_mask is None and isinstance(inputs, dict):
            attention_mask = inputs.get("attention_mask")
        generate_kwargs: dict[str, Any] = {}
+        if long_audio:
+            generate_kwargs["return_timestamps"] = True
        if attention_mask is not None:
            generate_kwargs["attention_mask"] = attention_mask.to(_device)
        with torch.no_grad():
            try:
-                forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe")
+                prompt_kwargs = {"no_timestamps": False} if long_audio else {}
+                forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe", **prompt_kwargs)
                predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced, **generate_kwargs)
            except Exception:
                predicted_ids = _en_model.generate(
@@ -548,14 +556,13 @@ def _expand_seed_tts_turn_outputs(
        turn_pcm = getattr(output, "tts_turn_pcm_bytes", None)
        session_pcm = getattr(output, "tts_output_pcm_bytes", None)
        for turn_index, turn in enumerate(turns):
-            expanded_requests.append(
-                replace(
-                    request,
-                    prompt=turn.target_text,
-                    seed_tts_utterance_id=turn.utterance_id,
-                    seed_tts_turns=(),
-                )
+            turn_request = replace(
+                request,
+                seed_tts_utterance_id=turn.utterance_id,
+                seed_tts_turns=(),
            )
+            turn_request.prompt = turn.target_text
+            expanded_requests.append(turn_request)
            turn_output = copy(output)
            if isinstance(turn_pcm, list) and turn_index < len(turn_pcm):
                turn_output.tts_output_pcm_bytes = turn_pcm[turn_index]
@@ -0,0 +1,64 @@
+# MOSS-TTS Local Transformer v1.5: both stages on one H200.
+# Select one free physical GPU with CUDA_VISIBLE_DEVICES; devices below are logical.
+# Capacity and graph capture require more than 80 GiB of GPU memory.
+async_chunk: true
+
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      connector_get_sleep_s: 0.005
+      connector_get_max_wait_first_chunk: 1000
+      connector_get_max_wait: 300
+      initial_codec_chunk_frames: 1
+      codec_chunk_frames: 15
+      codec_left_context_frames: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 256
+    gpu_memory_utilization: 0.60
+    kv_cache_memory_bytes: 34359738368  # 32 GiB; leave room for codec state and graphs.
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 65536
+    max_model_len: 4096
+    devices: "0"
+    compilation_config:
+      # Capture mixed prefill/decode batches by total scheduled tokens.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 160, 192, 224, 256, 320, 384, 448, 512]
+      max_cudagraph_capture_size: 512
+    output_connectors:
+      to_stage_1: shm
+    default_sampling_params:
+      temperature: 1.7
+      top_p: 0.8
+      top_k: 25
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 256
+    gpu_memory_utilization: 0.30
+    enforce_eager: false
+    enable_prefix_caching: false
+    async_scheduling: true
+    compilation_config:
+      # Codec execution-batch buckets; frame counts come from the connector.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 160, 192, 256]
+      cudagraph_num_of_warmups: 1
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    hf_overrides:
+      codec_attention_backend: triton
+    input_connectors:
+      from_stage_0: shm
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
@@ -17,9 +17,10 @@ from vllm_omni.platforms import current_omni_platform

 # Pointwise fusions are launch/memory-traffic reductions whose first execution
 # for every runtime key is checked bit-for-bit; failures use the ordinary model
-# math. Both paths remain SM90-only until other architectures are validated.
-# The FNA schedule is performance-tuned for SM90 and has no silent fallback.
-_CUDA_COMPUTE_CAPABILITIES = frozenset({90})
+# math. The FNA schedule is performance-tuned for SM90 and has no silent
+# fallback, so qualify the two paths independently.
+_FUSION_CUDA_COMPUTE_CAPABILITIES = frozenset({90, 100, 103})
+_FNA_CUDA_COMPUTE_CAPABILITIES = frozenset({90})


@dataclass(frozen=True)
@@ -66,9 +67,10 @@ def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | Non
    if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available():
        return None
    device_index = device.index if device.index is not None else torch.accelerator.current_device_index()
-    if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES:
+    capability = _cuda_compute_capability(int(device_index))
+    if capability not in _FUSION_CUDA_COMPUTE_CAPABILITIES:
        return None
-    return _operator_set(has_tilelang())
+    return _operator_set(capability in _FNA_CUDA_COMPUTE_CAPABILITIES and has_tilelang())


 def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None:
@@ -78,7 +80,7 @@ def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None:


 def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool:
-    """Allow self-verifying pointwise fusions on validated SM90 devices."""
+    """Allow self-verifying pointwise fusions on qualified CUDA devices."""

    return _eligible_operators(tensor) is not None

@@ -11,7 +11,7 @@ import os
 import threading
 import time
 from collections.abc import Callable, Iterator, Mapping, Sequence
-from contextlib import contextmanager
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass
 from typing import Any, cast

@@ -149,12 +149,9 @@ class StageRuntime:
        self.stage_pools: list[StagePool] = []
        self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
        self._spawn_device_lock = threading.Lock()
-        # Serialize all LLM replica spawning + handshake across device groups
-        # to prevent ZMQ port-allocation races (get_engine_zmq_addresses) and
-        # CUDA-context conflicts when multiple engine core subprocesses
-        # initialize simultaneously on different GPUs.  Matches the old
-        # AsyncOmniEngine._initialize_llm_replica pattern which used a single
-        # ``llm_stage_launch_lock`` for all replicas.
+        # Serialize process spawning and process-global environment overlays.
+        # Readiness waits use per-device initialization protection so different
+        # device groups can initialize concurrently.
        self._replica_launch_lock = threading.Lock()
        self._init_visible_devices_baseline: str | None = None

@@ -778,43 +775,25 @@ class StageRuntime:
                spawn_device_lock=self._spawn_device_lock,
                omni_parallel_stage_init=self._parallel_stage_init,
            )
-            # G2 launch lock serializes engine-core *spawning* across replicas
-            # (ZMQ port-allocation races + simultaneous CUDA context init).
-            #   * Default path: hold it across the whole launch context manager,
-            #     whose __exit__ waits for READY — this serializes the full child
-            #     init (spawn + load + profile + KV + capture).
-            #   * Parallel path: hold it only around the spawn (__enter__); run
-            #     the READY-wait (__exit__) outside the lock so replicas init
-            #     concurrently, coordinated by the child SH/EX device locks.
-            if self._parallel_stage_init:
-                # The per-stage runtime.env overlay mutates os.environ
-                # (process-global), so it must be applied under the launch lock
-                # and only needs to cover the spawn (__enter__) — children
-                # inherit the env at spawn time; the READY-wait needs no env.
-                g2_start = time.perf_counter()
+            # Only process spawning and the runtime.env overlay need the global
+            # launch lock. The launch context's exit waits for READY: holding the
+            # lock there serializes model loading and compilation even across
+            # different GPUs. Default initialization keeps its per-device EX
+            # locks until READY; parallel_stage_init uses child phase locks.
+            g2_start = time.perf_counter()
+            with ExitStack() as launch_stack:
                with self._replica_launch_lock:
-                    with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
-                        resources = launch_cm.__enter__()
-                g2_spawned = time.perf_counter()
-                launch_cm.__exit__(None, None, None)
-                logger.debug(
-                    "[stage_init] Stage-%s G2 spawn(locked)=%.3fs, READY(unlocked)=%.3fs",
-                    plan.metadata.stage_id,
-                    g2_spawned - g2_start,
-                    time.perf_counter() - g2_spawned,
-                )
-            else:
-                g2_start = time.perf_counter()
-                with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
                    g2_locked = time.perf_counter()
-                    with launch_cm as resources:
-                        pass
-                logger.debug(
-                    "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn+READY=%.3fs",
-                    plan.metadata.stage_id,
-                    g2_locked - g2_start,
-                    time.perf_counter() - g2_locked,
-                )
+                    with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
+                        resources = launch_stack.enter_context(launch_cm)
+                g2_spawned = time.perf_counter()
+            logger.debug(
+                "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs",
+                plan.metadata.stage_id,
+                g2_locked - g2_start,
+                g2_spawned - g2_locked,
+                time.perf_counter() - g2_spawned,
+            )

            logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
            if resources is None:
@@ -339,6 +339,7 @@ class OmniBase(PDDisaggregationMixin):
        allow_delta_coercion: bool = False,
    ) -> Sequence[Any]:
        """Resolve request parameters; pipeline sampling constraints override caller values."""
+        normalized: Sequence[Any]
        if sampling_params_list is None:
            normalized = self.default_sampling_params_list
            # Set the output kind to delta since no params were specified
@@ -354,7 +355,9 @@ class OmniBase(PDDisaggregationMixin):
        if len(normalized) != self.num_stages:
            raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}")

-        if sampling_params_list is not None:
+        # Streaming coercion may also change a constrained output kind in
+        # the defaults (for example, an internal TTS stage's FINAL_ONLY).
+        if sampling_params_list is not None or allow_delta_coercion:
            normalized = [
                self._apply_sampling_constraints(params, constraints)
                for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True)
@@ -364,7 +367,7 @@ class OmniBase(PDDisaggregationMixin):
    @staticmethod
    def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]:
        """Extract each stage's required sampling settings from runtime configs."""
-        constraints_list = []
+        constraints_list: list[dict[str, Any]] = []
        for stage_config in stage_configs:
            constraints = getattr(stage_config, "sampling_constraints", {})
            if not isinstance(constraints, Mapping):
@@ -398,7 +401,7 @@ class OmniBase(PDDisaggregationMixin):
        req_state = self.request_states.get(request_id)
        prom = getattr(self, "prom_metrics", None)
        metrics = getattr(req_state, "metrics", None)
-        if metrics is None or prom is None:
+        if req_state is None or metrics is None or prom is None:
            return
        if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False):
            return
@@ -6,7 +6,7 @@ Both variants share the same build/validate flow (``_build_moss_tts_params``
 handles each); they are registered under distinct model-type names.
 """

-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast

 from vllm.inputs import tokens_input

@@ -61,7 +61,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
        # the process-wide speaker cache.
        encoder = build_reference_encoder(
            self._get_moss_processor(),
-            variant=self._moss_variant,
+            variant=cast(str, self._moss_variant),
            speaker_cache=self._speaker_cache,
        )
        self._moss_ref_encoder = encoder
@@ -101,7 +101,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):

        return await encode_request_references(
            self._get_moss_ref_encoder(),
-            request.ref_audio,
+            cast(str, request.ref_audio),
            request.ref_audio_2 if two_speaker else None,
            resolve_ref_audio=self._resolve_ref_audio,
            get_artifact_key=self._get_resolved_ref_audio_artifact_key,
@@ -179,15 +179,16 @@ class _MossTTSAdapterBase(ARTTSAdapter):

        v = self._moss_variant

+        params: dict[str, Any]
        # ---- Legacy nano path (unchanged) ----
        if v is None:  # moss_tts_nano
-            params: dict[str, Any] = {
+            params = {
                "text": [request.input or ""],
                "mode": ["voice_clone"],
            }
            if request.max_new_tokens is not None:
                params["max_new_frames"] = [request.max_new_tokens]
-            wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio)
+            wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
            params["prompt_audio_array"] = [[wav_list, sr]]
            params["ref_audio_cache_key"] = cache_key
            return params
@@ -201,13 +202,13 @@ class _MossTTSAdapterBase(ARTTSAdapter):
        # talker reads for short prompts; full Realtime support needs a
        # separate processor.from_module path which we don't wire here.
        if v == "realtime":
-            params: dict[str, Any] = {
+            params = {
                "text": [request.input or ""],
                "mode": ["voice_clone"],
            }
            if request.max_new_tokens is not None:
                params["max_new_frames"] = [request.max_new_tokens]
-            wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio)
+            wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
            params["prompt_audio_array"] = [[wav_list, sr]]
            params["ref_audio_cache_key"] = cache_key
            return params
@@ -254,13 +255,20 @@ class _MossTTSAdapterBase(ARTTSAdapter):
        # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special
        # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly
        # audio_pad_code outside the reference block).
-        user_msg = proc.build_user_message(**user_kwargs)
-        batch = proc(conversations=[[user_msg]], mode="generation")
+        if v == "local" and request.ref_text and request.ref_text.strip():
+            reference = user_kwargs.pop("reference")
+            user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "")
+            user_msg = proc.build_user_message(**user_kwargs)
+            assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
+            batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
+        else:
+            user_msg = proc.build_user_message(**user_kwargs)
+            batch = proc(conversations=[[user_msg]], mode="generation")
        unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
        text_ids: list[int] = unified[:, 0].tolist()
        audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64)

-        params: dict[str, Any] = {
+        params = {
            "prompt_token_ids": text_ids,
            "codes": {"ref": audio_codes},
        }
@@ -278,9 +278,9 @@ def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool):
    if params.output_kind == target_type:
        return params
    elif is_streaming and params.output_kind == RequestOutputKind.FINAL_ONLY:
-        logger.warning("Request appears to be streaming, but got request type final only!")
+        logger.debug("Coercing FINAL_ONLY output to DELTA for streaming")
    elif not is_streaming and params.output_kind == RequestOutputKind.DELTA:
-        logger.warning("Request appears to not be streaming, but got request type delta!")
+        logger.debug("Coercing DELTA output to FINAL_ONLY for non-streaming")

    if not params.skip_clone:
        params = params.clone()
@@ -48,7 +48,7 @@ class StageRequestStats:
    finish_reason: str | None = None
    request_id: str | None = None
    postprocess_time_ms: float = 0.0
-    diffusion_metrics: dict[str, float] = None
+    diffusion_metrics: dict[str, float] | None = None
    audio_generated_frames: int = 0
    audio_sample_rate: int = 0
    audio_duration_s: float = 0.0
@@ -181,11 +181,11 @@ class OrchestratorAggregator:
        self.e2e_total_ms = 0.0
        self.e2e_total_tokens = 0
        self.e2e_count = 0
-        self.e2e_done = set()
+        self.e2e_done: set[str] = set()
        self.wall_start_ts = float(wall_start_ts)
        self.last_finish_ts = float(wall_start_ts)
-        self.stage_first_ts = [None for _ in range(self.num_stages)]
-        self.stage_last_ts = [None for _ in range(self.num_stages)]
+        self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
        self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict(
            lambda: defaultdict(float)
        )  # {request_id: {stage_id:accumulated_gen_time_ms}}
@@ -618,6 +618,7 @@ class OrchestratorAggregator:
        final_output_type: str | None = None,
    ) -> None:
        stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type)
+        assert stats.stage_id is not None
        self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out)
        if stats.stage_id == 0:
            self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in)
@@ -775,10 +776,8 @@ class OrchestratorAggregator:
            final_stage_id_map = self.final_stage_id_for_e2e

        stage_wall_time_ms = [
-            ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0)
-            if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None)
-            else 0.0
-            for i in range(self.num_stages)
+            ((last - first) * 1000.0) if first is not None and last is not None else 0.0
+            for first, last in zip(self.stage_first_ts, self.stage_last_ts, strict=True)
        ]

        overall_summary = {
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -801,7 +804,16 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
        else:
            attn_bias = None

-        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+        streaming_attention = getattr(self, "_streaming_attention", None)
+        if (
+            streaming_attention is not None
+            and attn_bias is not None
+            and q.dtype == torch.bfloat16
+            and q.shape[-1] == 64
+        ):
+            x = streaming_attention(q, k, v, attn_bias)
+        else:
+            x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
        x = x.transpose(1, 2).reshape(B, T, self.embed_dim)
        x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)

@@ -145,15 +145,22 @@ class _MossCodecStreamSession:
        terminal_slots = terminal_slots or set()
        if not terminal_slots.issubset(slot_codes):
            raise ValueError("terminal_slots must be a subset of the decode batch slots.")
-        step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()}
-        if len(step_lengths) != 1:
-            raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}")
-        (step_t,) = step_lengths
+        step_t = max(int(codes.shape[1]) for codes in slot_codes.values())
+        if any(int(codes.shape[1]) != step_t and slot not in terminal_slots for slot, codes in slot_codes.items()):
+            raise ValueError("Only terminal codec rows may be padded to a larger step length")
        slots = list(slot_codes)
        if any(slot not in self._leased_slots for slot in slots):
            raise RuntimeError(f"Streaming decode references an unleased state slot: {slots}")
        codes_step = torch.stack(
-            [slot_codes[slot].to(device=self._device, dtype=torch.long) for slot in slots],
+            [
+                torch.nn.functional.pad(
+                    slot_codes[slot].to(device=self._device, dtype=torch.long),
+                    (0, step_t - int(slot_codes[slot].shape[1])),
+                )
+                if int(slot_codes[slot].shape[1]) != step_t
+                else slot_codes[slot].to(device=self._device, dtype=torch.long)
+                for slot in slots
+            ],
            dim=1,
        )
        state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
@@ -197,9 +204,9 @@ class _MossCodecStreamSession:
            self._reset_slot_ids(terminal_slot_ids)

        audio = audio_tensor.detach().to("cpu", torch.float32)
-        audio_length = step_t * self._samples_per_frame
        out: dict[int, torch.Tensor] = {}
        for row, slot in enumerate(slots):
+            audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame
            out[slot] = audio[row, ..., :audio_length].contiguous()
        return out

@@ -518,6 +525,8 @@ class MossTTSCodecDecoder(nn.Module):
        outputs: dict[int, torch.Tensor] = {}
        grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {}
        max_step_frames = max(1, int(self._stream_max_step_frames))
+        graph = session._cudagraph_wrapper
+        coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity

        for output_index, request_id, codes_nq_t, finished in items:
            slot = self._stream_req_slots.get(request_id)
@@ -548,9 +557,13 @@ class MossTTSCodecDecoder(nn.Module):
                    )
                continue

-            grouped.setdefault(int(codes_nq_t.shape[1]), []).append(
-                (output_index, request_id, slot, codes_nq_t, finished)
-            )
+            frame_count = int(codes_nq_t.shape[1])
+            # Only combine tails that already used the same padded graph.
+            # Preserve exact-size graphs (including T=1) and the eager path.
+            group_frames = frame_count
+            if coalesce_tails and finished:
+                group_frames = graph._select_frame_size(frame_count, allow_padding=True) or frame_count
+            grouped.setdefault(group_frames, []).append((output_index, request_id, slot, codes_nq_t, finished))

        for group in grouped.values():
            plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
@@ -756,6 +769,17 @@ class MossTTSCodecDecoder(nn.Module):
        codec.eval()
        if device.type != "cpu":
            codec.decoder.to(dtype=torch.bfloat16)
+        attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa")
+        if attention_backend != "sdpa":
+            if attention_backend != "triton" or device.type != "cuda":
+                raise ValueError(f"Unsupported codec attention backend/device: {attention_backend}/{device.type}")
+            from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerMultiheadAttention
+            from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
+
+            for module in codec.decoder.modules():
+                if isinstance(module, MossAudioTokenizerMultiheadAttention):
+                    module._streaming_attention = masked_attention
+            logger.info("Enabled Triton masked attention for the streaming codec decoder")
        build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None)
        if callable(build_decode_lut):
            lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -101,8 +103,13 @@ class _MossTTSLocalAttention(nn.Module):
        odd = x[..., 1::2]
        return torch.stack((-odd, even), dim=-1).reshape_as(x)

-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call."""
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        """Run a causal prefix, or one new token using frame-local K/V."""
        batch_size, seq_len, _ = hidden_states.shape
        qkv = self.c_attn(hidden_states)
        query, key, value = qkv.split(self.embed_dim, dim=-1)
@@ -110,14 +117,22 @@ class _MossTTSLocalAttention(nn.Module):
        key = key.view(batch_size, seq_len, self.n_head, self.head_dim)
        value = value.view(batch_size, seq_len, self.n_head, self.head_dim)

-        cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = self._rope_cos_sin(position + seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = cos[:, position:], sin[:, position:]
        query = query * cos + self._rotate_half(query) * sin
        key = key * cos + self._rotate_half(key) * sin

        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
-        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True)
+        if kv_cache is not None:
+            assert seq_len == 1, "Cached Local Depth execution consumes one token at a time"
+            k_cache, v_cache = kv_cache
+            k_cache[:, :, position : position + 1].copy_(key)
+            v_cache[:, :, position : position + 1].copy_(value)
+            key = k_cache[:, :, : position + 1]
+            value = v_cache[:, :, : position + 1]
+        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
        return self.c_proj(attn_output)

@@ -141,8 +156,13 @@ class _MossTTSLocalBlock(nn.Module):
        self.ln_2 = nn.LayerNorm(hidden_size, eps=eps)
        self.mlp = _MossTTSLocalMLP(hidden_size, inner_size)

-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states), kv_cache, position)
        hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
        return hidden_states

@@ -156,11 +176,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
        and codebook-0's head (``audio_lm_heads[0]``) simultaneously.
      - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is
        re-embedded (``audio_embeddings[c]``) and appended as the next
-        position, re-prefilling the block over the growing (<=n_vq) sequence
-        with a fresh causal mask each call -- mathematically identical to
-        incremental KV-cache decoding since attention is strictly causal and
-        only the last position is ever read, but avoids any cache plumbing
-        given the trivially short sequence length.
+        position. K/V for earlier positions are reused within the frame, so
+        each position's projections and MLP run once. A new frame overwrites
+        every cache position before reading it, including CUDA Graph replay.
    """

    def __init__(self, gpt2_config, hidden_size: int | None = None) -> None:
@@ -174,11 +192,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
        self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
        self._compiled_forward_prefix = None

-    def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
-        hidden_states = seq_embeds
-        for block in self.h:
-            hidden_states = block(hidden_states)
-        return self.ln_f(hidden_states)
+    def _forward_prefix(
+        self,
+        seq_embeds: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        return self.ln_f(self.h[0](seq_embeds, kv_cache, position))

    def setup_compile(self) -> None:
        if self._compiled_forward_prefix is not None:
@@ -189,14 +209,18 @@ class MossTTSLocalDepthTransformer(nn.Module):
            return
        self._compiled_forward_prefix = torch.compile(
            self._forward_prefix,
-            dynamic=False,
+            # Share batch/position shapes instead of exhausting Dynamo's
+            # recompilation budget while warming request-size graph buckets.
+            dynamic=True,
            options={"epilogue_fusion": False},
        )
-        logger.info("MOSS-TTS local depth prefix enabled with torch.compile")
+        logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile")

-    def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
+    def _run_prefix(
+        self, seq_embeds: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor], position: int
+    ) -> torch.Tensor:
        forward_prefix = self._compiled_forward_prefix or self._forward_prefix
-        return forward_prefix(seq_embeds)
+        return forward_prefix(seq_embeds, kv_cache, position)

    @torch.no_grad()
    def generate_frame(
@@ -240,9 +264,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
        for block in self.h:
            block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype)

-        embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
-        embeds[:, 0, :] = backbone_last_hidden.to(dtype)
-        hidden = self._run_prefix(embeds[:, :1, :])
+        attn = self.h[0].attn
+        cache_shape = (batch_size, attn.n_head, n_vq, attn.head_dim)
+        kv_cache = (
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+        )
+        hidden = self._run_prefix(backbone_last_hidden[:, None, :].to(dtype), kv_cache, 0)
        local_hidden = hidden[:, 0, :]

        binary_logits = local_text_lm_head(local_hidden).float()
@@ -292,9 +320,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
            codes[:, channel_index] = channel_token

            if channel_index + 1 < n_vq:
-                embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype)
-                hidden = self._run_prefix(embeds[:, : channel_index + 2, :])
-                local_hidden = hidden[:, channel_index + 1, :]
+                embeds = audio_embeddings[channel_index](channel_token).to(dtype)[:, None, :]
+                hidden = self._run_prefix(embeds, kv_cache, channel_index + 1)
+                local_hidden = hidden[:, 0, :]

        return should_continue, codes

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -1671,20 +1674,19 @@ class MossTTSLocalTalkerForGeneration(nn.Module):

        self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
        self._batch_state_spans = kwargs.get("request_token_spans")
-        should_values: list[torch.Tensor] = []
-        have_mtp_output = False
+        should_codes: list[torch.Tensor] = []
+        should_indices: list[int] = []
        per_req_codes: list[torch.Tensor] = []
        have_codes = False
-        for info in info_dicts:
+        for index, info in enumerate(info_dicts):
            audio_codes = (info.get("audio_codes", {}) or {}) if isinstance(info, dict) else {}
            current = audio_codes.pop("current", None)
            if isinstance(current, torch.Tensor) and current.numel() > 0:
                codes = current.to(device=hidden.device, dtype=torch.long)
                if codes.dim() == 1:
                    codes = codes.unsqueeze(0)
-                should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1)
-                should_values.append(should_continue.reshape(-1)[:1])
-                have_mtp_output = True
+                should_codes.append(codes[:1])
+                should_indices.append(index)
                # Do not boolean-index a CUDA tensor here.  Even for the
                # common single-row case, ``codes[should_continue]`` must
                # discover a dynamic output size on the host and therefore
@@ -1694,9 +1696,19 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                per_req_codes.append(codes)
                have_codes = True
            else:
-                should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool))
                per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
-        self._batch_should_continue = torch.cat(should_values, dim=0) if have_mtp_output and should_values else None
+        self._batch_should_continue = None
+        if should_codes:
+            active = torch.cat(should_codes, dim=0).ne(self.audio_pad_token_id).any(dim=-1)
+            if len(should_codes) == len(info_dicts):
+                self._batch_should_continue = active
+            else:
+                self._batch_should_continue = torch.ones(len(info_dicts), device=hidden.device, dtype=torch.bool)
+                # Keep mixed prefill/decode output packing asynchronous.
+                indices = torch.tensor(should_indices, device="cpu", dtype=torch.long).to(
+                    hidden.device, non_blocking=True
+                )
+                self._batch_should_continue.index_copy_(0, indices, active)

        if not have_codes:
            return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
@@ -1,8 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
 """Pipeline topology for all MOSS-TTS variants (2-stage: talker → codec)."""

+from vllm.sampling_params import RequestOutputKind
+
 from vllm_omni.config.stage_config import (
    PipelineConfig,
    StageExecutionType,
@@ -100,6 +105,10 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
            sampling_constraints={
                "detokenize": False,
                "stop_token_ids": [151645],
+                # The worker connector streams codes directly to the codec.
+                # This internal stage only needs to publish its terminal
+                # result; codec audio output remains incremental.
+                "output_kind": RequestOutputKind.FINAL_ONLY,
            },
        ),
        StagePipelineConfig(
@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Masked streaming attention preserving ring-cache position semantics."""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _attention(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    mask_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    ms0: tl.constexpr,
+    ms2: tl.constexpr,
+    ms3: tl.constexpr,
+    num_heads: tl.constexpr,
+    q_len: tl.constexpr,
+    kv_len: tl.constexpr,
+    head_dim: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    dims = tl.arange(0, head_dim)
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3,
+        mask=rows[:, None] < q_len,
+        other=0,
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    for first in range(tl.cdiv(kv_len, block_n)):
+        cols = first * block_n + tl.arange(0, block_n)
+        k = tl.load(
+            k_ptr + batch * ks0 + head * ks1 + cols[None, :] * ks2 + dims[:, None] * ks3,
+            mask=cols[None, :] < kv_len,
+            other=0,
+        )
+        allowed = tl.load(
+            mask_ptr + batch * ms0 + rows[:, None] * ms2 + cols[None, :] * ms3,
+            mask=(rows[:, None] < q_len) & (cols[None, :] < kv_len),
+            other=0,
+        )
+        score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+        score = tl.where(allowed & (cols[None, :] < kv_len), score, -float("inf"))
+        next_max = tl.maximum(maximum, tl.max(score, axis=1))
+        safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+        rescale = tl.exp(maximum - safe_max)
+        p = tl.exp(score - safe_max[:, None])
+        denominator = denominator * rescale + tl.sum(p, axis=1)
+        v = tl.load(
+            v_ptr + batch * vs0 + head * vs1 + cols[:, None] * vs2 + dims[None, :] * vs3,
+            mask=cols[:, None] < kv_len,
+            other=0,
+        )
+        acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+        maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * q_len + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        mask=rows[:, None] < q_len,
+    )
+
+
+@torch.library.custom_op("moss_streaming::masked_attention", mutates_args=())
+def masked_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    b, h, t, d = q.shape
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    # Small query tiles expose parallelism for short streaming chunks. Larger
+    # tiles reuse K/V for steady chunks. Keep 64-key softmax blocks to retain
+    # the same accumulation granularity across both paths.
+    bm = 16 if t <= 32 else 64
+    num_warps = 8 if k.shape[2] <= 256 else 4
+    _attention[(triton.cdiv(t, bm), b * h)](
+        q,
+        k,
+        v,
+        mask,
+        out,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        mask.stride(0),
+        mask.stride(2),
+        mask.stride(3),
+        h,
+        t,
+        k.shape[2],
+        d,
+        bm,
+        64,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+    return out
+
+
+@masked_attention.register_fake
+def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass, field
 from typing import Any

@@ -52,7 +55,7 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
    """

    multimodal_outputs: list[dict[str, object]] | None = None
-    inter_stage_outputs: list[dict[str, object]] | None = None
+    inter_stage_outputs: list[dict[str, Any] | None] | None = None
    # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
    # The Scheduler can safely free the block tables for these requests.
    kv_extracted_req_ids: list[str] | None = None
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""XPU-specific patches for upstream vLLM behavior."""
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _patch_w8a16_fp8_nd_activation() -> None:
+    """Restore the N-D output shape for the XPU W8A16 FP8 linear kernel.
+
+    ``XPUW8A16FP8LinearKernel.apply_weights`` hands its activation straight to
+    ``torch.ops._xpu_C.fp8_gemm_w8a16``, skipping the flatten-then-reshape every
+    other ScaledMM kernel does. The op's registered fake always returns 2-D, so a
+    3-D activation makes inductor bake in a 2-D ``assert_size_stride`` that the
+    rank-3 result trips, killing FLUX.1 ``--quantization fp8`` at startup warmup.
+    Only >2-D inputs are affected, so LLM (token-flattened, 2-D) paths are not.
+    """
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
+            XPUW8A16FP8LinearKernel,
+        )
+    except ImportError:
+        logger.debug("XPUW8A16FP8LinearKernel not available; skipping FP8 shape patch.")
+        return
+
+    _original_apply_weights = XPUW8A16FP8LinearKernel.apply_weights
+
+    def _patched_apply_weights(self, layer, x, bias=None):
+        output_shape = (*x.shape[:-1], layer.weight.shape[1])
+        x = x.reshape(-1, x.shape[-1])
+        return _original_apply_weights(self, layer, x, bias).view(*output_shape)
+
+    XPUW8A16FP8LinearKernel.apply_weights = _patched_apply_weights
+
+
+def apply_patches() -> None:
+    """Apply all XPU-specific patches."""
+    _patch_w8a16_fp8_nd_activation()
@@ -11,6 +11,7 @@ from vllm.platforms.xpu import XPUPlatform

 from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.platforms.xpu.patch import apply_patches

 logger = init_logger(__name__)

@@ -31,6 +32,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):

    _omni_enum = OmniPlatformEnum.XPU

+    def __init__(self):
+        super().__init__()
+        apply_patches()
+
    @classmethod
    def get_omni_ar_worker_cls(cls) -> str:
        return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker"
@@ -15,7 +15,7 @@ from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
-from typing import Any, NamedTuple
+from typing import Any, NamedTuple, cast

 import numpy as np
 import torch
@@ -87,14 +87,39 @@ def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any:
        return value.detach().clone()
    if isinstance(value, dict):
        return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()}
-    if isinstance(value, list):
-        return [_clone_cuda_tensor_payload(v, sources) for v in value]
-    if isinstance(value, tuple):
-        return tuple(_clone_cuda_tensor_payload(v, sources) for v in value)
+    if isinstance(value, (list, tuple)):
+        if (
+            value
+            and all(
+                isinstance(v, torch.Tensor)
+                and v.is_cuda
+                and v.dtype == value[0].dtype
+                and v.device == value[0].device
+                and v.layout == torch.strided
+                and v.is_contiguous()
+                for v in value
+            )
+            # Keep frame/state packing bounded; large payloads use individual snapshots.
+            and sum(v.numel() * v.element_size() for v in value) <= 1024 * 1024
+        ):
+            packed = torch.cat([v.detach().reshape(-1) for v in value])
+            sources.append(packed)
+            return _PackedTensorPayload(packed, [v.shape for v in value], isinstance(value, tuple))
+        items = [_clone_cuda_tensor_payload(v, sources) for v in value]
+        return tuple(items) if isinstance(value, tuple) else items
    return value


 def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
+    if isinstance(value, _PackedTensorPayload):
+        flat = _copy_tensor_payload_to_cpu(value.tensor, pin_memory)
+        items = []
+        offset = 0
+        for shape in value.shapes:
+            count = shape.numel()
+            items.append(flat[offset : offset + count].view(shape))
+            offset += count
+        return tuple(items) if value.is_tuple else items
    if isinstance(value, torch.Tensor):
        if value.device.type != "cuda":
            return value
@@ -110,6 +135,12 @@ def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
    return value


+class _PackedTensorPayload(NamedTuple):
+    tensor: torch.Tensor
+    shapes: list[torch.Size]
+    is_tuple: bool
+
+
 class _AsyncCPUPayloadSnapshot:
    def __init__(
        self,
@@ -181,7 +212,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
        if kwargs:
            raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")

-        self._model_runner_output = None
+        self._model_runner_output: OmniModelRunnerOutput | None = None
        self._invalid_req_indices = invalid_req_indices

        self.async_copy_ready_event = torch.Event()
@@ -211,7 +242,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
            self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
            self.async_copy_ready_event.record()

-        self._model_runner_output_builder = model_runner_output_builder
+        self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder
        self._background_exception: BaseException | None = None
        self._background_thread: threading.Thread | None = None
        self._cuda_device = cuda_device
@@ -226,6 +257,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
        if self._model_runner_output is not None:
            return
        with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"):
+            assert self._model_runner_output_builder is not None
            self._model_runner_output = self._model_runner_output_builder()
        self._model_runner_output_builder = None

@@ -321,6 +353,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
    outputs per request while keeping Async output semantics.
    """

+    execute_model_state: ExecuteModelState | None
+    kv_extracted_req_ids: list[str] | None
+
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
@@ -436,6 +471,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
        if not isinstance(info, dict):
            return None
        val = info.get("omni_final_stage_id")
+        if val is None:
+            return None
        try:
            return int(val)
        except (TypeError, ValueError):
@@ -1485,7 +1522,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                    )
            selected_indices.append(idx)
            selected_req_ids.append(req_id)
-            selected_req_infos.append(req_info)
+            selected_req_infos.append(cast(dict[str, Any], req_info))

        if not selected_indices:
            return multimodal_outputs
@@ -1552,7 +1589,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex

    def _build_multimodal_outputs(
        self,
-        per_req_payloads: list[dict[str, object] | None] | None,
+        per_req_payloads: Sequence[dict[str, object] | None] | None,
    ) -> list[dict[str, torch.Tensor] | None] | None:
        """Build per-request multimodal output payloads (dedicated channel).

@@ -1648,8 +1685,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex

    def _build_omni_step_outputs(
        self,
-        pooler_inter: list[dict[str, object]],
-        pooler_client: list[dict[str, object]],
+        pooler_inter: Sequence[dict[str, object] | None] | None,
+        pooler_client: Sequence[dict[str, object] | None] | None,
        *,
        defer_full_payload_d2h: bool,
    ) -> tuple[
@@ -1935,6 +1972,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                    )
                    pooler_output.append(flatten_payload(payload))

+        pooler_inter: Sequence[dict[str, object] | None] | None
+        pooler_client: Sequence[dict[str, object] | None] | None
        pooler_output = pooler_output or []
        if self._async_chunk and stage_sends_async_output(self.model_config):
            pooler_inter, pooler_client = partition_payload_list(pooler_output)
@@ -78,6 +78,8 @@ def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dic


 class OmniGPUModelRunner(GPUModelRunner):
+    intermediate_tensors: IntermediateTensors | None
+
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_intermediate_buffer: dict[str, dict[str, Any]] = {}
@@ -842,7 +844,9 @@ class OmniGPUModelRunner(GPUModelRunner):
            return None

    @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict:
+    def extract_multimodal_outputs(
+        self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput
+    ) -> tuple[Any, Any]:
        if (
            hasattr(self.model, "have_multimodal_outputs")
            and self.model.have_multimodal_outputs
@@ -1407,14 +1411,14 @@ class OmniGPUModelRunner(GPUModelRunner):
    def _process_additional_information_updates(
        self,
        hidden_states: torch.Tensor,
-        multimodal_outputs: object,
+        multimodal_outputs: Any,
        num_scheduled_tokens_np: np.ndarray,
        scheduler_output: "SchedulerOutput",
        combined_hidden_states: dict[str, torch.Tensor] | None = None,
-        combined_multimodal_outputs: dict[str, object] | None = None,
+        combined_multimodal_outputs: dict[str, Any] | None = None,
        req_ids_filter: set[str] | None = None,
        req_ids: list[str] | None = None,
-        query_start_loc_cpu: object | None = None,
+        query_start_loc_cpu: Any = None,
    ) -> None:
        """Process model-provided per-request updates and merge into model_intermediate_buffer."""
        req_ids = req_ids if req_ids is not None else self.input_batch.req_ids
@@ -1476,9 +1480,9 @@ class OmniGPUModelRunner(GPUModelRunner):
    def _collect_additional_information_for_prefill(
        self,
        num_scheduled_tokens_np: np.ndarray,
-    ) -> dict[str, dict]:
+    ) -> None:
        """Overlay per-request prompt_embeds for the prefill portion and collect
-        additional_information slices for this step. Returns a map req_id -> dict."""
+        additional_information slices for this step."""
        for req_index, req_id in enumerate(self.input_batch.req_ids):
            req_state = self.requests[req_id]
            pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
@@ -1707,9 +1711,9 @@ class OmniGPUModelRunner(GPUModelRunner):

            # Overlay custom prompt_embeds per request for the prompt portion;
            # collect additional_information (tensor/list) for prefill portion only
-            decode_req_ids = []
-            decode_start_offsets = []
-            decode_batch_items = []
+            decode_req_ids: list[str] = []
+            decode_start_offsets: list[int] = []
+            decode_batch_items: list[tuple[str, int, dict[str, Any]]] = []
            batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None)

            def flush_decode_batch() -> None:
@@ -1717,6 +1721,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                if not decode_batch_items:
                    return

+                assert callable(batch_decode_preprocess)
                req_ids_b = [item[0] for item in decode_batch_items]
                start_offsets_b = [item[1] for item in decode_batch_items]
                req_infos_b = [item[2] for item in decode_batch_items]
@@ -1761,7 +1766,7 @@ class OmniGPUModelRunner(GPUModelRunner):

                # mimo-audio check
                req_state = self.requests.get(req_id)
-                req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)
+                req_infos = cast(dict[str, Any], self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id))

                start_offset = int(self.query_start_loc.cpu[req_index])
                sched_tokens = int(num_scheduled_tokens_np[req_index])
@@ -1968,11 +1973,29 @@ class OmniGPUModelRunner(GPUModelRunner):
        if start_offsets is None:
            id_to_index = self.input_batch.req_id_to_index
            start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids]
-        for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)):
-            inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            if code_predictor_codes is not None:
-                update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-                self._update_intermediate_buffer(req_id, update_dict)
+        if start_offsets == list(range(decode_batch_size)):
+            inputs_embeds[:decode_batch_size].copy_(req_embeds[:decode_batch_size])
+        else:
+            # A device tensor constructor waits for the preceding MTP graph.
+            # Enqueue the small host index transfer without synchronizing it.
+            offsets = torch.tensor(start_offsets, device="cpu", dtype=torch.long).to(
+                inputs_embeds.device, non_blocking=True
+            )
+            inputs_embeds.index_copy_(0, offsets, req_embeds[:decode_batch_size])
+        if code_predictor_codes is not None:
+            if out_key in getattr(self.model, "gpu_resident_buffer_keys", set()):
+                # One owned batch snapshot protects all rows from graph reuse.
+                owned_codes = code_predictor_codes[:decode_batch_size].detach().clone()
+                for req_id, row in zip(decode_req_ids, owned_codes.split(1), strict=True):
+                    req_state = self.requests.get(req_id)
+                    if req_state is not None:
+                        existing = self.model_intermediate_buffer.setdefault(req_id, {})
+                        existing.setdefault(out_key[0], {})[out_key[1]] = row
+                        req_state.additional_information_cpu = existing
+            else:
+                for idx, req_id in enumerate(decode_req_ids):
+                    update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
+                    self._update_intermediate_buffer(req_id, update_dict)

    def _model_forward(
        self,