diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index 2c58a7c3f..f51f13f1a 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -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 diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 3772ba59e..749f35f5a 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -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", diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index d7e14ed64..f54a58fa4 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -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 diff --git a/.buildkite/npu/bootstrap-upload-steps.yml b/.buildkite/npu/bootstrap-upload-steps.yml index 64b97f895..398fbc6dc 100644 --- a/.buildkite/npu/bootstrap-upload-steps.yml +++ b/.buildkite/npu/bootstrap-upload-steps.yml @@ -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 diff --git a/.buildkite/npu/scripts/build-npu-image.sh b/.buildkite/npu/scripts/build-npu-image.sh new file mode 100644 index 000000000..415fa7d5e --- /dev/null +++ b/.buildkite/npu/scripts/build-npu-image.sh @@ -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}" diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 01fb32879..d475aad5b 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -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 diff --git a/.claude/skills/review-pr/SKILL.md b/.claude/skills/review-pr/SKILL.md index 0b19b2b37..5907e13f3 100644 --- a/.claude/skills/review-pr/SKILL.md +++ b/.claude/skills/review-pr/SKILL.md @@ -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 diff --git a/.claude/skills/review-pr/references/checks/perf-verification.md b/.claude/skills/review-pr/references/checks/perf-verification.md index bbfe5a147..027e91589 100644 --- a/.claude/skills/review-pr/references/checks/perf-verification.md +++ b/.claude/skills/review-pr/references/checks/perf-verification.md @@ -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. diff --git a/.claude/skills/review-pr/references/process/general-checks.md b/.claude/skills/review-pr/references/process/general-checks.md index 88dfd405e..da86eabc7 100644 --- a/.claude/skills/review-pr/references/process/general-checks.md +++ b/.claude/skills/review-pr/references/process/general-checks.md @@ -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 : :` 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, diff --git a/.claude/skills/review-pr/references/process/review-execution.md b/.claude/skills/review-pr/references/process/review-execution.md index 7c0305254..c78532e9d 100644 --- a/.claude/skills/review-pr/references/process/review-execution.md +++ b/.claude/skills/review-pr/references/process/review-execution.md @@ -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 :`, and existing CI evidence; report all executable validation as a gap. diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 3ebacdc98..0d91c0b32 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -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 \ diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 deleted file mode 100644 index d341d456d..000000000 --- a/docker/Dockerfile.npu.a3 +++ /dev/null @@ -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 [] diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 029a905fd..0a150fdd8 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -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 \ diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 deleted file mode 100644 index fe3741d04..000000000 --- a/docker/Dockerfile.npu.ci.a3 +++ /dev/null @@ -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 [] diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index d578b4b02..885ebd528 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -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** diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md index 469e69cac..762863bb3 100644 --- a/docs/contributing/ci/test_writing_guide.md +++ b/docs/contributing/ci/test_writing_guide.md @@ -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 diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 0019d00bd..9ff477522 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -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] diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index 2f00d386b..894e2bf0e 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -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 diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 5dbe58314..5141a37c2 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -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: ""`). ### 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. diff --git a/pyproject.toml b/pyproject.toml index 3ba902977..0393c9c62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md index b9d32947b..efd88c068 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md +++ b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md @@ -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: +- 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. -``` diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index e9cb798fe..8d37039d8 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -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/) diff --git a/recipes/README.md b/recipes/README.md index 6b80f8a03..ea4576201 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 | diff --git a/tests/benchmarks/test_seed_tts_dataset_variants.py b/tests/benchmarks/test_seed_tts_dataset_variants.py index 31f66bd7d..17ee6df4c 100644 --- a/tests/benchmarks/test_seed_tts_dataset_variants.py +++ b/tests/benchmarks/test_seed_tts_dataset_variants.py @@ -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 diff --git a/tests/diffusion/layers/test_activation.py b/tests/diffusion/layers/test_activation.py index c3205502f..321f2ab8f 100644 --- a/tests/diffusion/layers/test_activation.py +++ b/tests/diffusion/layers/test_activation.py @@ -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( diff --git a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py index 6ffff3136..5fca0e5e6 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py +++ b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py @@ -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), diff --git a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py index d3d885d57..aa022ed61 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py @@ -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) diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 0da5b6dd3..5793ad95c 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -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" diff --git a/tests/engine/test_stage_runtime_launch_concurrency.py b/tests/engine/test_stage_runtime_launch_concurrency.py new file mode 100644 index 000000000..f880b9386 --- /dev/null +++ b/tests/engine/test_stage_runtime_launch_concurrency.py @@ -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) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index a2c0ace0b..a17a6ccbd 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -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) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 4f9fc592b..6fe04fdb7 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -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, *, diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py index 90c386ee0..566ba4595 100644 --- a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py +++ b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py @@ -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) diff --git a/tests/model_executor/models/moss_tts/test_streaming_attention.py b/tests/model_executor/models/moss_tts/test_streaming_attention.py new file mode 100644 index 000000000..3134575ca --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_attention.py @@ -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")) diff --git a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py new file mode 100644 index 000000000..fff97cd96 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py @@ -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. diff --git a/tests/worker/test_batched_omni_output.py b/tests/worker/test_batched_omni_output.py new file mode 100644 index 000000000..d0b18b3c1 --- /dev/null +++ b/tests/worker/test_batched_omni_output.py @@ -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 == {} diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 9dc99548a..100933870 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -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] diff --git a/vllm_omni/deploy/moss_tts_local_h200.yaml b/vllm_omni/deploy/moss_tts_local_h200.yaml new file mode 100644 index 000000000..6fb5e010d --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_h200.yaml @@ -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 diff --git a/vllm_omni/diffusion/models/ltx2/ops/__init__.py b/vllm_omni/diffusion/models/ltx2/ops/__init__.py index 93f2ff672..7907fba56 100644 --- a/vllm_omni/diffusion/models/ltx2/ops/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/ops/__init__.py @@ -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 diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 5d7d07653..53400df5e 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -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: diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 0aa6832d0..264b37fb5 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index e264c7af5..07e4bb197 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -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}, } diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index e1efd7ac4..39131fd58 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -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() diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index ef86a5a35..e38d83133 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -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 = { diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 380096af5..4dbedbd19 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -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) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 231e33857..c27949310 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -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 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index e2d46ce2c..ba19c02d1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -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 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 0495d79d8..9d2d2d171 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -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={}) diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index f499654a4..0d6151313 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -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( diff --git a/vllm_omni/model_executor/models/moss_tts/streaming_attention.py b/vllm_omni/model_executor/models/moss_tts/streaming_attention.py new file mode 100644 index 000000000..1acce7710 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/streaming_attention.py @@ -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) diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 90197f6e6..e74931a20 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -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 diff --git a/vllm_omni/platforms/xpu/patch.py b/vllm_omni/platforms/xpu/patch.py new file mode 100644 index 000000000..841cf2503 --- /dev/null +++ b/vllm_omni/platforms/xpu/patch.py @@ -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() diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 3d2cd2d62..df4e7fd55 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -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" diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 9a30209ee..6625effe8 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index eca15c16e..80d55020f 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -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,