From 3bc3f1a7d0a139df23e19566e502276e5d998e49 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 11:03:14 -0700 Subject: [PATCH] [Model] Optimize MOSS Local 1.5 MRV2 serving and progressive chunks (#8213) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin Co-authored-by: Yueqian Lin --- docs/configuration/environment_variables.md | 6 +- docs/configuration/stage_configs.md | 47 ++ recipes/OpenMOSS/MOSS-TTS.md | 297 ++++++++++ tests/config/test_config_factory.py | 68 ++- tests/config/test_environment_variables.py | 4 +- .../sched/test_generation_batch_coalescing.py | 253 ++++++++ .../sched/test_generation_scheduler_init.py | 102 ++++ .../test_omni_sched_deferred_free_fence.py | 5 + tests/engine/test_stage_engine_core_proc.py | 66 +++ .../openai_api/test_moss_local_prompt.py | 85 +++ .../openai_api/test_tts_adapter.py | 39 ++ tests/entrypoints/test_pd_disaggregation.py | 3 + tests/metrics/test_emit_calls.py | 4 +- tests/metrics/test_stats.py | 39 ++ .../moss_tts/test_codec_async_output.py | 338 +++++++++++ .../moss_tts/test_codec_ring_headroom.py | 146 +++++ .../moss_tts/test_first_chunk_fast_path.py | 547 ++++++++++++++++++ .../moss_tts/test_local_code_rows_snapshot.py | 32 + .../moss_tts/test_local_compiled_sampler.py | 37 ++ .../moss_tts/test_local_determined_tokens.py | 124 ++++ .../models/moss_tts/test_local_eager_mtp.py | 166 ++++++ .../models/moss_tts/test_local_model_state.py | 286 +++++++++ .../moss_tts/test_local_prefill_batch.py | 86 +++ .../moss_tts/test_local_short_attention.py | 82 +++ .../models/moss_tts/test_reference_encoder.py | 233 ++++++++ .../moss_tts/test_reference_encoder_graphs.py | 320 ++++++++++ .../moss_tts/test_shared_reference_encoder.py | 339 +++++++++++ .../models/moss_tts/test_slot_attention.py | 301 ++++++++++ ...t_streaming_decoder_wrapper_concurrency.py | 86 +++ .../moss_tts/test_streaming_terminal_batch.py | 32 +- .../test_moss_raw_async_chunk.py | 80 +++ .../test_moss_raw_empty_terminal.py | 93 +++ .../test_moss_tts_async_chunk.py | 137 +++++ tests/worker/test_omni_connector_mixin.py | 106 +++- tests/worker_v2/test_first_audio_sender.py | 16 + tests/worker_v2/test_omni_ar_model_runner.py | 31 + tests/worker_v2/test_omni_model_state.py | 55 ++ .../config/environment_variable_inventory.py | 18 + .../core/sched/omni_generation_scheduler.py | 144 +++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 14 +- .../core/sched/omni_scheduling_coordinator.py | 7 +- vllm_omni/deploy/moss_tts_local.yaml | 3 + vllm_omni/deploy/moss_tts_local_mrv2.yaml | 30 + .../moss_tts_local_mrv2_high_concurrency.yaml | 53 ++ .../moss_tts_local_mrv2_low_latency.yaml | 66 +++ .../omni_connector_payload_transport.py | 19 +- vllm_omni/engine/stage_engine_core_proc.py | 24 +- vllm_omni/entrypoints/omni_base.py | 2 +- .../entrypoints/openai/tts_adapters/base.py | 13 +- .../openai/tts_adapters/moss_tts.py | 105 +++- vllm_omni/metrics/stats.py | 99 ++-- .../models/moss_tts/audio_tokenizer_v2.py | 131 ++++- .../cuda_graph_streaming_decoder_wrapper.py | 35 +- .../models/moss_tts/first_chunk_fast_path.py | 333 +++++++++++ .../models/moss_tts/local_model_state.py | 447 ++++++++++++++ .../models/moss_tts/local_short_attention.py | 100 ++++ .../moss_tts/modeling_moss_tts_codec.py | 477 +++++++++++++-- .../moss_tts/modeling_moss_tts_local_depth.py | 31 +- .../moss_tts/modeling_moss_tts_talker.py | 26 + .../models/moss_tts/pipeline.py | 2 + .../models/moss_tts/reference_encoder.py | 448 ++++++++++++-- .../moss_tts/reference_encoder_graphs.py | 399 +++++++++++++ .../moss_tts/shared_reference_encoder.py | 334 +++++++++++ .../models/moss_tts/slot_attention.py | 475 +++++++++++++++ vllm_omni/model_executor/output_snapshot.py | 125 ++++ .../stage_input_processors/moss_tts.py | 72 +-- vllm_omni/worker_v2/first_audio_sender.py | 19 +- vllm_omni/worker_v2/model_states/__init__.py | 6 + vllm_omni/worker_v2/omni_ar_model_runner.py | 24 + vllm_omni/worker_v2/output_snapshot.py | 94 +-- 70 files changed, 8519 insertions(+), 347 deletions(-) create mode 100644 tests/core/sched/test_generation_batch_coalescing.py create mode 100644 tests/core/sched/test_generation_scheduler_init.py create mode 100644 tests/entrypoints/openai_api/test_moss_local_prompt.py create mode 100644 tests/model_executor/models/moss_tts/test_codec_async_output.py create mode 100644 tests/model_executor/models/moss_tts/test_codec_ring_headroom.py create mode 100644 tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py create mode 100644 tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py create mode 100644 tests/model_executor/models/moss_tts/test_local_compiled_sampler.py create mode 100644 tests/model_executor/models/moss_tts/test_local_determined_tokens.py create mode 100644 tests/model_executor/models/moss_tts/test_local_eager_mtp.py create mode 100644 tests/model_executor/models/moss_tts/test_local_model_state.py create mode 100644 tests/model_executor/models/moss_tts/test_local_prefill_batch.py create mode 100644 tests/model_executor/models/moss_tts/test_local_short_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py create mode 100644 tests/model_executor/models/moss_tts/test_shared_reference_encoder.py create mode 100644 tests/model_executor/models/moss_tts/test_slot_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml create mode 100644 vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py create mode 100644 vllm_omni/model_executor/models/moss_tts/local_model_state.py create mode 100644 vllm_omni/model_executor/models/moss_tts/local_short_attention.py create mode 100644 vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py create mode 100644 vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py create mode 100644 vllm_omni/model_executor/models/moss_tts/slot_attention.py create mode 100644 vllm_omni/model_executor/output_snapshot.py diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 9f0e1cc91..a872eaf3c 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -214,7 +214,7 @@ collection, examples, bug reports, or logs. ## Model-specific variables -The inventory includes 66 variables read by a single model or pipeline family. They are +The inventory includes 93 variables read by a single model or pipeline family. They are not listed as public usage options here because doing so would turn implementation escape hatches into an accidental compatibility contract. @@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the | Disposition | Count | Required outcome | | --- | ---: | --- | -| Promote | 38 | Move a stable setting into typed stage or model configuration. | +| Promote | 64 | Move a stable setting into typed stage or model configuration. | | Request scope | 5 | Move request-varying behavior into a declared request-option schema. | | External | 0 | Retain only when a supported third-party library owns the contract. | -| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. | +| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. | | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. | The disposition is a migration target, not a statement that the existing diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 2cf7327dd..214eb3187 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -514,3 +514,50 @@ for that comparison; this switch only controls automatic MPS management. MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on the intended GPU, and warm the complete pipeline before measuring performance. + +## MOSS-TTS Local 1.5 with Model Runner V2 + +`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using +the shared runtime introduced for Qwen3-TTS: + +```bash +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml +``` + +This profile inherits the batching, codec graph buckets and 1-frame/15-frame +chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local +Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp` +capabilities while retaining its V1 `talker_mtp` implementation, sampling +defaults and explicit request-seed handling. Codec chunks use the native data +plane; the internal Talker retains its final-only orchestrator output policy. + +On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration. +The codec retains its own CUDA Graph capture/replay, including batch buckets +through 64, but disables Inductor compilation (`compilation_config.mode: 0`) +of the stateful decoder to avoid lengthy compilation at startup. This is +independent of `enforce_eager`; the profile keeps `enforce_eager: false`. +It explicitly selects `cudagraph_mode: FULL`, which works without Inductor; +the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear +the codec's capture buckets when compilation is disabled. + +For sustained C128 serving on a large-memory CUDA GPU, select the +separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128 +stream slots per stage, a 32 GiB Talker KV budget, and the codec's +`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs. +The bounded KV budget leaves room for codec state and graphs; it does not +guarantee that 128 maximum-length prompts fit simultaneously. See the +[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention) +for activation, backend comparisons, memory requirements and benchmark commands. + +Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1. +NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable +MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio; +set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and +`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM. + +Event-driven orchestration remains independently selectable with +`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical +when comparing these modes. Model-runner selection does not change the +orchestration default or enable experimental reference encoding, chunk ramps, +generation-output draining or MPS. diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index 8d37039d8..c9ec7670b 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses audio-reference generation. Reference transcripts must match the reference audio; they are not style instructions. +For the optional CUDA MRV2 runner, see the +[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2). +The default Local deployment continues to use V1. + ## References - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/) @@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \ - No `ref_audio` required or accepted for MOSS-SoundEffect. - Input field maps to the `ambient_sound` parameter in the upstream processor. - Rate: ~12.5 tokens per second; longer descriptions produce longer audio. + +## Local 1.5 MRV2 and slot attention + +`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with +an explicit deploy profile. The default Local profile continues to use V1. +Both profiles below preserve 1-frame initial and 15-frame steady codec chunks +and the model's sampling defaults. + +```bash +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml +``` + +This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned +CUDA graphs without compiling the codec with Inductor. + +For sustained high concurrency on a large-memory CUDA GPU, use: + +```bash +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +The high-concurrency profile places both stages on one GPU, sets each stage's +capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses +mixed FULL CUDA graphs with both its token budget and capture limit set to +512, so new prompt chunks remain within graph coverage. It also enables +`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and +continuation control stay in GPU request slots, gathered through MRV2's +batch mapping. Prefill conditioning and explicit per-request sampling seeds +retain their original behavior. Published audio codes own their storage, so +slot reuse cannot overwrite an in-flight output. Set the override to `false` +to compare with the generic Omni model state. The profile selects +`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels +disabled, and codec CUDA graph buckets through 128. The configuration was +validated on one H200. The earlier capacity-256 variant used about 110 GiB of +sampled peak GPU memory, including loading and warmup. Reduce the stage capacities, +graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec +compilation can take several minutes; subsequent starts can reuse the AOT +cache. Graph capture is still performed at startup. + +The high-concurrency profile also coalesces ready codec streams with +`connectors.shm.extra.generation_min_batch_size: 32` and +`generation_max_wait_ms: 12`. The target controls when to stop waiting; it +neither caps execution at 32 nor guarantees every codec group has 32 rows. +First and final chunks share the bounded window at high concurrency. When +fewer than 32 requests remain, dispatch is immediate. Pending outputs are +retired before waiting, with the original deadline preserved, so waiting does +not prevent the previous batch from releasing its in-flight state. Cancellation +and input notifications keep their ordinary scheduler bookkeeping. This option +requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to +`0` to disable coalescing; other profiles keep the immediate-dispatch default. + +The high-concurrency profile enables two additional Local model-state +optimizations under the CUDA stage-0 `hf_overrides`, alongside +`mrv2_gpu_slot_state: true`: + +- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines + CPU reference-code slices and their destination positions into one pinned + upload. Host staging is reused only after the upload event completes; + GPU-resident references keep the original path. +- `mrv2_direct_tokens: true` returns the already determined text token after + the Local audio/stop decision. It avoids constructing and sampling the + full text vocabulary. Audio-code and binary-stop sampling are unchanged. + Requests requiring distribution metadata or token constraints, and + unsupported execution modes, retain the normal sampler. + +Other profiles leave both flags disabled. Set both flags explicitly when +comparing their combined effect. Deployment +inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also +retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as +throughput: faster prompt preparation can change competition between the +Talker and codec sharing the GPU. + +For the single-H200 capacity-128 throughput baseline, keep both switches on. +In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch +rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct +tokens alone had no reliable gain. The combined path had 685 ms mean first +audio versus 679 ms across six control rounds, with lower mean completion time. +A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus +3.78% for the control. This sample does not establish speaker similarity or +exclude small quality changes. At concurrency 256, the combined path had +raised mean first-audio latency by about 0.33 s, so other profiles leave the +switches off until their latency and quality are checked. + +The optional connector setting `generation_coalescing_policy: idle_wait` +waits for an inbox notification when no stream is runnable, all receivers +are parked and registered, and no output needs retirement. The first ready +arrival starts a fresh coalescing window, preserving the batch budget. +Control messages, including cancellations, can wait up to two windows +instead of one (24 ms at the 12 ms setting). The default policy remains +`fixed`; neither policy changes the global orchestration default. + +The codec backends differ in state access: + +- `sdpa` uses PyTorch attention after gathering and updating the ring cache. +- `triton` replaces the attention calculation, retaining the ring-cache + gather/copy and explicit mask construction. +- `triton_slot` writes surviving K/V directly into request slots, attends + directly to the ring, and advances the active slot offsets in three + ordered kernels. Graph-padding rows do not advance persistent state. + +The slot path preserves the existing chunk-complete ring semantics, including +retaining the final cache-capacity tokens when a chunk exceeds the ring. +It does not introduce another cache owner or change request-slot lifetime. +Zero-length padding rows skip the attention computation and emit zeros. + +The high-concurrency profile includes intermediate codec batch buckets 6, 12 +and 24. To compare with power-of-two buckets, change the stage-1 +`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`. +Keep the maximum bucket equal to the state capacity to retain terminal-tail +coalescing. Smaller buckets reduce padding work but require more graphs; +measure complete serving runs before selecting them for a deployment. + +For a codec output-path comparison, set +`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous +codec output path; `1` selects GPU output snapshots on a private codec +stream. Both configurations use MRV2 and keep streaming audio responses. +The option changes input ownership, metadata staging and snapshot handling +as well as transfer scheduling, so its timing difference is not just D2H +copy time. + +Those chunk-complete semantics truncate the causal window: every ring holds +exactly `context` entries and a whole chunk is written before attention runs, +so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the +tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/ +400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so +its first 80 tokens attend to nothing and no chunk sees the previous one there. +Padding a terminal tail to 15 frames therefore decodes it differently from an +exact-length decode; the difference is deterministic (identical in fp32 and +fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each +ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming +then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails +equal exact tails. It increases state/graph memory and capture time; historical C256 experiments +observed about a 2% throughput cost. It is off by default and its perceptual +quality benefit has not been established. Terminal tails are padded +into the regular 15-frame graph bucket in either mode. +The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths +retain their existing attention implementation. Event-driven orchestration +remains independently selectable and its default is unchanged. + +### Optional progressive chunks + +Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks +before steady decoding, for example: + +```yaml +connectors: + shm: + extra: + codec_chunk_ramp: [1, 4, 15] +``` + +The ramp's first entry overrides `initial_codec_chunk_frames`. After the last +entry, chunks use `codec_chunk_frames`. Each request advances independently; +the final partial chunk is flushed and an empty terminal carries no codec +tokens. The codec captures every ramp length, sizes the maximum execution +step for the largest entry, and uses the ramp's first entry for its dedicated +first-chunk graph. Additional shapes increase startup and graph memory. + +Ramps are opt-in. They can reduce early playback underrun while increasing +codec calls, first-audio latency or total generation time. The measured C128 +throughput presets retain 1→15; their published figures do not establish an +additional benefit from enabling ramps. The processor and graph-shape tests +run without GPU or model weights: + +```bash +CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \ + tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \ + -m 'core_model and cpu' --run-level core_model +``` + +### Low-latency and reference-encoding options + +`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and +uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP. +Its codec first-chunk fast path uses dedicated graphs and a stream handoff +before the regular decoder reuses the request slot. A gate limits contention +with regular codec work. Slot waits are bounded at 30 seconds. A failed or +stalled decode raises an error rather than replaying a partially advanced +slot; a failed output enqueue retains owned PCM for regular delivery. +Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch +policy; it does not change the client concurrency or stage capacities. + +Reference encoding runs in the API layer, independently of MRV2. Enable +reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default +eager encoder is retained when the option is absent. The graph path uses the +loaded tokenizer's encoder/quantizer, length buckets and optional compilation +(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed +attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention. +Compilation and attention changes need not produce bit-identical codes. + +For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared +reference-code storage and, by default, a single encoder host with four +workers. Use a dedicated directory per service, checkpoint and encoding +configuration. Workers have their own graph resources; all graph captures +complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains +separate encoders while sharing codes. More workers/graphs consume memory; +they do not imply more GPUs. Example after selecting an available GPU and +configuring any operator-managed MPS daemon: + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \ +VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \ +VLLM_OMNI_MOSS_REF_GRAPHS=1 \ +VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \ +VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \ +VLLM_OMNI_MOSS_REF_INFLIGHT=4 \ +VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \ +VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --api-server-count 4 \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +This example is a configurable serving profile, not the complete launch +configuration of a historical benchmark. In particular, the low-latency YAML +uses synchronous stage-1 scheduling; the retained source46 experiment used +asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder +environment settings. Do not infer a throughput result from the YAML alone. + +### Cold versus hot reference measurements + +A cold reference misses the encoded-reference cache: the request pays for +reference parsing/preparation and GPU encoding before speech generation. +A hot reference reuses reference codes but still synthesizes new speech; +generated audio is not cached. Pre-registering a reference moves encoding +into registration and can make subsequent synthesis hot, but registration +latency is still part of the first-use cost. + +Model startup/compilation is separate from reference coldness. EN1088 contains +repeated references, so its first pass is not uniformly cold. For an all-unique +cold pass, use its 666 first-occurrence references, keep order, warm execution +with disjoint references, then repeat the same 666 requests. A second pass +through independent API-local caches is not guaranteed hot; verify cache +coverage before comparing it with a shared-cache result. Report requests/s +alongside audio-s/s and retain failed, empty, long and near-silent outputs. + +### Reproduce the serving benchmark + +Use the complete Seed-TTS English test set, including reference audio and text. +Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088 +entries. Run the native benchmark after the server is ready: + +```bash +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000 +export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2 +export SEED_TTS_WER_EVAL=0 + +for concurrency in 128; do + for phase in warm r1 r2; do + vllm bench serve --omni \ + --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \ + --backend openai-audio-speech --endpoint /v1/audio/speech \ + --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \ + --seed-tts-locale en --disable-shuffle --num-prompts 1088 \ + --num-warmups 0 --ready-check-timeout-sec 0 \ + --output-len 256 --max-concurrency "$concurrency" \ + --request-rate inf --seed 42 \ + --extra-body '{"task_type":"Base","max_new_tokens":256}' \ + --save-result --save-detailed --result-dir results/moss-local \ + --result-filename "c${concurrency}-${phase}.json" + done +done +``` + +Discard the complete `warm` pass and combine measured runs as total generated +audio seconds divided by total benchmark duration. Require 1088 successes, +zero failures and 1088 nonempty-audio metric samples in each run. The output +cap matches the benchmark protocol; success alone does not establish speech +quality or that every sentence ended before the cap. Dataset seed 42 does +not fix an independent sampling seed for every request. + +For an attention-only comparison, copy the high-concurrency YAML beside the +original to preserve relative `base_config` resolution. Change its +codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`, +and compare `codec_attention_backend: triton` against `triton_slot`. Keep all +other settings, warmup and client concurrency identical. Restart the server +between configurations and use separate result directories. Comparing the +compiled slot profile to uncompiled Triton includes both changes. + +Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages, +and no model weights: + +```bash +python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \ + tests/model_executor/models/moss_tts/test_streaming_attention.py \ + -m 'core_model and cuda' --run-level=core_model +``` diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index db97f990d..64f7cd995 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2928,9 +2928,17 @@ class TestPlatformOverrides: assert deploy.stages[0].engine_extras["block_size"] == 128 - @pytest.mark.parametrize("platform", ["cuda", "npu"]) - def test_recommended_native_runner_platform_defaults(self, platform): - filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts" + @pytest.mark.parametrize( + "filename,pipeline_key", + [ + ("qwen3_tts_mrv2.yaml", "qwen3_tts"), + ("moss_tts_local_mrv2.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"), + ], + ) + @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) + def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key): deploy = load_deploy_config(Path(get_deploy_config_path(filename))) deploy = _apply_platform_overrides(deploy, platform=platform) expect_v2 = platform == "cuda" @@ -2946,6 +2954,60 @@ class TestPlatformOverrides: # v2 only engages the native plane on stages declaring support. assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages) + def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self): + pipeline = resolve_pipeline_config("moss_tts_local") + base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name)) + candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml")) + base = load_deploy_config(base_path) + candidate = load_deploy_config(candidate_path) + assert base.model_runner == "v1" + assert candidate.stages == base.stages + assert candidate.connectors == base.connectors + assert candidate.async_chunk == base.async_chunk + cuda = _apply_platform_overrides(candidate, platform="cuda") + stages = merge_pipeline_deploy(pipeline, cuda) + assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512 + codec_args = stages[1].yaml_engine_args + assert codec_args["max_num_seqs"] == 64 + assert codec_args["enforce_eager"] is False + assert codec_args["compilation_config"] == { + "mode": 0, + "cudagraph_mode": "FULL", + "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64], + "cudagraph_num_of_warmups": 1, + } + for platform in ("npu", "xpu", "rocm", "musa"): + # Platform resolution mutates the deployment; each server loads + # a fresh profile before applying its own platform overrides. + fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform) + original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform) + assert fallback.stages == original.stages + for variant in ("moss_tts_delay", "moss_tts_realtime"): + assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages) + + def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self): + pipeline = resolve_pipeline_config("moss_tts_local") + path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml")) + deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda") + talker, codec = merge_pipeline_deploy(pipeline, deploy) + assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128 + assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True + assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True + assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True + assert talker.yaml_engine_args["max_num_batched_tokens"] == 512 + assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3 + assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot" + compilation = codec.yaml_engine_args["compilation_config"] + assert compilation["mode"] == 3 and compilation["backend"] == "inductor" + assert compilation["cudagraph_mode"] == "FULL" + assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False} + base_path = Path(get_deploy_config_path("moss_tts_local.yaml")) + for platform in ("npu", "xpu", "rocm", "musa"): + actual = _apply_platform_overrides(load_deploy_config(path), platform=platform) + expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform) + assert actual.stages == expected.stages + @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)]) def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native): from unittest.mock import Mock diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index eb3db91ec..58e2a31f1 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 77, + EnvironmentVariableCategory.MODEL_SPECIFIC: 95, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 7, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 48, + ModelEnvironmentVariableDisposition.PROMOTE: 66, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 17, diff --git a/tests/core/sched/test_generation_batch_coalescing.py b/tests/core/sched/test_generation_batch_coalescing.py new file mode 100644 index 000000000..489e6a819 --- /dev/null +++ b/tests/core/sched/test_generation_batch_coalescing.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bounded readiness coalescing must not allocate or consume request state.""" + +import queue +from types import SimpleNamespace + +import pytest +from vllm.v1.core.sched.interface import PauseState +from vllm.v1.request import RequestStatus + +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler +from vllm_omni.outputs import OmniConnectorOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch): + events = [notification("0"), notification("1", "2", "3")] + s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])]) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.006) + assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012]) + + +@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"]) +def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending): + s, clock = scheduler(monkeypatch) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + if pending == "in_flight": + s.requests["0"].num_in_flight_tokens = 1 + elif pending == "admission": + s.requests["0"].status = RequestStatus.WAITING + elif pending == "registration": + s.input_coordinator.pending_chunk_registrations = [object()] + elif pending == "terminal": + s._pending_data_plane_terminal_req_ids = {"0"} + else: + s.requests.clear() + assert s._drain_omni_connector_outputs() == [] + assert clock[0] == 10 + assert not s._omni_connector_output_inbox.waits + + +def test_idle_wait_without_arrivals_is_bounded(monkeypatch): + s, clock = scheduler(monkeypatch) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + assert s._drain_omni_connector_outputs() == [] + assert clock[0] == pytest.approx(10.012) + + +def notification(*ids, terminal=()): + return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal)) + + +class Inbox: + def __init__(self, clock, initial=(), delayed=()): + self.clock = clock + self.initial = list(initial) + self.delayed = list(delayed) + self.waits = [] + + def get_nowait(self): + if self.initial: + return self.initial.pop(0) + raise queue.Empty + + def get(self, timeout): + self.waits.append(timeout) + if not self.delayed or self.delayed[0][0] > timeout: + self.clock[0] += timeout + raise queue.Empty + delay, value = self.delayed.pop(0) + self.clock[0] += delay + return value + + +def scheduler(monkeypatch, initial=(), delayed=()): + clock = [10.0] + monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0]) + s = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + s._generation_max_wait_s = 0.012 + s._generation_min_batch_size = 4 + s._pause_state = PauseState.UNPAUSED + s.requests = { + str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8) + } + s.running = list(s.requests.values()) + s.input_coordinator = SimpleNamespace( + requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set() + ) + s._latest_omni_connector_output = None + s._omni_connector_output_inbox = Inbox(clock, initial, delayed) + return s, clock + + +def test_arrivals_form_target_batch_in_order(monkeypatch): + events = [notification("0"), notification("1", "2"), notification("3")] + s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])]) + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.005) + assert not s.input_coordinator.requests_with_ready_chunks + assert all(r.num_in_flight_tokens == 0 for r in s.requests.values()) + + +def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch): + events = [notification("0", "gone", "1", "2"), notification("0")] + s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])]) + s.requests["1"].is_finished = lambda: True + s.requests["2"].is_finished = lambda: True + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.012) + assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009]) + + +@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"]) +def test_immediate_dispatch(monkeypatch, reason): + event = notification("0") + s, clock = scheduler(monkeypatch, [event]) + if reason == "enough": + event.chunk_ready_req_ids.update(["1", "2", "3"]) + elif reason == "empty": + event.chunk_ready_req_ids.clear() + elif reason == "low": + s.requests = {"0": s.requests["0"]} + elif reason == "paused": + s._pause_state = PauseState.PAUSED_ALL + elif reason == "off": + s._generation_max_wait_s = 0 + assert s._drain_omni_connector_outputs() == [event] + assert not s._omni_connector_output_inbox.waits + assert clock[0] == 10 + + +def test_terminal_arrival_keeps_original_deadline(monkeypatch): + events = [notification("0"), notification("1", terminal=["1"])] + s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])]) + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.012) + + +def test_pending_ready_and_legacy_notification_preserved(monkeypatch): + s, _ = scheduler(monkeypatch) + s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"]) + last = notification("3") + s._latest_omni_connector_output = last + assert s._drain_omni_connector_outputs() == [last] + assert s._latest_omni_connector_output is None + assert not s._omni_connector_output_inbox.waits + + +def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch): + s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))]) + assert len(s._drain_omni_connector_outputs()) == 2 + assert clock[0] == pytest.approx(10.012) + s._omni_connector_output_inbox.initial.append(notification("3")) + assert len(s._drain_omni_connector_outputs()) == 2 + assert clock[0] == pytest.approx(10.024) + + +def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker): + event = notification("0", "cancelled") + event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}} + s, _ = scheduler(monkeypatch, [event]) + s._generation_max_wait_s = 0 + s.waiting = [] + s.input_coordinator._async_chunk = True + s.input_coordinator.update_request_metadata = mocker.Mock() + s.input_coordinator.process_pending_chunks = mocker.Mock() + s._consume_pending_connector_output("generation") + s.input_coordinator.update_request_metadata.assert_called_once_with( + s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation" + ) + s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set()) + + +@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))]) +def test_invalid_config_rejected(monkeypatch, target, wait): + from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler + + def initialize(s, *args, **kwargs): + s.max_num_running_reqs = 256 + s.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + retains_state_across_chunks=True, + stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}}, + ) + ) + + monkeypatch.setattr(VLLMScheduler, "__init__", initialize) + monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None) + with pytest.raises(ValueError): + OmniGenerationScheduler() + + +@pytest.mark.parametrize("terminal", [False, True]) +def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal): + event = notification("0", terminal=["0"] if terminal else []) + s, clock = scheduler(monkeypatch, [event]) + s.running = s.running[1:] + assert s._drain_omni_connector_outputs() == [event] + assert clock[0] == pytest.approx(10.012) + + +def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch): + events = [notification("0"), notification("1")] + s, clock = scheduler(monkeypatch, events[:1]) + s.requests["7"].num_in_flight_tokens = 1 + assert s._drain_omni_connector_outputs() == events[:1] + assert s._generation_defer_batch + assert not s._omni_connector_output_inbox.waits + s.input_coordinator.requests_with_ready_chunks.add("0") + # Simulate EngineCore retiring the previous batch after five milliseconds. + clock[0] += 0.005 + s.requests["7"].num_in_flight_tokens = 0 + s._omni_connector_output_inbox.initial.append(events[1]) + assert s._drain_omni_connector_outputs() == events[1:] + assert not s._generation_defer_batch + assert s._omni_connector_output_inbox.waits == pytest.approx([0.007]) + assert clock[0] == pytest.approx(10.012) + assert s._generation_batch_deadline is None + + +def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch): + s, clock = scheduler(monkeypatch, [notification("0")]) + s.requests["7"].num_in_flight_tokens = 1 + s._generation_batch_deadline = clock[0] - 0.001 + assert len(s._drain_omni_connector_outputs()) == 1 + assert not s._generation_defer_batch + assert not s._omni_connector_output_inbox.waits + + +def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker): + from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler + + request = _chunk_request("pending") + s = _make_generation_scheduler(request) + s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True) + allocate = mocker.spy(s.kv_cache_manager, "allocate_slots") + output = s.schedule() + assert output.total_num_scheduled_tokens == 0 + assert not output.num_scheduled_tokens + allocate.assert_not_called() + assert list(s.waiting) == [request] + assert request.num_computed_tokens == request.num_in_flight_tokens == 0 + assert s.chunk_transfer_adapter.restore_called diff --git a/tests/core/sched/test_generation_scheduler_init.py b/tests/core/sched/test_generation_scheduler_init.py new file mode 100644 index 000000000..ade94af17 --- /dev/null +++ b/tests/core/sched/test_generation_scheduler_init.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise generation constructor policy after actual deploy/platform resolution.""" + +from types import SimpleNamespace + +import pytest + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config.pipeline_registry import resolve_pipeline_config +from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def construct_scheduler(monkeypatch, mocker): + # Only upstream cache allocation and connector construction are omitted. + # Keep the actual Omni I/O initialization, native decision and constructor. + mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter") + + def upstream_init(self, config): + self.vllm_config = config + self.max_num_running_reqs = config.scheduler_config.max_num_seqs + + monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init) + return OmniGenerationScheduler + + +def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128): + return SimpleNamespace( + model_config=SimpleNamespace( + use_v2_model_runner=native, + supports_native_mrv2_data_plane=True, + retains_state_across_chunks=stateful, + async_chunk=True, + stage_id=1, + stage_connector_config={"extra": extras or {}}, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp), + scheduler_config=SimpleNamespace(max_num_seqs=capacity), + ) + + +@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"]) +@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) +def test_moss_profile_generation_constructor_after_platform_resolution(construct_scheduler, mocker, profile, platform): + deploy = _apply_platform_overrides( + load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml")), platform=platform + ) + pipeline = resolve_pipeline_config("moss_tts_local") + codec = merge_pipeline_deploy(pipeline, deploy)[1] + extras = deploy.connectors["shm"]["extra"] + config = _config( + native=codec.yaml_engine_args["use_v2_model_runner"], + stateful=pipeline.stages[1].retains_state_across_chunks, + capacity=codec.yaml_engine_args["max_num_seqs"], + extras=extras, + ) + warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning") + scheduler = construct_scheduler(config) + # The raw extras persist through fallback; the live constructor must cope. + assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12 + if platform == "cuda": + assert scheduler._native_data_plane and scheduler.input_coordinator is not None + assert scheduler._generation_min_batch_size == 32 + assert scheduler._generation_max_wait_s == 0.012 + assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0) + warning.assert_not_called() + if profile == "low_latency": + assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1 + else: + assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None + assert scheduler._generation_min_batch_size == 1 + assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0 + warning.assert_called_once() + + +@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)]) +def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp): + config = _config( + stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12} + ) + with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"): + construct_scheduler(config) + + +@pytest.mark.parametrize( + "extras,message", + [ + ({"generation_min_batch_size": 0}, "generation_min_batch_size"), + ({"generation_min_batch_size": 129}, "generation_min_batch_size"), + ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"), + ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"), + ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"), + ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"), + ], +) +def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message): + with pytest.raises(ValueError, match=message): + construct_scheduler(_config(native=False, extras=extras)) diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index 9088c6bc2..c353ced3f 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -126,6 +126,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request sched.max_num_running_reqs = 8 sched._native_data_plane = False sched._retains_state_across_chunks = False + # Unset MagicMock attributes are truthy. Keep optional scheduling modes + # disabled so this fixture actually schedules its one-shot request. + sched._generation_defer_batch = False + sched._generation_first_chunk_only = False + sched._generation_max_regular_batch = 0 sched._pause_state = PauseState.UNPAUSED sched.requests = {request.request_id: request} sched.running = [request] diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index 217f13e53..19c24b306 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +from typing import Any from unittest.mock import patch import pytest @@ -25,11 +26,15 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder( expected, ): executor = UniProcExecutor.__new__(UniProcExecutor) + executor.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1) + ) executor.driver_worker = mocker.Mock() runner = executor.driver_worker.worker.model_runner runner.model.first_frame_decoder = object() if first_decoder else None runner.model.stream_decoder = object() if stream_decoder else None runner.model.stream_first_audio = stream_first_audio + runner._omni_data_plane = None assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected assert runner.model_state.set_first_audio_sink.call_count == int(expected) @@ -54,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields(): assert current_wave == 3 assert result.external_req_id == "external" assert result.additional_information == {"conditioning": "payload"} + + +def test_codec_uses_generic_first_audio_binding(monkeypatch): + import queue + + import torch + + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + from vllm_omni.engine import stage_engine_core_proc as module + + calls: dict[str, Any] = {} + + def hook(*args): + return False + + def bind(sink): + calls["sink"] = sink + return hook + + plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value)) + model = SimpleNamespace(bind_first_chunk_fast_path=bind) + runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane) + executor = SimpleNamespace( + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)), + driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)), + ) + monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace) + outputs: queue.Queue = queue.Queue() + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)}) + assert module._bind_first_audio_sink(executor, outputs, scheduler) + assert calls["hook"] is hook + calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000)) + client, batch = outputs.get_nowait() + assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output + + +def test_first_audio_binding_preserves_talker_marker(monkeypatch): + import queue + + import torch + + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + from vllm_omni.engine import stage_engine_core_proc as module + + calls: dict[str, Any] = {} + runner = SimpleNamespace( + model=SimpleNamespace(first_frame_decoder=object()), + model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)), + ) + executor = SimpleNamespace( + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)), + driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)), + ) + monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace) + outputs: queue.Queue = queue.Queue() + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)}) + assert module._bind_first_audio_sink(executor, outputs, scheduler) + calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000)) + assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY] + executor.vllm_config.parallel_config.tensor_parallel_size = 2 + assert not module._bind_first_audio_sink(executor, outputs, scheduler) diff --git a/tests/entrypoints/openai_api/test_moss_local_prompt.py b/tests/entrypoints/openai_api/test_moss_local_prompt.py new file mode 100644 index 000000000..8fb4c62a8 --- /dev/null +++ b/tests/entrypoints/openai_api/test_moss_local_prompt.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The Local-v1.5 continuation prompt fast path equals the processor layout.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Processor: + """Segment API of the model's processor with a deterministic tokenizer.""" + + def __init__(self): + self.model_config = SimpleNamespace( + n_vq=3, + im_start_token_id=1, + im_end_token_id=2, + audio_start_token_id=3, + audio_pad_token_id=1024, + audio_assistant_slot_token_id=5, + ) + self.encoded = [] + + def _encode_text(self, text): + self.encoded.append(text) + return [100 + ord(c) % 50 for c in text] + + def _user_prompt_prefix_ids(self): + return [1] + self._encode_text("user\n") + self._encode_text("") + + def _user_prompt_after_reference_ids(self, language, fields): + return self._encode_text(f"lang={fields.get('language', language)}") + + def _assistant_prompt_prefix_ids(self): + return self._encode_text("") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant") + + def reference(self, text, language, codes): + # Mirrors the processor's _build_continuation_codes + _pad (one item). + fields = {} if language is None else {"language": language} + ids = ( + self._user_prompt_prefix_ids() + + self._encode_text("None") + + self._user_prompt_after_reference_ids(language, fields) + + self._encode_text(text) + + self._assistant_prompt_prefix_ids() + + [3] + ) + text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long) + text_rows[:, 0] = torch.tensor(ids) + audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long) + audio_rows[:, 0] = 5 + audio_rows[:, 1:] = codes + unified = torch.cat([text_rows, audio_rows]) + return unified[:, 0].tolist(), unified[:, 1:].contiguous() + + +@pytest.mark.parametrize("language", [None, "English"]) +def test_matches_processor_and_caches_constant_segments(language): + proc = _Processor() + codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32) + kwargs = {"text": "ref words target words"} + if language is not None: + kwargs["language"] = language + ids, audio = _local_continuation_prompt(proc, kwargs, [codes]) + expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long()) + assert ids == expected_ids + assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio) + + proc.encoded.clear() + _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes]) + assert proc.encoded == ["other"] + + +def test_unsupported_layouts_fall_back(): + proc = _Processor() + codes = torch.zeros((2, 3), dtype=torch.long) + assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 92f13462e..f17fdbb65 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker) audio_tokenizer.to.assert_called_once() +@pytest.mark.parametrize( + "variant,graphs_env,prepares", + [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)], +) +def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch): + if graphs_env is None: + monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env) + adapter = _moss_adapter_with_stage_configs([], mocker) + adapter._moss_variant = variant + encoder = mocker.Mock() + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + asyncio.run(adapter.warmup()) + assert encoder.prepare.called is prepares + + +def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch): + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1") + adapter = _moss_adapter_with_stage_configs([], mocker) + encoder = mocker.Mock() + encoder.prepare.side_effect = RuntimeError("capture failed") + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + asyncio.run(adapter.warmup()) + encoder.prepare.assert_called_once_with() + + +def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch): + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1") + from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError + + adapter = _moss_adapter_with_stage_configs([], mocker) + encoder = mocker.Mock() + encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling") + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + with pytest.raises(SharedReferenceEncoderStartupError): + asyncio.run(adapter.warmup()) + + def test_qwen3_tts_metadata(): assert Qwen3TTSAdapter.backend == "ar" assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index e9a3ba54c..bf56da2b3 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch): def build_and_log_summary(self): return "Fake summary" + def log_timing_summary(self): + pass + monkeypatch.setattr( "vllm_omni.entrypoints.omni.OrchestratorAggregator", _FakeOrchestratorAggregator, diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index 02a98fc00..515aaeb8d 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -111,7 +111,7 @@ class TestFailureCounterWiring: obj.request_states["req-1"] = SimpleNamespace( metrics=SimpleNamespace( e2e_done=set(), - build_and_log_summary=lambda: None, + log_timing_summary=lambda: None, ), failure_recorded=False, ) @@ -127,7 +127,7 @@ class TestFailureCounterWiring: obj.request_states["req-1"] = SimpleNamespace( metrics=SimpleNamespace( e2e_done=set(), - build_and_log_summary=lambda: None, + log_timing_summary=lambda: None, ), failure_recorded=False, ) diff --git a/tests/metrics/test_stats.py b/tests/metrics/test_stats.py index 8b18c5dff..f72cd0867 100644 --- a/tests/metrics/test_stats.py +++ b/tests/metrics/test_stats.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 pytest @@ -155,3 +158,39 @@ def test_build_and_log_summary_multiple_requests() -> None: r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2") assert len(r1_stage_entry["stages"]) == 2 assert len(r2_stage_entry["stages"]) == 1 + + +def test_log_timing_summary_skips_tables_unless_debug(monkeypatch) -> None: + from vllm_omni.metrics import stats + + agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0) + agg.stage_first_ts[0] = 0.0 + agg.stage_last_ts[0] = 0.02 + agg.on_stage_metrics( + 0, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=3, + num_tokens_out=3, + stage_gen_time_ms=20.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_finalize_request(0, "r1", req_start_ts=0.0) + tables, lines = [], [] + + def record_table(*args, **kwargs): + tables.append(args) + return "" + + monkeypatch.setattr(stats, "_format_table", record_table) + monkeypatch.setattr(stats.logger, "info", lambda msg, *args: lines.append(msg % args)) + monkeypatch.setattr(stats.logger, "isEnabledFor", lambda level: False) + agg.log_timing_summary() + assert not tables + assert len(lines) == 1 and lines[0].startswith("[OmniTiming] req=r1 total=") diff --git a/tests/model_executor/models/moss_tts/test_codec_async_output.py b/tests/model_executor/models/moss_tts/test_codec_async_output.py new file mode 100644 index 000000000..d855ef9d3 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_async_output.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Waveform ownership across graph replays and asynchronous output copies.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot + +pytestmark = [pytest.mark.core_model] + + +class CausalCodec(nn.Module): + downsample_rate = 1 + + def __init__(self, channels): + super().__init__() + self.weight = nn.Parameter(torch.zeros(())) + self.channels = channels + self.config = SimpleNamespace(codebook_size=1024) + + def initialize_decoder_state_pool(self, capacity, scratch): + self.state = torch.zeros(capacity + scratch, device=self.weight.device) + + def reset_decoder_state_slots(self, slots): + self.state[slots] = 0 + + def decode_streaming_batch(self, codes, lengths, slots, valid_rows): + audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None] + self.state[slots] = audio[:, -1] + channels = torch.arange(1, self.channels + 1, device=audio.device) + return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths) + + +class ReplayGraph: + """A real CUDA graph with deliberately reused output storage.""" + + batch_sizes = [1] + + def __init__(self, codec): + self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long) + self.slots = torch.zeros(1, device="cuda", dtype=torch.long) + self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long) + self.valid = torch.ones(1, device="cuda", dtype=torch.bool) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid) + torch.cuda.current_stream().wait_stream(stream) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid) + codec.state.zero_() + + def decode(self, codes, slots, *, allow_frame_padding): + if codes.shape[1:] != (1, 3): + return None + self.codes.copy_(codes) + self.slots.copy_(slots) + self.graph.replay() + return self.output.audio, self.output.audio_lengths, 1 + + +def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True): + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(n_vq=2, model_type=model_type), + async_chunk=streaming, + use_v2_model_runner=v2, + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + ) + decoder = MossTTSCodecDecoder(vllm_config=config) + decoder._codec = CausalCodec(channels).to(device) + decoder._n_channels = channels + decoder._stream_max_step_frames = 3 + return decoder + + +@pytest.mark.cpu +@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])]) +def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames): + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"), + async_chunk=True, + use_v2_model_runner=True, + stage_connector_config={ + "extra": { + "initial_codec_chunk_frames": 1, + "codec_chunk_frames": 15, + "codec_chunk_ramp": ramp, + } + }, + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + ) + decoder = MossTTSCodecDecoder(vllm_config=config) + assert decoder._streaming_graph_frame_sizes == frames + assert decoder._stream_max_step_frames == 15 + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "v2,streaming,model_type", + [ + (False, True, "moss_tts_local"), + (True, False, "moss_tts_local"), + (True, True, "moss_tts_realtime"), + (True, True, "moss_tts_local"), + ], +) +def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type): + decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type) + assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local") + session = decoder._ensure_stream_session() + slot = session.acquire() + assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu" + + +@pytest.mark.cuda +@pytest.mark.parametrize("v2", [False, True]) +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("graph", [False, True]) +@torch.no_grad() +def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph): + from vllm_omni.outputs import OmniModelRunnerOutput + from vllm_omni.worker_v2.omni_generation_model_runner import ( + OmniGenerationAsyncOutput, + OmniGenerationModelRunner, + _contains_cuda_tensor, + ) + + decoder = make_decoder(v2=v2, channels=channels, device="cuda") + # Real Local-v1.5 weights replace the class's 24 kHz default at load time. + decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32) + session = decoder._ensure_stream_session() + if graph: + session._cudagraph_wrapper = ReplayGraph(decoder._codec) + copy_stream = torch.cuda.Stream() + pending = [] + # Do not materialize any output until later forwards have replayed the + # graph, reset terminal/cancelled state, and reused those state slots. + for round_index in range(4): + value = round_index + 1 + codes = torch.cat( + [ + torch.full((14,), value, device="cuda", dtype=torch.long), + torch.full((6,), value + 10, device="cuda", dtype=torch.long), + ] + ) + info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]] + output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6]) + # The runner may overwrite its input buffer before codec completion. + codes.fill_(-100) + payload = output.multimodal_outputs + assert _contains_cuda_tensor(payload) == v2 + assert isinstance(payload, PackedOutputSnapshot) == v2 + if v2: + # Exactly one waveform slab, regardless of request/chunk count. + assert sum(slab.is_cuda for _, slab in payload._groups) == 1 + result = OmniModelRunnerOutput( + req_ids=["tail", "empty", "live"], + req_id_to_index={"tail": 0, "empty": 1, "live": 2}, + sampled_token_ids=[[], [], []], + ) + pending.append( + OmniGenerationAsyncOutput( + model_runner_output=result, + multimodal_outputs=payload, + num_reqs=3, + main_stream=torch.cuda.current_stream(), + copy_stream=copy_stream, + ) + ) + else: + pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3)) + decoder.on_requests_finished(["live"]) + assert decoder._stream_req_slots == {} + assert not session._leased_slots + + for round_index, item in enumerate(pending): + rows = item.get_output().multimodal_outputs if v2 else item + for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]): + expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32) + if channels > 1: + expected = expected[None] * torch.arange(1, channels + 1)[:, None] + torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0) + assert rows[index]["sr"].item() == 48_000 + + +@pytest.mark.cuda +@torch.no_grad() +def test_owned_input_copy_clamps_entire_ragged_batch(): + from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm + + decoder = make_decoder(device="cuda", channels=2) + codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda") + output = decoder.forward( + codes, + runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")], + seq_token_counts=[6, 0, 2], + ) + codes.fill_(-100) + cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True) + torch.accelerator.synchronize() + expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]]) + torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0) + assert cpu["model_outputs"][1].numel() == 0 + torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0) + decoder.on_requests_finished(["a", "b"]) + + +@pytest.mark.cuda +@torch.no_grad() +def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals(): + decoder = make_decoder(device="cuda") + session = decoder._ensure_stream_session() + # Force reuse every call, with an outstanding transfer on the first call. + # This delay tests host-buffer ownership, never performance or overlap. + session._metadata_ring = session._metadata_ring[:1] + slots = [session.acquire() for _ in range(3)] + pending = [] + for step in range(8): + order = slots[step % 3 :] + slots[: step % 3] + terminal = {order[1]} + torch.cuda._sleep(2_000_000) + ids, tails = session._stage_slot_ids(order, terminal) + pending.append((ids.clone(), tails.clone(), order, list(terminal))) + session.close() + for ids, tails, expected_ids, expected_tails in pending: + assert ids.cpu().tolist() == expected_ids + assert tails.cpu().tolist() == expected_tails + + +@pytest.mark.cuda +@torch.no_grad() +def test_mixed_terminal_reset_does_not_reset_live_slot(): + decoder = make_decoder(device="cuda", channels=1) + session = decoder._ensure_stream_session() + first, second = session.acquire(), session.acquire() + codes = torch.ones(2, 3, device="cuda") + initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first}) + following = session.step({second: codes, first: codes}) + torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0) + torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0) + torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0) + session.close() + + +@pytest.mark.cuda +@pytest.mark.parametrize("explicit_copy_stream", [False, True]) +@torch.no_grad() +def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream): + from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm + + decoder = make_decoder(device="cuda", channels=1) + decoder._codec_stream = torch.cuda.Stream() + with torch.cuda.stream(decoder._codec_stream): + torch.cuda._sleep(20_000_000) + codes = torch.ones(6, device="cuda", dtype=torch.long) + output = decoder.forward( + codes, + runtime_additional_information=[{"request_id": "live"}], + seq_token_counts=[6], + ) + codes.fill_(-100) + payload = output.multimodal_outputs + assert isinstance(payload, PackedOutputSnapshot) + assert payload.producer_event is not None + # No explicit consumer wait here: the shared copier must honor readiness. + copy_stream = torch.cuda.Stream() + with torch.cuda.stream(copy_stream): + cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True) + decoder.on_requests_finished(["live"]) + copy_stream.synchronize() + torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("channels", [1, 2]) +@torch.no_grad() +def test_short_terminal_row_crops_each_channel_into_output_slab(channels): + decoder = make_decoder(device="cuda", channels=channels) + session = decoder._ensure_stream_session() + short, live = session.acquire(), session.acquire() + codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")} + slab = torch.empty(channels * 7, device="cuda") + buffers = { + short: slab[: channels * 2].view(channels, 2), + live: slab[channels * 2 :].view(channels, 5), + } + result = session.step(codes, terminal_slots={short}, output_buffers=buffers) + for slot, frames, value in [(short, 2, 1), (live, 5, 3)]: + expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None] + expected = expected * torch.arange(1, channels + 1)[:, None] + assert result[slot].data_ptr() == buffers[slot].data_ptr() + torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0) + session.close() + + +@pytest.mark.cuda +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"]) +@torch.no_grad() +def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout): + decoder = make_decoder(device="cuda", channels=channels) + session = decoder._ensure_stream_session() + slots = [session.acquire() for _ in range(4)] + slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda") + if layout == "gapped": + buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)} + frames = 6 + elif layout == "adjacent": + buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)} + frames = 6 + else: + # Interleaved with other requests' output: the short views must not + # overwrite the untouched part of another row/channel. + buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)} + frames = 3 + codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)} + result = session.step(codes, output_buffers=buffers) + for row, slot in enumerate(slots): + expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None] + expected = expected * torch.arange(1, channels + 1)[:, None] + assert result[slot].data_ptr() == buffers[slot].data_ptr() + torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0) + if layout == "strided": + assert torch.all(slab[:, :, 3:] == -1) + elif layout == "gapped": + assert torch.all(slab[[2, 5]] == -1) + session.close() diff --git a/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py b/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py new file mode 100644 index 000000000..aa3fb1797 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Chunked streaming decode must match whole-sequence decode. + +``RingKVCache.complete`` writes a whole chunk before attending. With a ring of +exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its +own queries should still see; for T > context it evicts part of the chunk +itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of +headroom per layer so streaming reproduces the non-streaming decoder. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel +from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import ( + MossAudioTokenizerConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TRANSFORMER = { + "module_type": "Transformer", + "d_model": 16, + "num_heads": 2, + "num_layers": 1, + "dim_feedforward": 32, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 1.0, + "conv_layout": True, +} + + +def tiny_codec() -> MossAudioTokenizerModel: + # 64 Hz mono input, 4-sample code frames (16 code frames per second). + # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer + # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12 + # tokens at the second transformer, more than its 8-token ring. + config = MossAudioTokenizerConfig( + sampling_rate=64, + downsample_rate=4, + number_channels=1, + enable_channel_interleave=False, + encoder_kwargs=[ + {"module_type": "PatchedPretransform", "patch_size": 4}, + {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5}, + ], + decoder_kwargs=[ + {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5}, + {"module_type": "PatchedPretransform", "patch_size": 2}, + {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25}, + {"module_type": "PatchedPretransform", "patch_size": 2}, + ], + quantizer_kwargs={ + "input_dim": 8, + "rvq_dim": 8, + "output_dim": 8, + "num_quantizers": 2, + "codebook_size": 16, + "codebook_dim": 4, + "quantizer_type": "rlfq", + }, + ) + torch.manual_seed(0) + model = MossAudioTokenizerModel(config).eval() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.5) + model.quantizer.build_decode_lut(2, dtype=torch.float32) + return model + + +def stream_decode(model, codes, chunk_frames, *, headroom): + n_q, _, total = codes.shape + model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0) + try: + slot = torch.zeros(1, dtype=torch.long) + valid = torch.ones(1, dtype=torch.bool) + parts = [] + for start in range(0, total, chunk_frames): + chunk = codes[:, :, start : start + chunk_frames] + lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long) + parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio) + return torch.cat(parts, dim=-1) + finally: + model.close_decoder_state_pool() + + +@torch.no_grad() +def test_ring_headroom_sizes_each_layer_for_one_chunk(): + model = tiny_codec() + transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")] + assert [m.tokens_per_input_frame for m in transformers] == [1, 2] + assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8] + model.initialize_decoder_state_pool(1, 0, chunk_frames=6) + try: + rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers] + assert rings == [8 + 6, 8 + 12] + finally: + model.close_decoder_state_pool() + model.initialize_decoder_state_pool(1, 0) + try: + rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers] + assert rings == [8, 8] + finally: + model.close_decoder_state_pool() + + +@torch.no_grad() +def test_chunked_streaming_matches_whole_sequence_only_with_headroom(): + model = tiny_codec() + torch.manual_seed(1) + codes = torch.randint(0, 16, (2, 1, 18)) + reference = model._decode_frame(codes).audio + with_headroom = stream_decode(model, codes, 6, headroom=True) + legacy = stream_decode(model, codes, 6, headroom=False) + assert with_headroom.shape == reference.shape == legacy.shape + torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5) + # The legacy context-sized ring lets the 12-token chunk evict keys inside + # its own window, so streaming diverges from whole-sequence decoding. + assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3) + + +@torch.no_grad() +def test_padded_terminal_tail_matches_exact_tail_with_headroom(): + model = tiny_codec() + torch.manual_seed(2) + codes = torch.randint(0, 16, (2, 1, 14)) + history, tail = codes[:, :, :12], codes[:, :, 12:] + padded_tail = torch.nn.functional.pad(tail, (0, 4)) # 2 real frames padded to a 6-frame step + outputs = {} + for name, last in [("exact", tail), ("padded", padded_tail)]: + model.initialize_decoder_state_pool(1, 0, chunk_frames=6) + try: + slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool) + for chunk in (history[:, :, :6], history[:, :, 6:]): + model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid) + audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio + outputs[name] = audio[..., : 2 * model.downsample_rate] + finally: + model.close_decoder_state_pool() + torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5) diff --git a/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py new file mode 100644 index 000000000..85e53e80b --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The first-chunk fast path emits audio out of band; the main path resumes the stream.""" + +import queue +import threading +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session +from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink + +pytestmark = pytest.mark.core_model + +N_VQ, SAMPLES, CHANNELS = 2, 4, 2 + + +class _Session: + def __init__(self, capacity): + self.free = list(range(capacity)) + self.reset_waits = [] + + def acquire(self, *, oldest=False): + return self.free.pop(0 if oldest else -1) if self.free else None + + def order_after_reset(self, slot, stream): + self.reset_waits.append(slot) + + def release(self, slot, **kwargs): + self.free.append(slot) + + +class _Wrapper: + batch_sizes = [1, 2, 4, 8] + + def __init__(self, device): + self.state = torch.zeros(16, device=device) + self.calls = [] + self.gate = threading.Event() + self.gate.set() + + def decode(self, codes, slots): + self.gate.wait() + self.calls.append(int(codes.shape[1])) + base = codes.sum(0)[:, 0].float() + 100 * slots.float() + channel = torch.arange(CHANNELS, device=codes.device).float() + audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device) + self.state[slots] = 1 + return audio.to(torch.bfloat16), None, int(codes.shape[1]) + + +@pytest.fixture +def cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device("cuda") + + +def _fast(device, capacity=4): + session, wrapper = _Session(capacity), _Wrapper(device) + fast = MossFirstChunkFastPath( + session, + wrapper, + n_vq=N_VQ, + frames=1, + codebook_size=1024, + samples_per_frame=SAMPLES, + n_channels=CHANNELS, + sample_rate=torch.tensor(24000, dtype=torch.int32), + device=device, + ) + return fast, session, wrapper + + +class _CallbackSink: + def __init__(self, callback): + self.callback = callback + + def prepare(self, request_ids): + callback = self.callback + + class Delivery: + routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids} + + def __call__(self, ids, rows, sr): + for rid, row in zip(ids, rows, strict=True): + callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr}) + + def fail(self, ids): + pass + + return Delivery() + + +@pytest.mark.cuda +def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda): + fast, session, wrapper = _fast(cuda) + emitted = {} + done = threading.Event() + + def sink(client, req_id, payload): + emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"])) + if len(emitted) == 2: + done.set() + + slots: dict[str, int] = {} + assert not fast.submit("r0", "k0", [1, 2], slots) # unbound: main path owns it + fast.bind(_CallbackSink(sink)) + wrapper.gate.clear() + assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots) + assert fast.submit("r1", "k1", [5, 7], slots) + assert not fast.submit("r0", "k0", [1, 2], slots) # already has a slot + assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots) # not a 1-frame chunk + assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3] + wrapper.gate.set() + assert done.wait(10) + fast.close() + + assert wrapper.calls in ([1, 1], [2]) + for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)): + got_client, wav, sr = emitted[req_id] + assert got_client == client and sr == 24000 and wav.device.type == "cpu" + expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES) + torch.testing.assert_close(wav, expected) + assert sorted(session.reset_waits) == [0, 1] + + # The main path consumes the marker once and is ordered after the decode. + assert fast.take_decoded("k0") and not fast.take_decoded("k0") + fast.order_after(0) + fast.order_after(0) # no pending handoff: no-op + assert wrapper.state[:2].tolist() == [1, 1] + + +@pytest.mark.cuda +def test_order_after_blocks_until_decode_finishes(cuda): + fast, _, wrapper = _fast(cuda) + fast.bind(_CallbackSink(lambda *args: None)) + wrapper.gate.clear() + assert fast.submit("r0", "k0", [1, 2], {}) + finished = threading.Event() + + def wait_for_slot(): + fast.order_after(0) + finished.set() + + waiter = threading.Thread(target=wait_for_slot) + waiter.start() + assert not finished.wait(0.2) + wrapper.gate.set() + assert finished.wait(10) + waiter.join() + fast.close() + + +@pytest.mark.cuda +def test_no_free_slot_leaves_chunk_to_main_path(cuda): + fast, _, _ = _fast(cuda, capacity=0) + fast.bind(_CallbackSink(lambda *args: None)) + slots: dict[str, int] = {} + assert not fast.submit("r0", "k0", [1, 2], slots) + assert slots == {} and not fast.take_decoded("k0") + fast.close() + + +class _FakeFast: + frames = 1 + + def __init__(self, decoded, unsent=None): + self.decoded = set(decoded) + self.unsent = dict(unsent or {}) + self.ordered = [] + self.forgotten = [] + + def take_decoded(self, key): + if key in self.decoded: + self.decoded.remove(key) + return True + return False + + def get_request_slot(self, key, slots): + return slots.get(key) + + def order_after(self, slot): + self.ordered.append(slot) + + def take_unsent_audio(self, key): + return self.unsent.pop(key, None) + + def forget(self, key): + self.forgotten.append(key) + + def gate(self): + self.gated = getattr(self, "gated", 0) + 1 + + +def _decoder(s, fast): + decoder = object.__new__(MossTTSCodecDecoder) + nn.Module.__init__(decoder) + decoder._stream_max_step_frames = 15 + decoder._stream_state_capacity = 8 + decoder._stream_req_slots = {} + decoder._ensure_stream_session = lambda: s + decoder._first_chunk_fast_path = fast + decoder._stream_session = s + decoder._codec_stream = None + return decoder + + +@pytest.mark.cpu +def test_main_path_skips_fast_decoded_chunk_and_resumes_state(): + s = cpu_session() + reference = cpu_session() + fast = _FakeFast({"a"}) + decoder = _decoder(s, fast) + # The fast path leased slot 0 and decoded frame 0 of "a" into it. + slot = s.acquire(oldest=True) + decoder._stream_req_slots["a"] = slot + s.step({slot: torch.full((2, 1), 3.0)}) + ref_slot = reference.acquire() + reference.step({ref_slot: torch.full((2, 1), 3.0)}) + + out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)]) + assert 0 not in out and out[1].tolist() == [[2.0]] + assert fast.ordered == [slot] + + out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)]) + expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot] + torch.testing.assert_close(out[0], expected) + assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot] + + +@pytest.mark.cpu +def test_abort_orders_release_after_fast_decode(): + s = cpu_session() + fast = _FakeFast({"a"}) + decoder = _decoder(s, fast) + slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True) + decoder.on_requests_finished({"a"}) + assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {} + assert slot in s._free_stream_slots + + +@pytest.mark.cpu +def test_main_path_returns_unsent_first_chunk_without_decoding_twice(): + s = cpu_session() + first = torch.full((1, 1), 7.0) + fast = _FakeFast({"a"}, {"a": first}) + decoder = _decoder(s, fast) + slot = s.acquire(oldest=True) + decoder._stream_req_slots["a"] = slot + s.step({slot: torch.full((2, 1), 3.0)}) + + out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)]) + torch.testing.assert_close(out[0], first) + assert fast.unsent == {} and fast.ordered == [slot] + + # A payload containing a later frame must concatenate it after the + # preserved first samples, while advancing state only for the later frame. + first = torch.full((1, 1), 7.0) + fast.decoded.add("b") + fast.unsent["b"] = first + slot_b = s.acquire(oldest=True) + decoder._stream_req_slots["b"] = slot_b + s.step({slot_b: torch.full((2, 1), 3.0)}) + out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)]) + torch.testing.assert_close(out[0][..., :1], first) + assert out[0].shape[-1] == 2 + + +@pytest.mark.cpu +def test_oldest_acquire_reuses_least_recently_released_slot(): + s = cpu_session() + a, b = s.acquire(), s.acquire() + s.release(a) + s.release(b) + assert s.acquire() == b + assert s.acquire(oldest=True) == 7 # never-used slots were freed before any release + assert s._free_stream_slots[-1] == a + + +@pytest.mark.cuda +def test_failed_decode_reports_error_without_replaying_state(cuda): + fast, _, wrapper = _fast(cuda) + + def fail(codes, slots): + raise RuntimeError("boom") + + wrapper.decode = fail + fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected"))) + slots: dict[str, int] = {} + assert fast.submit("r0", "k0", [1, 2], slots) + with pytest.raises(RuntimeError, match="first-chunk decode failed"): + fast.order_after(slots["k0"]) + assert fast.take_decoded("k0") + fast.close() + + +@pytest.mark.cuda +def test_failed_output_handoff_retains_decoded_samples(cuda): + fast, _, _ = _fast(cuda) + + def reject(*args): + raise RuntimeError("output queue closed") + + fast.bind(_CallbackSink(reject)) + slots: dict[str, int] = {} + assert fast.submit("r0", "k0", [1, 2], slots) + fast.order_after(slots["k0"]) + assert fast.take_decoded("k0") + wav = fast.take_unsent_audio("k0") + assert wav is not None and wav.shape == (CHANNELS, SAMPLES) + assert fast.take_unsent_audio("k0") is None + fast.close() + + +@pytest.mark.cuda +def test_gate_orders_main_stream_after_inflight_fast_decode(cuda): + fast, _, _ = _fast(cuda) + side = torch.cuda.Stream() + with torch.cuda.stream(side): + torch.cuda._sleep(200_000_000) + pending = torch.cuda.Event() + pending.record(side) + fast._inflight = pending + fast.gate() # disabled: no wait + unordered = torch.cuda.Event() + unordered.record() + fast._gate_main = True + fast.gate() + ordered = torch.cuda.Event() + ordered.record() + unordered.synchronize() + assert not pending.query() and not ordered.query() + ordered.synchronize() + assert pending.query() + + +def _cpu_fast(monkeypatch, *, scheduler=None): + # Replace only CUDA allocations, retaining the real lifecycle and worker. + import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module + + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor) + monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None) + monkeypatch.setattr(torch.cuda, "Event", lambda: None) + monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None)) + session = _Session(2) + fast = MossFirstChunkFastPath( + session, + SimpleNamespace(batch_sizes=[1, 2]), + n_vq=2, + frames=1, + codebook_size=1024, + samples_per_frame=4, + n_channels=1, + sample_rate=torch.tensor(24000), + device=torch.device("cpu:0"), + handoff_timeout_s=0.01, + ) + outputs: queue.Queue = queue.Queue() + scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False) + return fast, session, sink, outputs + + +def _cancel_during_admission(fast, session, scheduler, sink): + entered, release = threading.Event(), threading.Event() + + class PausedSink: + def prepare(self, ids): + delivery = sink.prepare(ids) + if delivery.routes: + entered.set() + assert release.wait(5) + return delivery + + fast.bind(PausedSink()) + decoder = _decoder(session, fast) + errors = [] + finished = threading.Event() + finishing = threading.Event() + + def submit(): + try: + assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots) + except BaseException as error: + errors.append(error) + + def cancel(): + try: + finishing.set() + decoder.on_requests_finished({"r"}) + finished.set() + except BaseException as error: + errors.append(error) + + producer = threading.Thread(target=submit) + cleaner = threading.Thread(target=cancel) + producer.start() + try: + assert entered.wait(5) + assert decoder._stream_req_slots == {} + # Scheduler cancellation retires the route before the runner hook. + scheduler.requests.pop("r") + cleaner.start() + assert finishing.wait(5) + assert not finished.wait(0.1) + release.set() + producer.join(5) + cleaner.join(5) + assert not producer.is_alive() and not cleaner.is_alive() + assert not errors, errors + assert finished.is_set() + assert decoder._stream_req_slots == {} + assert fast._decoded == {} and fast._handoffs == {} + assert sorted(session.free) == [0, 1] + assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots) + + # A fresh request can use and return the same capacity immediately. + scheduler.requests["reuse"] = SimpleNamespace(client_index=3) + release.set() + assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots) + scheduler.requests.pop("reuse") + decoder.on_requests_finished({"reuse"}) + assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1] + finally: + release.set() + producer.join(5) + if cleaner.ident is not None: + cleaner.join(5) + fast.close() + + +@pytest.mark.cpu +def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch): + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler) + + # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode. + def decode(jobs): + for job in jobs: + job.handoff.done.set() + + monkeypatch.setattr(fast, "_decode", decode) + fast._handoff_timeout_s = 5.0 + _cancel_during_admission(fast, session, scheduler, sink) + + +@pytest.mark.cuda +def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda): + fast, session, _ = _fast(cuda, capacity=2) + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False) + _cancel_during_admission(fast, session, scheduler, sink) + + +@pytest.mark.cpu +def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch): + fast, _, _, _ = _cpu_fast(monkeypatch) + handoff = _SlotHandoff() + fast._handoffs[0] = handoff + for _ in range(2): + with pytest.raises(TimeoutError, match="slot 0"): + fast.order_after(0) + assert fast._handoffs[0] is handoff + handoff.done.set() + fast.order_after(0) + assert fast._handoffs == {} + + +@pytest.mark.cpu +def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch): + fast, session, sink, _ = _cpu_fast(monkeypatch) + fast.bind(sink) + fast.close() + slots: dict[str, int] = {} + assert not fast.submit("r", "external", [1, 2], slots) + assert slots == {} and session.free == [0, 1] + with pytest.raises(RuntimeError, match="closed"): + fast.bind(sink) + + +@pytest.mark.cpu +def test_missing_route_does_not_claim_slot(monkeypatch): + fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={})) + fast.bind(sink) + try: + assert not fast.submit("gone", "external", [1, 2], {}) + assert session.free == [0, 1] and outputs.empty() + finally: + fast.close() + + +@pytest.mark.cpu +def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch): + from vllm.v1.engine import FinishReason + + fast, _, sink, outputs = _cpu_fast(monkeypatch) + entered, release = threading.Event(), threading.Event() + + def fail(jobs): + entered.set() + assert release.wait(2) + raise RuntimeError("partial replay failed") + + fast._decode = fail + fast.bind(sink) + slots: dict[str, int] = {} + try: + assert fast.submit("r", "a", [1, 2], slots) + assert entered.wait(2) + assert fast.submit("r", "b", [3, 4], slots) + release.set() + fast._thread.join(2) + assert not fast._thread.is_alive() + for slot in slots.values(): + with pytest.raises(RuntimeError, match="decode failed"): + fast.order_after(slot) + assert all(fast.take_decoded(key) for key in slots) + assert not fast.submit("r", "c", [1, 2], slots) + assert outputs.qsize() == 2 + while not outputs.empty(): + assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR + finally: + release.set() + fast.close() + + +@pytest.mark.cpu +def test_platform_failure_wakes_pending_handoff(monkeypatch): + import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module + + fast, _, _, _ = _cpu_fast(monkeypatch) + fast._handoffs[0] = _SlotHandoff() + + def fail(device): + raise RuntimeError("device unavailable") + + monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail)) + fast._run() + assert fast._closed + with pytest.raises(RuntimeError, match="decode failed"): + fast.order_after(0) diff --git a/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py b/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py new file mode 100644 index 000000000..b9109c0e4 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Per-request code rows keep the runner payload layout with one host copy.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_rows_map_to_batch_positions_and_copy_once(): + codes = torch.arange(6, dtype=torch.long).reshape(2, 3) + snapshot = _CodeRowsSnapshot(codes, [2, 0], 4) + assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None + rows = snapshot["codes"]["audio"] + assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)] + assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:]) + + copies = [] + + def copy(tensor): + copies.append(tensor) + return tensor.clone() + + host = snapshot.copy_to_cpu(copy)["codes"]["audio"] + assert len(copies) == 1 and copies[0] is codes + codes.zero_() + assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]] + assert host[1].numel() == 0 and host[1].dtype == torch.long diff --git a/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py b/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py new file mode 100644 index 000000000..6c926c134 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.tts, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +def test_explicit_generator_bypasses_compiled_audio_sampler(): + from transformers import GPT2Config + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, + ) + + model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval() + heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16() + embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16() + binary = torch.nn.Linear(80, 2).cuda().bfloat16() + hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16) + generator = torch.Generator(device="cuda").manual_seed(17) + expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator) + + def forbidden(*args, **kwargs): + raise AssertionError("compiled path selected") + + model._compiled_audio_sampler = forbidden + generator.manual_seed(17) + actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator) + assert all(torch.equal(a, b) for a, b in zip(expected, actual)) + with pytest.raises(AssertionError, match="compiled path selected"): + model.generate_frame(hidden, heads, embeddings, binary, n_vq=2) diff --git a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py new file mode 100644 index 000000000..2d529cf26 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Determined text tokens retain the upstream sampling/counting contract.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = pytest.mark.core_model + + +@pytest.mark.cuda +@pytest.mark.parametrize("seed", [None, 17]) +@pytest.mark.parametrize("temperature", [0.0, 1.7]) +def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda") + reqs = RequestState(5, 32, 16, 0, 8, device) + sampler = Sampler(5, 8, device, reqs) + for i, length in enumerate([4, 1, 3]): + reqs.add_request(str(i), length, [2] * length, 0, 8) + slot = reqs.req_id_to_index[str(i)] + sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed)) + reqs.apply_staged_writes() + sampler.apply_staged_writes() + state = _state(MossLocalModelState, device) + state._direct_tokens = True + state.model._batch_state = None + batch = _batch(device, [2, 4, 3], [1, 2, 1]) + batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32) + batch.cu_num_logits_np = np.arange(4, dtype=np.int32) + batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device) + batch.expanded_idx_mapping = batch.idx_mapping + batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device) + batch.logits_indices = batch.query_start_loc[1:].long() - 1 + batch.positions = torch.arange(batch.num_tokens, device=device) + state.model._batch_should_continue = torch.tensor([False, True, True], device=device) + with torch.inference_mode(): + logits = state.model.compute_logits(torch.ones(3, 4, device=device)) + expected = sampler(logits, batch) + actual = state.sample_determined_tokens(batch, sampler) + assert actual is not None + for field in ["sampled_token_ids", "num_sampled", "num_rejected"]: + torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0) + assert actual.num_sampled.tolist() == [1, 0, 1] + snapshot = actual.sampled_token_ids.clone() + state.model._batch_should_continue.logical_not_() + torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "constraint", + [ + {"min_tokens": 2, "stop_token_ids": [1]}, + {"allowed_token_ids": [1, 2]}, + {"logit_bias": {1: -10.0}}, + ], +) +def test_real_sampler_constraints_disable_determined_tokens(constraint): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda") + reqs = RequestState(5, 32, 16, 0, 8, device) + sampler = Sampler(5, 8, device, reqs) + reqs.add_request("constrained", 1, [2], 0, 8) + slot = reqs.req_id_to_index["constrained"] + sampler.add_request(slot, 1, SamplingParams(**constraint)) + state = _state(MossLocalModelState, device) + state._direct_tokens = True + # Use the actual upstream request registration, including its combined + # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves. + assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"]) +def test_distribution_features_fall_back_before_using_gpu(feature): + sampler = object.__new__(Sampler) + sampler.compute_nans = feature == "nans" + sampler.return_sampling_mask = feature == "mask" + sampler.trace_replay_state = object() if feature == "trace" else None + sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None + sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"])) + sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"])) + sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"]))) + sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"])) + state = object.__new__(MossLocalModelState) + state._direct_tokens = True + assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"]) +def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature): + hook = mocker.Mock(return_value=None) + + class State: + def sample_determined_tokens(self, batch, sampler): + return hook(batch, sampler) + + runner = object.__new__(OmniARModelRunner) + runner.model_state = State() + runner.sampler = object() + runner.batch_sharder = object() if feature == "sharding" else None + runner.pp_handler = object() if feature == "pipeline" else None + batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1) + grammar = object() if feature == "grammar" else None + fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream") + hidden = object() + assert runner.sample(hidden, batch, grammar) == "upstream" + fallback.assert_called_once_with(hidden, batch, grammar) + assert hook.call_count == int(feature == "declined") diff --git a/tests/model_executor/models/moss_tts/test_local_eager_mtp.py b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py new file mode 100644 index 000000000..cf3893cbd --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Eager MTP emits the same frames one step earlier with identical backbone inputs.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState + +pytestmark = pytest.mark.core_model + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device(request.param) + + +def _run(state, device, schedule, stop_step): + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[3], prompt[0] = 4, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + embeds = [] + frames: dict[int, list[torch.Tensor]] = {3: [], 0: []} + for index, (slots, counts) in enumerate(schedule): + state.model.force_stop = index == stop_step + embed, payload, _ = _step(state, _batch(device, slots, counts), req_states) + embeds.append(embed) + rows = payload.get("codes", {}).get("audio", []) + for slot, row in zip(slots, rows): + if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()): + frames[slot].append((index, row.clone())) + for slot, count in zip(slots, counts): + computed[slot] += count + return embeds, frames + + +def test_eager_frames_match_canonical_one_step_earlier(device): + schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])] + canonical, eager = (_state(MossLocalModelState, device) for _ in range(2)) + eager._local_eager_mtp = True + for state in (canonical, eager): + _admit(state, 3, "long", 17) + _admit(state, 0, "short", 17) + # Stop both streams at the canonical step that would draw the 4th step's frames. + c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4) + e_embeds, e_frames = _run(eager, device, schedule, stop_step=3) + + for c, e in zip(c_embeds, e_embeds): + torch.testing.assert_close(c, e, rtol=0, atol=0) + for slot in (3, 0): + assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot]) + for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]): + assert ei == ci - 1 + torch.testing.assert_close(cf, ef, rtol=0, atol=0) + + +def test_stopped_stream_does_not_emit_again(device): + state = _state(MossLocalModelState, device) + state._local_eager_mtp = True + _admit(state, 0, "short", 17) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0] = 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + emitted = [] + for index, count in enumerate([2, 1, 1, 1]): + state.model.force_stop = index == 1 + _, payload, _ = _step(state, _batch(device, [0], [count]), req_states) + row = payload["codes"]["audio"][0] + emitted.append(bool(row.ne(state.model.audio_pad_token_id).any())) + computed[0] += count + state.model.force_stop = False + # Frame at the prefill step, stop at step 1, then asynchronous overrun steps. + assert emitted == [True, False, False, False] + + +def test_unmarked_decode_row_ignores_stale_slot_state(device): + canonical, eager = (_state(MossLocalModelState, device) for _ in range(2)) + eager._local_eager_mtp = True + for state in (canonical, eager): + _admit(state, 0, "short", 17) + # A decode row whose completing prefill was not observed in this state. + eager._keep_pool[0] = False + eager._eager_emb_pool[0].fill_(3.0) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0], computed[0] = 2, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + for state in (canonical, eager): + state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False} + c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states) + e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states) + torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0) + torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0) + + +@pytest.mark.cpu +def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker): + """The Local config must not enable the parent's incompatible eager hook.""" + from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner + + device = torch.device("cpu") + template = _state(MossLocalModelState, device) + model = template.model + model.config = SimpleNamespace(mrv2_eager_mtp=True) + + def initialize_parent(state, *args): + state.__dict__.update(template.__dict__) + state._eager_mtp, state._eager_rows = False, None + state._eager_state = EagerMTPState(state) + + monkeypatch.setattr(OmniModelState, "__init__", initialize_parent) + state = MossLocalModelState(SimpleNamespace(), model, None, device) + _admit(state, 0, "req", 17) + batch = _batch(device, [0], [2]) + batch.req_ids = ["req"] + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0] = 2 + _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)) + hook = mocker.spy(state, "run_eager_mtp") + + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.execute_model_state = SimpleNamespace( + input_batch=batch, + hidden_states=torch.zeros(2, 4), + finished_req_ids=set(), + ec_connector_output=None, + routed_experts=None, + ) + runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None + runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False + runner.model_config = SimpleNamespace(async_chunk=False) + runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text")) + runner.model_state, runner.model = state, model + runner.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)), + num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)), + prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)), + ) + runner.main_stream = runner.output_copy_stream = mocker.Mock() + runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock() + runner.sampler = None + runner.sample = mocker.Mock( + return_value=( + SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), + torch.ones(1), + torch.zeros(1), + ) + ) + runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={})) + runner.postprocess_sampled = mocker.Mock() + runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None)) + mock_out = SimpleNamespace(copy_event=None) + monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out)) + + assert runner.sample_tokens(None) is mock_out + hook.assert_called_once() + assert state._eager_mtp is False and state._eager_rows is None + assert state._local_eager_mtp is True diff --git a/tests/model_executor/models/moss_tts/test_local_model_state.py b/tests/model_executor/models/moss_tts/test_local_model_state.py new file mode 100644 index 000000000..4956d5c89 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_model_state.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare slot lifecycle with canonical Local hooks, including output ownership.""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from vllm.v1.worker.gpu.model_states.default import DefaultModelState + +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration +from vllm_omni.worker_v2.model_states import init_omni_model_state +from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState +from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +pytestmark = pytest.mark.core_model + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device(request.param) + + +def _model(device): + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + torch.nn.Module.__init__(model) + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model.model = torch.nn.Module() + weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16 + model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight) + model.hidden_size, model.n_vq = 4, 2 + model.audio_pad_token_id, model.text_vocab_size = 8, 8 + model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3 + model.talker_mtp_output_key = ("audio_codes", "current") + model.talker_mtp_graph_safe = False + model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")} + model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None + model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8 + model.force_stop = False + + def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs): + assert (temperature, top_k, top_p) == (1.7, 25, 0.8) + noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator) + codes = (hidden[:, :2].mul(16).long() + noise) % 8 + return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes + + model.local_transformer = SimpleNamespace(generate_frame=frame) + return model + + +def _state(cls, device): + state = object.__new__(cls) + state.model = _model(device) + state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype + state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True + state.scheduler_config = SimpleNamespace(max_num_seqs=5) + state.vllm_config = SimpleNamespace() + state.intermediate_buffer = OmniIntermediateBuffer(5) + state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype) + state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long) + state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long) + for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"): + setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype)) + state._mtp_runner = state._mtp_sample_uniforms = None + state._mtp_generators = {} + state._stream_pos = {} + state._stream_decode_event = None + state._eager_state = EagerMTPState(state) + if cls is MossLocalModelState: + state._init_slot_buffers(5, 4, device, state.dtype) + return state + + +def _admit(state, slot, name, seed): + params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None) + with patch.object(DefaultModelState, "add_request", return_value=None): + state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params)) + state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])} + + +def _batch(device, slots, counts): + starts = np.array([0, *np.cumsum(counts)], dtype=np.int32) + return SimpleNamespace( + idx_mapping_np=np.array(slots), + idx_mapping=torch.tensor(slots, device=device), + num_reqs=len(slots), + num_tokens=int(starts[-1]), + num_scheduled_tokens=np.array(counts), + query_start_loc_np=starts, + query_start_loc=torch.tensor(starts, device=device), + # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs. + input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device), + ) + + +def _step(state, batch, req_states, dispatcher=None): + inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]} + with ( + patch("vllm.forward_context.set_forward_context", return_value=nullcontext()), + patch( + "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext() + ), + torch.inference_mode(), + ): + state.run_preprocess(batch, inputs, req_states, dispatcher) + hidden = inputs["inputs_embeds"] + 0.25 + state.run_postprocess(hidden, batch) + _, payload = state.postprocess_model_output(hidden, batch, req_states) + last = batch.query_start_loc[1:] - 1 + tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1) + result = inputs["inputs_embeds"].clone(), payload, tokens + hidden.fill_(-100) # next graph replay can overwrite its output immediately + return result + + +@pytest.mark.parametrize("seed", [None, 17]) +@pytest.mark.parametrize("batch_prefill", [False, True]) +def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill): + reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState)) + candidate._batch_prefill = batch_prefill + for state in (reference, candidate): + _admit(state, 3, "long", seed) + _admit(state, 0, "short", seed) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[3], prompt[0] = 4, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + saved = [] + for index, (slots, counts) in enumerate( + [ + ([3, 0], [2, 2]), + ([0, 3], [1, 1]), + ([3, 0], [1, 1]), + ([3, 0], [1, 1]), + ([0, 3], [1, 1]), + ([3, 0], [1, 1]), + ] + ): + if index == 4: + # Cancellation/removal followed by reuse must not retain hidden, + # reference offset, control, or the old request's RNG stream. + for state in (reference, candidate): + state.remove_request("short") + assert "short" not in state._mtp_generators + _admit(state, 0, "replacement", seed) + computed[0], prompt[0] = 0, 1 + outputs = [] + for state in (reference, candidate): + state.model.force_stop = index == 3 + torch.manual_seed(91 + index) + outputs.append(_step(state, _batch(device, slots, counts), req_states)) + a, b = outputs + torch.testing.assert_close(a[0], b[0], rtol=0, atol=0) + torch.testing.assert_close(a[2], b[2], rtol=0, atol=0) + if index == 3: + assert b[2].tolist() == [3, 3] + assert a[1].keys() == b[1].keys() + for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])): + torch.testing.assert_close(ar, br, rtol=0, atol=0) + saved.append((br, br.clone())) + for slot, count in zip(slots, counts): + computed[slot] += count + old = reference.intermediate_buffer.buffers[slot] + new = candidate.intermediate_buffer.buffers[slot] + assert old["ref_offset"] == new["ref_offset"] + torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0) + assert "hidden_states" not in new and "audio_codes" not in new + candidate._codes_pool.fill_(8) + candidate._hidden_pool.zero_() + for current, snapshot in saved: + torch.testing.assert_close(current, snapshot, rtol=0, atol=0) + + +def test_slot_honors_external_stopping_control(device): + states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)] + for state in states: + _admit(state, 0, "stop", None) + req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5)) + _step(state, _batch(device, [0], [1]), req) + state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True + req.num_computed_tokens[:] = 1 + output = _step(state, _batch(device, [0], [1]), req) + assert output[2].tolist() == [3] + assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]] + + +@pytest.mark.cuda +def test_slot_mtp_graph_padding_reorder_and_owned_output(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)] + for state in states: + + def frame(hidden, *_args, **_kwargs): + return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8 + + state.model.local_transformer.generate_frame = frame + for slot in [0, 3, 4]: + _admit(state, slot, f"req{slot}", None) + candidate = states[1] + buffers = [ + candidate._mtp_input_ids[:4], + candidate._mtp_input_embeds[:4], + candidate._mtp_hidden[:4], + candidate._mtp_text_step[:4], + ] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream), torch.inference_mode(): + for _ in range(3): + candidate.model.mtp(*buffers) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph), torch.inference_mode(): + graph_output = candidate.model.mtp(*buffers) + + def replay(*args, **kwargs): + assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers)) + assert args[0].shape[0] == 4 + graph.replay() + return graph_output + + candidate._mtp_runner = replay + candidate._is_mtp_graph_runner = lambda: True + from vllm.config import CUDAGraphMode + + def dispatcher(bsz): + return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL) + + req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32)) + retained = [] + for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]: + out = [ + _step( + state, + _batch(torch.device("cuda"), slots, [1] * len(slots)), + req, + dispatcher if state is candidate else None, + ) + for state in states + ] + torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0) + torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0) + for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])): + torch.testing.assert_close(a, b, rtol=0, atol=0) + retained.append((b, b.clone())) + req.num_computed_tokens[slots] += 1 + with torch.inference_mode(): + for tensor in graph_output: + tensor.fill_(7) + for tensor, copy in retained: + torch.testing.assert_close(tensor, copy, rtol=0, atol=0) + + +@pytest.mark.cpu +def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch): + monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None) + dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None) + assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState) + dynamic.create_mrv2_model_state.assert_not_called() + + class Invalid: + has_preprocess = True + + def create_mrv2_model_state(self, *args): + return object() + + with pytest.raises(TypeError, match="must return a ModelState"): + init_omni_model_state(None, Invalid(), None, torch.device("cpu")) + + +@pytest.mark.cpu +@pytest.mark.parametrize("enabled", [False, True]) +def test_local_factory_default_and_option(monkeypatch, enabled): + monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None) + monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None) + model = _model(torch.device("cpu")) + model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled) + state = init_omni_model_state(None, model, None, torch.device("cpu")) + assert type(state) is (MossLocalModelState if enabled else OmniModelState) diff --git a/tests/model_executor/models/moss_tts/test_local_prefill_batch.py b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py new file mode 100644 index 000000000..8b6856f4a --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference batching preserves coordinates and asynchronous source ownership.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration + +pytestmark = pytest.mark.core_model + + +@pytest.mark.cpu +@pytest.mark.parametrize("flat", [False, True]) +def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat): + state = _state(MossLocalModelState, torch.device("cpu")) + state._batch_prefill = True + codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]]) + for slot, offset in [(3, 1), (0, 2)]: + state.intermediate_buffer.buffers[slot] = { + "req_id": str(slot), + "codes": {"ref": codes.flatten() if flat else codes}, + "ref_offset": offset, + } + batch = _batch(torch.device("cpu"), [3, 0], [2, 1]) + req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5)) + embeds = state._static_inputs_embeds[:3] + expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]]) + embed = mocker.spy(state.model.model.embed_tokens, "forward") + with torch.inference_mode(): + state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req) + torch.testing.assert_close(embeds, expected, rtol=0, atol=0) + assert embed.call_count == 1 + assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3 + assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3 + + +@pytest.mark.cuda +def test_pending_reference_uploads_survive_staging_reuse_and_overflow(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = _state(MossLocalModelState, torch.device("cuda")) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + outputs = [] + with torch.cuda.stream(stream), torch.inference_mode(): + # Keep the first H2D pending while CPU prepares all subsequent batches. + torch.cuda._sleep(100_000_000) + for value in range(1, 8): + embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype) + codes = torch.full((2, 2), value) + state._apply_reference_batch(embeds, [(1, codes)]) + outputs.append((embeds, value)) + codes.fill_(-1) + stream.synchronize() + assert len(state._prefill_staging) <= 3 + for output, value in outputs: + expected = torch.zeros_like(output) + expected[1:3] = value / 8 + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@pytest.mark.cuda +def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = _state(MossLocalModelState, torch.device("cuda")) + model = state.model + model.n_vq = 12 + model.audio_vocab_size = model.audio_pad_token_id = 1024 + model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16) + model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model) + state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16) + references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]] + with torch.inference_mode(): + actual = torch.randn_like(state._static_inputs_embeds) + expected = actual.clone() + for start, codes in references: + expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda")) + state._apply_reference_batch(actual, references) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/model_executor/models/moss_tts/test_local_short_attention.py b/tests/model_executor/models/moss_tts/test_local_short_attention.py new file mode 100644 index 000000000..212d32f36 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_short_attention.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.tts, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)]) +def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + torch.manual_seed(43) + query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16) + key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + for position in range(12): + # Non-contiguous prefixes retain the stride of the full 12-slot cache. + k, v = key[:, :, : position + 1], value[:, :, : position + 1] + before_k, before_v = key.clone(), value.clone() + expected = F.scaled_dot_product_attention(query, k, v) + actual = local_short_attention(query, k, v) + assert torch.equal(key, before_k) and torch.equal(value, before_v) + if position == 0: + assert torch.equal(actual, v) + torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01) + # Future cache contents must never affect an earlier position. + key[:, :, position + 1 :] += 100 + value[:, :, position + 1 :] -= 100 + assert torch.equal(local_short_attention(query, k, v), actual) + key.copy_(before_k) + value.copy_(before_v) + + +def test_compile_and_graph_replay_reads_fresh_inputs(): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16) + key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + compiled = torch.compile(local_short_attention, fullgraph=True) + for _ in range(3): + compiled(query, key, value) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = compiled(query, key, value) + for _ in range(3): + query.normal_() + key.normal_() + value.normal_() + graph.replay() + torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01) + + +def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch): + from transformers import GPT2Config + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, + ) + + monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1") + model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval() + attn = model.h[0].attn + attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16) + cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2)) + forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True) + with torch.inference_mode(): + for _ in range(2): + inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16) + # The second frame overwrites the previous frame's K/V before use. + for position in range(12): + expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:] + actual = forward(inputs[:, position : position + 1], cache, position) + torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02) diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder.py b/tests/model_executor/models/moss_tts/test_reference_encoder.py index cfd8fb5c2..06efaaf25 100644 --- a/tests/model_executor/models/moss_tts/test_reference_encoder.py +++ b/tests/model_executor/models/moss_tts/test_reference_encoder.py @@ -598,3 +598,236 @@ def test_reference_resampler_is_cached_and_matches_functional_resample(): assert info.misses == 1 and info.hits == 1 finally: _reference_resampler.cache_clear() + + +# --------------------------------------------------------------------------- # +# Inline data: URI fast path # +# --------------------------------------------------------------------------- # + + +class _DigestAudio(_FakeAudio): + """Resolver whose key for data: URIs is the SHA-1 digest, like serving.""" + + def __init__(self, *, digest_keys=True): + super().__init__() + self._digest_keys = digest_keys + + async def resolve(self, ref_str): + wav_list, sr, _ = await super().resolve(ref_str) + if not self._digest_keys: + return wav_list, sr, "rk:" + ref_str + import hashlib + + return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest() + + def artifact_key(self, cache_key): + return None + + +async def test_inline_reference_skips_repeat_resolution(make_encoder): + proc, audio = _FakeProcessor(), _DigestAudio() + ref = "data:audio/wav;base64,QUJD" + audio.register(ref, 5) + enc = make_encoder(proc) + + first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + first.fill_(0) # callers own their copy + second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + + assert audio.resolve_calls == [ref] + assert key1 == key2 + assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long)) + + +async def test_inline_index_requires_matching_resolve_key(make_encoder): + proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False) + ref = "data:audio/wav;base64,REVG" + audio.register(ref, 6) + enc = make_encoder(proc) + + for _ in range(2): + await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + assert audio.resolve_calls == [ref, ref] + + +async def test_inline_reference_shared_across_encoders(make_encoder, tmp_path_factory, monkeypatch): + shared_dir = tmp_path_factory.mktemp("ref") + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", str(shared_dir)) + ref = "data:audio/wav;base64,R0hJ" + first_audio, second_audio = _DigestAudio(), _DigestAudio() + first_audio.register(ref, 9) + second_audio.register(ref, 9) + first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor()) + + a, key_a = await first.encode(ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key) + b, key_b = await second.encode( + ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key + ) + + assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == [] + assert key_a == key_b and torch.equal(a, b) + + +# --------------------------------------------------------------------------- # +# Reference-encoder CUDA graphs: selection and processor-equivalent prep # +# --------------------------------------------------------------------------- # + + +class _GraphsStub: + def __init__(self, result): + self.result = result + self.calls: list[list[torch.Tensor]] = [] + + def encode(self, wavs): + self.calls.append(wavs) + return self.result + + +class _LoudnessProcessor(_FakeProcessor): + @staticmethod + def loudness_normalize(wav): + return wav * 2.0 + + +async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder): + proc = _LoudnessProcessor() + enc = make_encoder(proc) + graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)] + enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes) + prepared = [torch.ones(1, 100)] + assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == [] + enc._graphs = _GraphsStub([None]) # no graph fits this clip + out = enc._encode_prepared(prepared) + assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3 + + +async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder): + enc = make_encoder(_LoudnessProcessor()) + mono = enc._processor_prepare(torch.ones(100)) + assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0)) + many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4)) + assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0) + + +async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder): + torchaudio = pytest.importorskip("torchaudio") + proc = _LoudnessProcessor(sampling_rate=48000) + enc = make_encoder(proc) # works at 24 kHz, the Local-v1.5 reference rate + wav = torch.randn(1, 2400) + expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0 + # The encoder resamples the mono clip before duplicating it. Some CPU conv + # kernels round a one-row batch differently from a two-row one, so the two + # orders agree to float32 precision, not bit for bit. + torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize("role", ["", "host", "client"]) +async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role): + from types import SimpleNamespace + + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + enc = make_encoder(_FakeProcessor()) + enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role) + captured, hosts, ready = [], [], [] + monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True)) + monkeypatch.setattr(enc, "_host_workers", lambda: []) + helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True)) + + def make_host(*args, **kwargs): + hosts.append(args) + return helper + + monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host) + monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper) + enc.prepare() + enc.prepare() + # Every API process waits for the shared host before startup completes. + assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role] + assert len(ready) == (2 if role else 0) + + +async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch): + torchaudio = pytest.importorskip("torchaudio") + proc = _LoudnessProcessor(sampling_rate=48000) + enc = make_encoder(proc) # works at 24 kHz + sent = [] + + def record_prepared(prepared): + sent.extend(prepared) + return [None] * len(prepared) + + monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared) + wav = [float(v) for v in torch.randn(2400)] + enc._encode_batch_sync([(wav, _SR)]) + # Resampled where the request arrived (a shared encoder receives it this way); still one channel. + expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000) + torch.testing.assert_close(sent[0], expected, rtol=0, atol=0) + # The eager path is told the clip is already at the tokenizer rate. + calls = [] + + def record_encoding(wavs, sampling_rate, n_vq=None): + calls.append(sampling_rate) + return [None] + + monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding) + enc._graphs_enabled = False + enc._encode_prepared(sent) + assert calls == [48000] + + +async def test_batcher_runs_up_to_max_inflight_batches_at_once(): + inside = threading.Barrier(2, timeout=5) + seen: list[int] = [] + + def encode_batch(payload): + inside.wait() # two batches are encoding at the same time + seen.append(len(payload)) + return [torch.zeros(1) for _ in payload] + + batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2) + try: + first = asyncio.ensure_future(batcher.submit([1.0], _SR)) + await asyncio.sleep(0.05) # the first batch is in flight before the second clip arrives + await asyncio.gather(first, batcher.submit([2.0], _SR)) + finally: + await batcher.aclose() + assert seen == [1, 1] + + +def test_reference_stream_priority(monkeypatch): + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False) + assert re_mod.reference_stream_priority() == 0 + monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high") + expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0 + assert re_mod.reference_stream_priority() == expected + + +async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch): + torchaudio = pytest.importorskip("torchaudio") + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + proc = _LoudnessProcessor(sampling_rate=48000) + monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False) # CPU placement + enc = make_encoder(proc) + enc._gpu_prep = True + wav = torch.randn(1, 2400) + clip = enc._encoder_input(wav) + assert clip.shape == (1, 2400) # left at the working rate for the encoder's device + expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0 + torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6) + # The eager fallback is told the clip's own rate. + calls = [] + + def record_encoding(wavs, sampling_rate, n_vq=None): + calls.append(sampling_rate) + return [None] + + monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding) + enc._graphs_enabled = False + enc._encode_prepared([clip]) + assert calls == [24000] + assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP" diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py b/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py new file mode 100644 index 000000000..74969005e --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode.""" + +from contextlib import nullcontext + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import ( + MossReferenceEncoderGraphs, + _cached_reference_rope, + _reference_rope_factors, + _validate_unmasked_reference_encoder, + _windowed_attention, + install_windowed_attention, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +FRAME, N_VQ = 8, 3 + + +class _Stage(nn.Module): + """Causal per-frame stage: frames never mix, like a causal encoder layer.""" + + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + + def forward(self, x, lengths): + return torch.tanh(self.proj(x)), lengths + + +class _Quantizer(nn.Module): + def forward(self, hidden, lengths, n_q): + # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices. + codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)]) + return hidden, codes, lengths + + +class _Tokenizer(nn.Module): + number_channels = 2 + downsample_rate = FRAME + sampling_rate = 64 + + def __init__(self): + super().__init__() + self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)]) + self.quantizer = _Quantizer() + + def _flatten_channels_for_codec(self, audio, lengths): + batch, channels, samples = audio.shape + frames = samples // FRAME + values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1) + return values, lengths // FRAME + + def _codec_inference_autocast(self): + return nullcontext() + + def eager(self, wavs): + out = [] + for wav in wavs: + values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]])) + for module in self.encoder: + values, lengths = module(values, lengths) + _, codes, _ = self.quantizer(values, lengths, N_VQ) + out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long()) + return out + + +@pytest.fixture +def cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device("cuda", torch.accelerator.current_device_index()) + + +def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda): + from types import SimpleNamespace + + module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0)) + cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + torch.manual_seed(17) + # Prime once; all subsequent sequences and batches use an immutable prefix. + q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16) + k = torch.randn_like(q) + _cached_reference_rope(module, q, k, cache) + stored = next(iter(cache.values())) + for batch, length in [(1, 17), (4, 64), (2, 128)]: + q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16) + k = torch.randn_like(q) + before = (q.clone(), k.clone()) + got = _cached_reference_rope(module, q, k, cache) + cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda) + torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0) + torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0) + for value, source, saved in zip(got, (q, k), before): + real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1) + expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype) + torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0) + torch.testing.assert_close(source, saved, rtol=0, atol=0) + assert len(cache) == 1 + assert next(iter(cache.values())) is stored + + +def test_cached_rope_distinguishes_period_and_head_width(cuda): + from types import SimpleNamespace + + cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]: + module = SimpleNamespace(rope=SimpleNamespace(max_period=period)) + q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16) + _cached_reference_rope(module, q, q, cache) + assert len(cache) == 3 + + +@pytest.mark.parametrize("batched_transfer", [False, True]) +def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer): + torch.manual_seed(0) + tokenizer = _Tokenizer().to(cuda) + # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples. + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer + ) + graphs.capture() + assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)] + assert graphs.max_batch == 4 + + for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]: + wavs = [torch.randn(2, n * FRAME) for n in frames] + codes = graphs.encode(wavs) + expected = tokenizer.eager([w.to(cuda) for w in wavs]) + assert len(codes) == count + for got, want in zip(codes, expected): + assert got.device.type == "cpu" and got.dtype == torch.long + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + # Longer than every bucket: left to the caller; the rest still use graphs. + long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME) + codes = graphs.encode([long_clip, short_clip]) + assert codes[0] is None + torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0) + + # Wider than every batch: split, longest first, and returned in input order. + wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + # Result ownership includes the degenerate single-frame shape. Reusing + # the pinned output slab must not overwrite a previous caller's codes. + previous = graphs.encode([torch.randn(2, FRAME)])[0] + snapshot = previous.clone() + graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)]) + torch.testing.assert_close(previous, snapshot, rtol=0, atol=0) + + +def test_batch_audio_budget_limits_captured_graphs(cuda): + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0 + ) + graphs.capture() + # 4 x 4 s exceeds the 8 s budget; a single clip is always captured. + assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)] + wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +def test_finer_singleton_buckets_preserve_multiclip_batching(cuda): + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, + n_vq=N_VQ, + batch_sizes=(1, 2, 4), + bucket_seconds=(2.0, 4.0), + singleton_bucket_seconds=(1.0, 3.0), + ) + graphs.capture() + assert (1, 64) in graphs.captured and (1, 192) in graphs.captured + assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured + replay = graphs._replay + shapes = [] + + def record(entry, wavs): + shapes.append(tuple(entry.audio.shape)) + return replay(entry, wavs) + + graphs._replay = record + one = torch.randn(2, 20 * FRAME) + out = graphs.encode([one]) + assert shapes == [(1, 2, 192)] + torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0) + shapes.clear() + wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]] + out = graphs.encode(wavs) + assert shapes == [(4, 2, 256)] + for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +def test_compiled_capture_matches_eager(cuda): + torch.manual_seed(2) + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True + ) + graphs.capture() + assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)] + wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +class _Attention(nn.Module): + """The attention surface the tokenizer's MHA exposes to the windowed kernel.""" + + def __init__(self, embed_dim=128, heads=2, context=5): + super().__init__() + self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True + self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False) + + def _project_qkv(self, x): + b, t, _ = x.shape + return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0) + + def _apply_dense_rope(self, q, k): + return q, k + + def _forward_non_streaming_sdpa(self, x, input_lengths): + # Reference semantics: causal, distance < context, padded rows zeroed. + b, t, _ = x.shape + q, k, v = self._project_qkv(x) + pos = torch.arange(t, device=x.device) + delta = pos.view(-1, 1) - pos.view(1, -1) + mask = (delta >= 0) & (delta < self.context) + out = F.scaled_dot_product_attention(q, k, v, mask[None, None]) + valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1) + out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype)) + return out.transpose(1, 2).reshape(b, t, self.embed_dim) + + +@pytest.mark.parametrize("fa_version", [2, 3]) +def test_windowed_attention_matches_masked_sdpa(cuda, fa_version): + if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9: + pytest.skip("FA3 requires Hopper") + torch.manual_seed(1) + model = nn.Module() + model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)]) + attention = model.encoder[0] + x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16) + lengths = torch.tensor([40, 17, 3], device=cuda) + expected = attention._forward_non_streaming_sdpa(x, lengths) + assert install_windowed_attention(model, fa_version=fa_version) == 1 + actual = attention._forward_non_streaming_sdpa(x, lengths) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0 + # Non-bf16 inputs keep the tokenizer's own SDPA path. + x32 = x.float() + attention.float() + torch.testing.assert_close( + attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2 + ) + + +@pytest.mark.parametrize("fa_version", [2, 3]) +def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version): + if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9: + pytest.skip("FA3 requires Hopper") + torch.manual_seed(44) + attention = _Attention().to(cuda, torch.bfloat16) + x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16) + lengths = torch.tensor([40, 17, 1, 0], device=cuda) + kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version) + with torch.no_grad(): + masked = _windowed_attention(attention, x, lengths, **kwargs) + unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs) + poisoned = x.clone() + for row, length in enumerate([40, 17, 1, 0]): + poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30 + changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs) + for row, length in enumerate([40, 17, 1, 0]): + torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0) + torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0) + assert masked[row, length:].count_nonzero() == 0 + assert unmasked[1, 17:].count_nonzero() > 0 + # The FP32 fallback retains its original masking contract. + attention.float() + got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs) + expected = attention._forward_non_streaming_sdpa(x.float(), lengths) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_unmasked_queries_reject_unknown_encoder_before_patching(): + tokenizer = _Tokenizer() + with pytest.raises(ValueError, match="patch-downsample"): + install_windowed_attention(tokenizer, skip_padded_query_mask=True) + + +@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"]) +def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part): + from types import SimpleNamespace + + patch = type("MossAudioTokenizerPatchedPretransform", (), {})() + patch.is_downsample = bad_part != "upsample" + projected = type("MossAudioTokenizerProjectedTransformer", (), {})() + projected.transformer = SimpleNamespace( + layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))] + ) + quantizer_cls = type( + "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {} + ) + quantizer = quantizer_cls() + quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1) + tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer) + with pytest.raises(ValueError): + _validate_unmasked_reference_encoder(tokenizer) diff --git a/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py new file mode 100644 index 000000000..b8bd50184 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""API processes share one reference encoder through a Unix socket.""" + +import threading + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _codes(wav): + # Deterministic per clip: frames from the clip length, values from its first sample. + return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long) + + +class _Recorder: + def __init__(self, gate=None): + self.batches = [] + self.gate = gate + + def __call__(self, wavs): + if self.gate is not None: + self.gate.wait(5) + self.batches.append(len(wavs)) + return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs] + + +@pytest.fixture +def host_dir(tmp_path_factory, monkeypatch): + monkeypatch.setattr(shared, "_host_lock_handle", None) + # AF_UNIX limits the full socket path, including pytest directory names. + return str(tmp_path_factory.mktemp("ref")) + + +def test_only_one_process_per_directory_hosts(host_dir): + assert shared.elect_host(host_dir) is True + assert shared.elect_host(host_dir) is False # a second open file description cannot take the lock + + +def test_close_releases_queued_and_active_callers(host_dir): + entered, release = threading.Event(), threading.Event() + + def encode(wavs): + entered.set() + assert release.wait(5) + return [_codes(wav) for wav in wavs] + + host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)]) + results: dict[str, list[torch.Tensor | BaseException]] = {} + threads = [ + threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True) + for key in ("active", "queued") + ] + try: + threads[0].start() + assert entered.wait(2) + threads[1].start() + with host._ready: + assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2) + host.close() + for thread in threads: + thread.join(1) + assert not thread.is_alive(), "closed encoder left a request waiting" + assert all(isinstance(results[key][0], RuntimeError) for key in results) + finally: + release.set() + host.close() + for thread in threads: + if thread.ident is not None: + thread.join(1) + + +def test_worker_result_count_mismatch_fails_entire_batch(host_dir): + host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)]) + results = [] + caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True) + try: + caller.start() + caller.join(2) + assert not caller.is_alive(), "missing encoder result left a request waiting" + assert len(results) == 2 + assert all(isinstance(result, RuntimeError) for result in results) + finally: + host.close() + caller.join(1) + + +def test_client_round_trip_and_per_clip_errors(host_dir): + host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)]) + try: + client = shared.SharedReferenceEncoderClient(host_dir) + wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)] + out = client.encode(wavs) + assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0])) + assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1]) + assert torch.equal(out[2], _codes(wavs[2])) + # The host process's own clips skip the socket. + assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2])) + finally: + host.close() + + +def test_host_merges_clips_from_concurrent_clients(host_dir): + gate = threading.Event() + recorder = _Recorder(gate) + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)]) + try: + results = {} + + def send(name, value): + results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2) + + # Encoding is held until both clients' clips are queued. + warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)])) + warm.start() + threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))] + for thread in threads: + thread.start() + threading.Timer(0.5, gate.set).start() + for thread in (warm, *threads): + thread.join(10) + # Both clients' clips were encoded together, not one request at a time. + assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4 + assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"]) + finally: + host.close() + + +def test_client_without_host_raises(host_dir, monkeypatch): + monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3) + with pytest.raises((FileNotFoundError, ConnectionRefusedError)): + shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)]) + + +def test_host_window_collects_across_workers_without_reserving_singletons(host_dir): + recorder = _Recorder() + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500) + try: + barrier = threading.Barrier(3) + results = [] + + def send(value): + barrier.wait() + results.extend(host.encode([torch.full((1, 20), value)])) + + threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)] + for thread in threads: + thread.start() + barrier.wait() + for thread in threads: + thread.join(5) + assert not thread.is_alive() + assert recorder.batches == [2] + assert sorted(int(x[0, 0]) for x in results) == [1, 2] + finally: + host.close() + + +def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch): + from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client") + monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3) + proc = _FakeProcessor() + enc = re_mod.MossReferenceEncoder( + proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20) + ) + placed = [] + monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True)) + out = enc._encode_batch_sync([([7.0] * 50, _SR)]) + assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7 + + +def test_workers_encode_concurrently_after_their_warmups(host_dir): + inside, release = threading.Barrier(3, timeout=5), threading.Event() + warmed = [] + + def worker(name): + def encode(wavs): + inside.wait() # both workers are inside an encode at the same time + release.wait(5) + return [_codes(w) for w in wavs] + + return encode, lambda: warmed.append(name) + + host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")]) + try: + out = [] + + def submit(value): + out.append(host.encode([torch.full((1, 20), value)])) + + first = threading.Thread(target=submit, args=(1.0,)) + first.start() + # The first worker is busy by now, so the second clip goes to the other one. + threading.Event().wait(0.1) + second = threading.Thread(target=submit, args=(2.0,)) + second.start() + inside.wait() + release.set() + for thread in (first, second): + thread.join(5) + assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2] + finally: + host.close() + + +def test_client_requests_in_flight_use_their_own_connections(host_dir): + entered, release = threading.Semaphore(0), threading.Event() + + def worker(): + def encode(wavs): + entered.release() + release.wait(5) + return [_codes(w) for w in wavs] + + return encode, None + + host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()]) + try: + client = shared.SharedReferenceEncoderClient(host_dir) + out = [] + threads = [ + threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0) + ] + for thread in threads: + thread.start() + # A worker holds this request, still unanswered, before the next is + # sent. Sent together, one free worker may take both as one batch. + assert entered.acquire(timeout=5) + release.set() # both requests reached the host at once + for thread in threads: + thread.join(5) + assert sorted(int(o[0][0, 0]) for o in out) == [1, 2] + assert len(client._idle) == 2 # both connections are kept for reuse + finally: + host.close() + + +def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch): + entered, release = threading.Event(), threading.Event() + recorder = _Recorder() + warmed = [] + + def last_warmup(): + entered.set() + assert release.wait(5) + warmed.append("last") + + host = shared.SharedReferenceEncoderHost( + host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)] + ) + client = shared.SharedReferenceEncoderClient(host_dir) + try: + assert entered.wait(5) + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03) + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + client.wait_until_ready() + assert warmed == ["first"] and recorder.batches == [] and not client._idle + release.set() + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5) + host.wait_until_ready() + client.wait_until_ready() + assert warmed == ["first", "last"] and recorder.batches == [] + assert len(client._idle) == 1 + assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20))) + finally: + release.set() + host.close() + + +def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch): + release = threading.Event() + monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05) + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False) + host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))]) + timer = threading.Timer(0.2, release.set) + timer.start() + try: + # The first socket waits for readiness before starting its much shorter + # encode deadline. No API warmup call is required to protect this path. + client = shared.SharedReferenceEncoderClient(host_dir) + assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20))) + finally: + release.set() + timer.join(5) + host.close() + + +def test_failed_host_warmup_is_not_ready(host_dir): + def fail(): + raise RuntimeError("warmup failed") + + recorder = _Recorder() + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)]) + try: + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + shared.SharedReferenceEncoderClient(host_dir).wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.encode([torch.ones(1, 20)]) + assert recorder.batches == [] + finally: + host.close() + + +def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch): + from unittest.mock import Mock + + from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client") + enc = re_mod.MossReferenceEncoder( + _FakeProcessor(), + variant="local", + n_vq=_N_VQ, + sr_target=_SR, + speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20), + ) + enc._shared_client = Mock() + enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling") + local = Mock() + monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local) + with pytest.raises(shared.SharedReferenceEncoderStartupError): + enc._encode_shared_or_local([torch.ones(1, 20)]) + local.assert_not_called() diff --git a/tests/model_executor/models/moss_tts/test_slot_attention.py b/tests/model_executor/models/moss_tts/test_slot_attention.py new file mode 100644 index 000000000..3ee822371 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_slot_attention.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerMultiheadAttention, + MossAudioTokenizerRotaryEmbedding, + RingKVCache, + StreamingExecutionContext, +) +from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize( + "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)] +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("execution", ["graph", "compiled"]) +def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution): + torch.manual_seed(45) + batch, heads, states, dim = 3, 2, 5, 64 + projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype) + expected_cache = cache.clone() + end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0] + expected_end = end.clone() + slots = torch.tensor([0, 4, 2], device="cuda") + lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0] + + def candidate(): + return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True) + + graph = None + if execution == "compiled": + run = torch.compile(candidate, fullgraph=True) + else: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + candidate() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = candidate() + + for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]): + end.fill_(start) + expected_end.fill_(start) + cache.copy_(expected_cache) + projected.normal_() + slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17) + if graph is not None: + graph.replay() + else: + actual = run() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + torch.testing.assert_close(end, expected_end, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)]) +def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity): + q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)] + cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype) + end = torch.tensor([0, capacity + 3, 9, 0], device="cuda") + slots = torch.tensor([1, 3], device="cuda") + lengths = torch.zeros(2, device="cuda", dtype=torch.int32) + before_cache, before_end = cache.clone(), end.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity) + for active in [False, True, False, True]: + lengths.fill_(frames if active else 0) + expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity) + graph.replay() + torch.testing.assert_close(cache, before_cache, rtol=0, atol=0) + torch.testing.assert_close(end, before_end, rtol=0, atol=0) + torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025) + if not active: + assert out.count_nonzero() == 0 + + +def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity(): + torch.manual_seed(19) + capacity, frames = 400, 480 + state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32) + slots = torch.tensor([1, 0], device="cuda") + context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool)) + for start in [0, frames]: + k = torch.randn(2, 2, frames, 64, device="cuda") + v = torch.randn_like(k) + state.complete(k, v, execution_context=context) + retained = torch.arange(frames - capacity, frames, device="cuda") + destinations = (start + retained) % capacity + torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0) + torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0) + torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0) + + +def _ring_reference(q, k, v, cache, end, slots, lengths, context): + """Independent chronological ring oracle with FP32 attention arithmetic.""" + capacity = cache.shape[-2] + output = torch.zeros_like(q) + for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())): + if length == 0: + continue + start = int(end[slot]) + for token in range(max(0, length - capacity), length): + cache[0, slot, :, (start + token) % capacity] = k[row, :, token] + cache[1, slot, :, (start + token) % capacity] = v[row, :, token] + stop = start + length + physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long) + available = torch.arange(max(0, stop - capacity), stop, device=q.device) + physical_positions[available % capacity] = available + queries = torch.arange(start, stop, device=q.device) + distance = queries[:, None] - physical_positions[None] + mask = (physical_positions[None] >= 0) & (distance >= 0) + if context > 0: + mask &= distance < context + output[row, :, :length] = F.scaled_dot_product_attention( + q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None] + ).to(q.dtype) + end[slot] = stop + return output + + +@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)]) +@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"]) +def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution): + torch.manual_seed(29) + batch, heads, state_capacity = 3, 2, 5 + projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype) + # Strided metadata is allowed by the kernel contract. + end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0] + slots = torch.arange(batch * 2, device="cuda")[::2] + lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0] + context = capacity + attention = slot_ring_attention + if execution == "compiled": + # Functionalization must preserve mutations of the persistent ring + # and its offsets between calls, including slot reuse. + attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True) + if execution == "graph": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention(q, k, v, cache, end, slots, lengths, context) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context) + end.zero_() + expected_cache, expected_end = cache.clone(), end.clone() + for step in range(8): + slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + projected.normal_() + if step == 4: + end[1:3].zero_() + expected_end[1:3].zero_() + expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context) + if execution == "graph": + graph.replay() + else: + actual = attention(q, k, v, cache, end, slots, lengths, context) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + torch.testing.assert_close(end, expected_end, rtol=0, atol=0) + torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025) + padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None] + assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0 + + +@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)]) +def test_slot_attention_integrates_with_current_mha(frames, context): + torch.manual_seed(9) + reference = MossAudioTokenizerMultiheadAttention( + 128, + 2, + causal=True, + context=context, + rope=MossAudioTokenizerRotaryEmbedding(), + device="cuda", + dtype=torch.bfloat16, + ) + candidate = copy.deepcopy(reference) + candidate._slot_attention = slot_ring_attention + for model in (reference, candidate): + model._streaming_state = model._init_streaming_state(4) + for step in range(8): + slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda") + valid = torch.tensor([True, True, False], device="cuda") + execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid) + x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + expected = reference(x, x, x, execution_context=execution_context) + actual = candidate(x, x, x, execution_context=execution_context) + torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02) + cs, rs = candidate._streaming_state, reference._streaming_state + torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0) + if step == 4: + for model in (reference, candidate): + model._streaming_state.reset_slots(torch.tensor([0], device="cuda")) + + +@pytest.mark.parametrize( + "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)] +) +@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"]) +def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution): + torch.manual_seed(31) + batch, heads, state_capacity = 3, 2, 6 + projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype) + reference_cache = cache.clone() + end = torch.zeros(state_capacity, device="cuda", dtype=torch.long) + slots = torch.zeros(batch, device="cuda", dtype=torch.long) + lengths = torch.zeros(batch, device="cuda", dtype=torch.int32) + starts = torch.zeros(batch, device="cuda", dtype=torch.long) + attention = slot_ring_attention_rows + if execution == "compiled": + attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True) + if execution == "graph": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity) + cache.copy_(reference_cache) + for step in range(8): + slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + starts.copy_(end[slots]) + projected.normal_() + expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity) + if execution == "graph": + graph.replay() + else: + actual = attention(q, k, v, cache, starts, slots, lengths, capacity) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0) + + +@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)]) +def test_fused_rows_attention_integrates_with_mha_offsets(frames, context): + torch.manual_seed(11) + reference = MossAudioTokenizerMultiheadAttention( + 128, + 2, + causal=True, + context=context, + rope=MossAudioTokenizerRotaryEmbedding(), + device="cuda", + dtype=torch.bfloat16, + ) + candidate = copy.deepcopy(reference) + reference._slot_attention = slot_ring_attention + candidate._slot_attention_rows = slot_ring_attention_rows + for model in (reference, candidate): + model._streaming_state = model._init_streaming_state(4) + for step in range(8): + slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda") + valid = torch.tensor([True, True, False], device="cuda") + execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid) + x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + expected = reference(x, x, x, execution_context=execution_context) + actual = candidate(x, x, x, execution_context=execution_context) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + cs, rs = candidate._streaming_state, reference._streaming_state + torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, rtol=0, atol=0) + if step == 4: + for model in (reference, candidate): + model._streaming_state.reset_slots(torch.tensor([0], device="cuda")) diff --git a/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py b/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py new file mode 100644 index 000000000..79c4351a4 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""A second streaming-decoder graph set can replay concurrently with the first.""" + +import threading + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import ( + CUDAGraphStreamingDecoderWrapper, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +CAP, DIM = 8, 64 + + +class _Codec(nn.Module): + def __init__(self, device): + super().__init__() + self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM) + self.state = torch.zeros(CAP + 16, DIM, device=device) + + def decode_streaming_tensors(self, codes, lengths, slots, valid_rows): + x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None] + y = x @ self.weight + self.state.index_copy_(0, slots, y[:, -1]) + return y.transpose(1, 2).contiguous(), lengths + + def reset_decoder_state_slots(self, slots): + self.state.index_fill_(0, slots.to(self.state.device), 0) + + +def _wrapper(codec, **kwargs): + return CUDAGraphStreamingDecoderWrapper( + codec, + state_capacity=CAP, + num_quantizers=2, + vllm_config=None, + compiled_decode=codec.decode_streaming_tensors, + **kwargs, + ) + + +def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda", torch.accelerator.current_device_index()) + codec = _Codec(device) + main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3]) + fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True) + main.warmup(device) + fast.warmup(device) + assert main.is_ready and fast.is_ready + assert main._capture_stream is None and fast._capture_stream is not None + + codes_main = torch.randint(0, 5, (2, 4, 3), device=device) + codes_fast = torch.randint(0, 5, (2, 2, 1), device=device) + errors = [] + + def run(wrapper, codes, slots, iters): + stream = torch.cuda.Stream(device=device) + slot_ids = torch.tensor(slots, device=device) + try: + with torch.cuda.stream(stream): + for _ in range(iters): + codec.reset_decoder_state_slots(slot_ids) + audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone() + expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2) + stream.synchronize() + torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3) + except Exception as error: # surfaced below + errors.append(error) + + threads = [ + threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)), + threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish" + assert not errors, errors 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 index fff97cd96..5ccd758f3 100644 --- a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py +++ b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py @@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots(): assert [resumed[i].item() for i in (0, 1)] == [2, 36] -@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)]) -def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls): +@pytest.mark.parametrize( + "graph_capacity,expected_groups", + [ + # No graph, or graphs that do not cover the state capacity: exact tails. + (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}), + (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}), + # Full graph coverage: finished tails (3, 7) join the padded 8-frame + # bucket with the live 8-frame row; the finished 1-frame row keeps the + # exact 1-frame graph and the live 15-frame row stays alone. + (8, {8: [0, 1, 2], 1: [3], 15: [4]}), + ], +) +def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups): s = session() s._cudagraph_wrapper = ( None @@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa for i, (length, done) in enumerate(zip(lengths, finished, strict=True)) ] ) - assert step.call_count == expected_calls + # Slots are leased in request order, so slot == row here. + groups = {} + for call in step.call_args_list: + plan = call.args[0] + groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan) + terminal = call.kwargs["terminal_slots"] or set() + assert terminal == {slot for slot in plan if finished[slot]} + assert groups == expected_groups + assert step.call_count == len(expected_groups) + # Output is always cropped to the exact frame count, padded or not. assert [result[i].shape[-1] for i in range(5)] == lengths - if graph_capacity == 8: - assert list(step.call_args_list[0].args[0]) == [0, 1, 2] - assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1} - assert list(step.call_args_list[1].args[0]) == [3] # Preserve the exact one-frame graph. + # Finished rows released their slots; live rows keep theirs. + assert set(decoder._stream_req_slots) == {"2", "4"} diff --git a/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py b/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py new file mode 100644 index 000000000..9d9f6b7c1 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MOSS raw codec rows are buffered into first/regular chunks in codebook-major order.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.moss_tts import ( + _MOSS_AUDIO_PAD_CODE, + talker2codec_raw_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _manager(first=1, regular=3): + config = {"extra": {"codec_chunk_frames": regular, "initial_codec_chunk_frames": first}} + return SimpleNamespace(connector=SimpleNamespace(config=config)) + + +def _step(manager, frames, finished=False): + request = SimpleNamespace(request_id="r", external_req_id="r") + payload = talker2codec_raw_async_chunk(manager, {"codes": {"audio": frames}}, request, is_finished=finished) + if payload is not None and payload.codes is not None: + manager.put_req_chunk["r"] += 1 + manager.ramp_chunk_count["r"] += 1 + return payload + + +def test_first_regular_and_terminal_chunks(): + manager = _manager() + rows = torch.arange(6 * 4, dtype=torch.int32).reshape(6, 4) + first = _step(manager, rows[:1]) + assert first.meta.codec_chunk_frames == 1 + assert first.codes.audio.dtype == torch.int64 + assert torch.equal(first.codes.audio, rows[0].long()) + + assert _step(manager, rows[1:2]) is None + padded = torch.full((1, 4), _MOSS_AUDIO_PAD_CODE, dtype=torch.int32) + assert _step(manager, padded) is None # PAD rows are not audio + assert _step(manager, rows[2:3]) is None + regular = _step(manager, rows[3:4]) + assert regular.meta.codec_chunk_frames == 3 + assert torch.equal(regular.codes.audio, rows[1:4].long().T.reshape(-1)) + + tail = _step(manager, rows[4:6], finished=True) + assert tail.meta.codec_chunk_frames == 2 + assert torch.equal(tail.codes.audio, rows[4:6].long().T.reshape(-1)) + assert bool(tail.meta.finished) + assert "r" not in manager.code_prompt_token_ids + + +def test_buffered_rows_do_not_alias_producer_storage(): + manager = _manager(first=2, regular=2) + frames = torch.tensor([[7, 8, 9]]) + assert _step(manager, frames) is None + frames.fill_(0) + chunk = _step(manager, torch.tensor([[1, 2, 3]])) + assert torch.equal(chunk.codes.audio, torch.tensor([7, 1, 8, 2, 9, 3])) + + +@pytest.mark.parametrize("total", [1, 3, 5, 10, 13, 28, 31]) +def test_ramp_preserves_every_frame_and_terminal(total): + manager = _manager(first=1, regular=15) + manager.connector.config["extra"]["codec_chunk_ramp"] = [1, 4, 8, 15] + rows = torch.arange(total * 4, dtype=torch.int32).reshape(total, 4) + packets = [] + for index in range(total): + packet = _step(manager, rows[index : index + 1], finished=index == total - 1) + if packet is not None: + packets.append(packet) + rebuilt = torch.cat([packet.codes.audio.reshape(4, -1).T for packet in packets]) + assert torch.equal(rebuilt, rows.long()) + assert all(not bool(packet.meta.finished) for packet in packets[:-1]) + assert bool(packets[-1].meta.finished) + assert "r" not in manager.code_prompt_token_ids + if total == 31: + assert [packet.meta.codec_chunk_frames for packet in packets] == [1, 4, 8, 15, 3] diff --git a/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py b/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py new file mode 100644 index 000000000..a2277fa62 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Empty MOSS terminals must not execute prewarmed or previous codec tokens.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator +from vllm_omni.data_entry_keys import to_dict +from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import ( + _OmniConnectorPayloadTransportMixin as Transport, +) +from vllm_omni.model_executor.stage_input_processors.moss_tts import ( + _MOSS_AUDIO_PAD_CODE, + talker2codec_raw_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _request(tokens): + return SimpleNamespace( + request_id="r", + external_req_id="r", + prompt_token_ids=list(tokens), + num_prompt_tokens=len(tokens), + _all_token_ids=list(tokens), + _output_token_ids=[], + num_computed_tokens=0, + ) + + +@pytest.mark.parametrize("after_audio", [False, True]) +def test_empty_terminal_clears_native_codec_prompt(after_audio): + manager = SimpleNamespace( + connector=SimpleNamespace(config={"extra": {"initial_codec_chunk_frames": 1, "codec_chunk_frames": 3}}) + ) + source = _request([0] * 88) + previous = [0] * 88 + if after_audio: + packet = talker2codec_raw_async_chunk(manager, {"codes": {"audio": torch.ones(1, 12)}}, source) + manager.put_req_chunk["r"] += 1 + previous = packet.codes.audio.tolist() + terminal = talker2codec_raw_async_chunk( + manager, + {"codes": {"audio": torch.full((1, 12), _MOSS_AUDIO_PAD_CODE)}}, + source, + is_finished=True, + ) + payload = to_dict(terminal) + assert not Transport._payload_is_consumable(payload) + metadata = Transport._extract_scheduling_metadata(payload) + receiver = _request(previous) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata({"r": receiver}, {"r": metadata}, model_mode="generation") + assert receiver.prompt_token_ids == [] + assert receiver._all_token_ids == [] + assert receiver.num_prompt_tokens == receiver.num_computed_tokens == 0 + assert "r" in coordinator.input_terminal_req_ids + assert "r" not in manager.code_prompt_token_ids + + +@pytest.mark.parametrize("empty", [[], torch.empty(0, dtype=torch.long)]) +def test_explicit_empty_codes_clear_generation_prompt(empty): + request = _request([17, 18, 19]) + request._output_token_ids = [23] + request.num_computed_tokens = 3 + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata( + {"r": request}, {"r": {"code_predictor_codes": empty, "input_terminal": True}}, model_mode="generation" + ) + assert request.prompt_token_ids == request._all_token_ids == request._output_token_ids == [] + assert request.num_prompt_tokens == request.num_computed_tokens == 0 + assert "r" in coordinator.input_terminal_req_ids + + +@pytest.mark.parametrize("metadata", [{"input_terminal": True}, {"code_predictor_codes": None}]) +def test_absent_codes_preserve_generation_prompt(metadata): + request = _request([17, 18, 19]) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata({"r": request}, {"r": metadata}, model_mode="generation") + assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19] + + +def test_ar_prompt_is_not_replaced_by_empty_audio_codes(): + request = _request([17, 18, 19]) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata( + {"r": request}, {"r": {"code_predictor_codes": [], "input_terminal": True}}, model_mode="ar" + ) + assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19] diff --git a/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py new file mode 100644 index 000000000..4625c6399 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Progressive Local codec chunks preserve frames and agree with graph shapes.""" + +from collections import defaultdict +from dataclasses import dataclass, field +from unittest.mock import Mock + +import pytest +import torch +from transformers import PretrainedConfig +from vllm.config import CompilationConfig, ModelConfig, SchedulerConfig, VllmConfig + +from vllm_omni.data_entry_keys import OmniPayloadStruct +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.model_executor.stage_input_processors.moss_tts import talker2codec_raw_async_chunk + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass +class _Connector: + config: dict[str, dict[str, object]] + + +@dataclass +class _TransferManager: + connector: _Connector + code_prompt_token_ids: dict[str, list[object]] = field(default_factory=lambda: defaultdict(list)) + request_payload: dict[str, object] = field(default_factory=dict) + put_req_chunk: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + ramp_chunk_count: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + + +@dataclass +class _Request: + external_req_id: str + + +def _manager(ramp: object = None, *, first: int = 1) -> _TransferManager: + extra = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp} + return _TransferManager(_Connector({"extra": extra})) + + +def _emit( + manager: _TransferManager, req_id: str, value: int | None, finished: bool = False +) -> OmniPayloadStruct | None: + output = None if value is None else {"codes": {"audio": torch.full((1, 4), value)}} + payload = talker2codec_raw_async_chunk(manager, output, _Request(req_id), finished) + if payload is not None: + manager.put_req_chunk[req_id] += 1 + manager.ramp_chunk_count[req_id] += 1 + return payload + + +@pytest.mark.parametrize( + "ramp,first,expected", + [ + (None, 1, [1, 15, 15, 9]), + ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15, 10]), + ([1, 4, 15], 7, [1, 4, 15, 15, 5]), + ([4, 8, 15], 1, [4, 8, 15, 13]), + ([1, 20, 15], 1, [1, 20, 15, 4]), + ("1,4,15", 1, [1, 4, 15, 15, 5]), + ([1, 0, 15], 1, [1, 15, 15, 9]), + ([1], 1, [1, 15, 15, 9]), + ("invalid", 1, [1, 15, 15, 9]), + ], +) +def test_order_and_final_flush(ramp: object, first: int, expected: list[int]) -> None: + manager = _manager(ramp, first=first) + packets = [] + for index in range(40): + packet = _emit(manager, "a", index, finished=index == 39) + if packet is not None: + packets.append(packet) + assert [p.meta.codec_chunk_frames for p in packets] == expected + actual = torch.cat([p.codes.audio.reshape(4, -1).T for p in packets]) + assert torch.equal(actual, torch.arange(40).unsqueeze(1).expand(-1, 4)) + assert all(not bool(p.meta.finished) for p in packets[:-1]) + assert bool(packets[-1].meta.finished) + assert "a" not in manager.code_prompt_token_ids + + +def test_requests_progress_independently_and_empty_finish() -> None: + manager = _manager([1, 2, 4, 8, 15]) + for name in ("a", "b"): + first = _emit(manager, name, 0) + assert first is not None + assert first.meta.codec_chunk_frames == 1 + assert _emit(manager, "a", 1) is None + final_a = _emit(manager, "a", None, finished=True) + assert final_a is not None + assert final_a.meta.codec_chunk_frames == 1 + assert bool(final_a.meta.finished) + final_b = _emit(manager, "b", None, finished=True) + assert final_b is not None + assert final_b.meta.codec_chunk_frames == 0 + assert bool(final_b.meta.finished) + assert final_b.codes.audio.numel() == 0 + assert not manager.code_prompt_token_ids + + +@pytest.mark.parametrize( + "ramp,first,expected,max_step,effective_first", + [ + (None, 1, [1, 15], 15, 1), + ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15], 15, 1), + ([1, 20, 15], 1, [1, 15, 20], 20, 1), + ([4, 8, 15], 1, [4, 8, 15], 15, 4), + ([1, 4, 15], 7, [1, 4, 15], 15, 1), + ([1, 0, 15], 1, [1, 15], 15, 1), + ], +) +def test_graph_shapes_and_fast_path_follow_sender( + ramp: object, first: int, expected: list[int], max_step: int, effective_first: int +) -> None: + connector = _manager(ramp, first=first).connector.config + connector["extra"]["codec_first_chunk_fast_path"] = 1 + config = Mock( + spec=VllmConfig, + model_config=Mock( + spec=ModelConfig, + hf_config=PretrainedConfig(model_type="moss_tts_local"), + async_chunk=True, + use_v2_model_runner=True, + enforce_eager=False, + stage_connector_config=connector, + ), + scheduler_config=Mock(spec=SchedulerConfig, max_num_seqs=8), + compilation_config=Mock(spec=CompilationConfig, cudagraph_capture_sizes=[1, 2, 4, 8]), + ) + codec = MossTTSCodecDecoder(vllm_config=config) + assert codec._streaming_graph_frame_sizes == expected + assert codec._stream_max_step_frames == max_step + assert codec._first_chunk_fast + assert codec._initial_stream_chunk_frames == effective_first diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 9006d58c8..52b7d61dd 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -39,6 +39,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class MockConnector: + sender_host: str | None = None + sender_zmq_port: int | None = None + """In-memory connector for testing (mimics OmniConnectorBase).""" def __init__(self, stage_id: int = 0): @@ -187,7 +190,7 @@ def test_init_payload_connector_ownership(role, custom_func, expected): @pytest.mark.parametrize("async_chunk", [False, True]) @pytest.mark.parametrize("rank", [0, 1]) -def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): +def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank, mocker): host = MixinHost() host.init_omni_connectors(_make_model_config(async_chunk=async_chunk)) # No background threads: deterministic interleaved polling. @@ -196,7 +199,7 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): connector = host._omni_connector connector.sender_host = "configured-host" connector.sender_zmq_port = 50051 - connector.get = MagicMock(return_value=None) + get = mocker.patch.object(connector, "get", return_value=None) first = _make_request("r1", "external1") first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"} second = _make_request("r2", "external2") @@ -209,14 +212,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): for req_id in ("r2", "r1", "r2"): assert not host._poll_single_request(req_id) if rank == 1: - connector.get.assert_not_called() + get.assert_not_called() else: - assert connector.get.call_args_list == [ + assert get.call_args_list == [ unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}), unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}), unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}), ] - connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1) + get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1) assert host._poll_single_request("r2") assert "r2" not in host._pending_load_reqs assert "r1" in host._pending_load_reqs @@ -225,14 +228,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): assert not host._pending_load_reqs # Reused internal ID without sender info cannot inherit either endpoint. replacement = _make_request("r1", "external-new") - connector.get.reset_mock(return_value=True) - connector.get.return_value = None + get.reset_mock(return_value=True) + get.return_value = None host.register_chunk_recv(replacement) assert not host._poll_single_request("r1") if rank == 0: - connector.get.assert_called_once_with("0", "1", "external-new_0_0") + get.assert_called_once_with("0", "1", "external-new_0_0") else: - connector.get.assert_not_called() + get.assert_not_called() assert connector.sender_host == "configured-host" assert connector.sender_zmq_port == 50051 assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"} @@ -1437,6 +1440,58 @@ class TestAsyncPayloadLifecycle(unittest.TestCase): host.shutdown_omni_connectors() + def test_first_chunk_hook_sees_only_chunk_zero_before_staging(self): + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"), + ) + host._omni_connector = MagicMock() + host._stage_id = 1 + host._async_chunk = True + host._model_mode = "gen" + seen = [] + + def hook(req_id, handle, payload): + # The chunk is not yet visible to the scheduler when the hook runs. + seen.append((req_id, handle.external_req_id, payload["codes"]["audio"], "r1" in host._finished_load_reqs)) + + host.set_first_chunk_hook(hook) + host.register_chunk_recv(_make_request("r1", "ext-r1")) + chunk = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}} + host._omni_connector.get.return_value = (chunk, 1) + self.assertTrue(host._poll_single_request("r1")) + self.assertEqual(seen, [("r1", "ext-r1", [1, 2], False)]) + self.assertIn("r1", host._finished_load_reqs) + + host.get_omni_connector_output() + host._local_stage_payload_cache.clear() + host._local_request_metadata.clear() + host._omni_connector.get.return_value = ({"codes": {"audio": [3, 4]}, "meta": {}}, 1) + self.assertTrue(host._poll_single_request("r1")) + self.assertEqual(len(seen), 1) + + host.shutdown_omni_connectors() + + def test_first_chunk_hook_skips_terminal_first_chunk(self): + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"), + ) + host._omni_connector = MagicMock() + host._stage_id = 1 + host._async_chunk = True + host._model_mode = "gen" + hook = MagicMock() + host.set_first_chunk_hook(hook) + host.register_chunk_recv(_make_request("r1")) + host._omni_connector.get.return_value = ( + {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(True)}}, + 1, + ) + self.assertTrue(host._poll_single_request("r1")) + hook.assert_not_called() + host.shutdown_omni_connectors() + def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self): host = MixinHost() host.init_omni_connectors( @@ -1480,6 +1535,39 @@ class TestAsyncPayloadLifecycle(unittest.TestCase): host.shutdown_omni_connectors() +@pytest.mark.parametrize("cancel_at", ["connector_get", "first_chunk_hook"]) +def test_cancelled_chunk_receive_does_not_recreate_delivery_state(mocker, cancel_at): + host = MixinHost() + host.init_omni_connectors(_make_model_config(async_chunk=True, worker_type="gen")) + host._stage_id = 1 + host._omni_connector = mocker.MagicMock() + host.register_chunk_recv(_make_request("r1")) + seen = [] + payload = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}} + + def get(*args, **kwargs): + if cancel_at == "connector_get": + host.cleanup_finished_request("r1") + return payload, 1 + + def hook(req_id, handle, chunk): + seen.append(req_id) + host.cleanup_finished_request(req_id) + + host._omni_connector.get.side_effect = get + host.set_first_chunk_hook(hook) + try: + assert not host._poll_single_request("r1") + assert seen == (["r1"] if cancel_at == "first_chunk_hook" else []) + assert host.get_local_stage_payload("r1") is None + assert "r1" not in host._get_req_chunk + assert "r1" not in host._pending_load_reqs + output = host.get_omni_connector_output() + assert output.chunk_ready_req_ids == set() and output.request_metadata == {} + finally: + host.shutdown_omni_connectors() + + class TestRankAwareKVRouting(unittest.TestCase): def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost: host = MixinHost() diff --git a/tests/worker_v2/test_first_audio_sender.py b/tests/worker_v2/test_first_audio_sender.py index 181f01d8b..8734dbf1e 100644 --- a/tests/worker_v2/test_first_audio_sender.py +++ b/tests/worker_v2/test_first_audio_sender.py @@ -235,3 +235,19 @@ def test_sender_copy_failure_reports_all_prepared_routes(): assert client == 3 assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)] assert output_queue.empty() + + +def test_final_codec_delivery_omits_upstream_marker_and_freezes_route(): + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + + outputs: queue.Queue = queue.Queue() + scheduler = _scheduler(r=3) + delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"]) + scheduler.requests.clear() + pcm = torch.ones(4) + delivery(["r"], [pcm], torch.tensor(24000)) + client, batch = outputs.get_nowait() + assert client == 3 + assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output + assert batch.outputs[0].multimodal_output["model_outputs"] is pcm + assert batch.outputs[0].finish_reason is None diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index 0bbebf491..cd0c09877 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -15,6 +15,7 @@ from vllm.v1.outputs import RoutedExpertsTensors from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot @@ -144,6 +145,7 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: runner.model = SimpleNamespace() runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 + runner._last_multimodal_snapshot_slot = None waited: list[object] = [] runner.main_stream = SimpleNamespace(wait_event=waited.append) source = torch.tensor([[7, 8]], dtype=torch.long) @@ -161,6 +163,35 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: assert waited == [copy_event] +def test_producer_snapshot_is_not_repacked() -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace() + runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] + runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 + runner._last_multimodal_snapshot_slot = None + source = torch.tensor([[7, 8]], dtype=torch.float32) + snapshot = pack_output_snapshot({"audio": source}, {}, max_buckets=1) + assert isinstance(snapshot, PackedOutputSnapshot) + + retained = runner._retain_multimodal_outputs(snapshot) + + assert retained is snapshot + assert runner._last_multimodal_snapshot_slot is None + assert runner._async_mm_snapshot_pending == [False] + + +def test_fresh_per_step_outputs_are_not_repacked() -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace(mm_outputs_fresh_per_step=True) + runner._last_multimodal_snapshot_slot = None + outputs = {"model_outputs": torch.ones(2, 4)} + + assert runner._retain_multimodal_outputs(outputs) is outputs + assert runner._last_multimodal_snapshot_slot is None + + def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None: slot: dict[tuple[Any, ...], torch.Tensor] = {} omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot) diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 1cbef67e8..86764b691 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -224,6 +224,61 @@ def test_static_decode_embeddings_refresh_from_input_ids(): assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True +def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker): + """Exercise MRV2 dispatch through the real Local hook, output and logits.""" + from vllm.sampling_params import SamplingParams + from vllm.v1.worker.gpu.states import RequestState + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSLocalTalkerForGeneration, + ) + + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + torch.nn.Module.__init__(model) + model.model = torch.nn.Module() + model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4)) + model.hidden_size = 4 + model.n_vq = 2 + model.audio_pad_token_id = 8 + model.audio_assistant_slot_token_id = 2 + model.im_end_token_id = 3 + model.text_vocab_size = 8 + model.talker_mtp_output_key = ("audio_codes", "current") + model.talker_mtp_graph_safe = False + model.gpu_resident_buffer_keys = set() + model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None + model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float() + frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]]))) + model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame) + + state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True) + state.model = model + for idx in range(2): + state.intermediate_buffer.buffers[idx] = { + "req_id": f"r{idx}", + "audio_state": {"is_stopping": False}, + "hidden_states": {"last": torch.full((4,), float(idx + 1))}, + } + batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1]) + inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)} + request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None) + state.run_preprocess(batch, inputs, request_state) + + frame.assert_called_once() + torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4])) + assert frame.call_args.kwargs["temperature"] == 1.7 + assert frame.call_args.kwargs["top_k"] == 25 + assert frame.call_args.kwargs["top_p"] == 0.8 + torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4])) + _, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state) + assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]] + assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3] + # An explicit local seed must reach the request-owned MRV2 generator. + params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99) + assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17 + + @pytest.mark.parametrize("owned", [False, True]) def test_batched_postprocess_gpu_snapshot_writeback(owned): # batch=[1, 0]: last-token indices follow the reordered batch, the scalar diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 06dcb3ef4..965e8d340 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -182,6 +182,24 @@ _MODEL_PROMOTE = ( "VLLM_GEPARD_FIRST_CHUNK_FRAMES", "VLLM_GEPARD_LOOKBACK_FRAMES", "VLLM_OMNI_FISH_KVCACHE_ATTN", + "VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", + "VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", + "VLLM_OMNI_MOSS_REF_ATTN", + "VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", + "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS", + "VLLM_OMNI_MOSS_REF_CACHE_ROPE", + "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", + "VLLM_OMNI_MOSS_REF_COMPILE", + "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS", + "VLLM_OMNI_MOSS_REF_FA_VERSION", + "VLLM_OMNI_MOSS_REF_GPU_PREP", + "VLLM_OMNI_MOSS_REF_GRAPHS", + "VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", + "VLLM_OMNI_MOSS_REF_INFLIGHT", + "VLLM_OMNI_MOSS_REF_SHARED_ENCODER", + "VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", + "VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", + "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY", "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER", "VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS", "VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES", diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cefa2a37a..b103776fa 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -3,7 +3,9 @@ from __future__ import annotations +import math import os +import queue import time from collections import defaultdict from typing import TYPE_CHECKING, Any @@ -58,6 +60,48 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._init_omni_io_scheduling_state() self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] + connector_config = getattr(model_config, "stage_connector_config", None) + extra = ( + connector_config.get("extra", connector_config) + if isinstance(connector_config, dict) + else getattr(connector_config, "extra", None) + ) or {} + self._generation_min_batch_size = int(extra.get("generation_min_batch_size", 1)) + self._generation_max_wait_s = float(extra.get("generation_max_wait_ms", 0)) / 1000 + self._generation_coalescing_policy = extra.get("generation_coalescing_policy", "fixed") + self._generation_max_regular_batch = int(extra.get("generation_max_regular_batch", 0) or 0) + if self._generation_max_regular_batch < 0: + raise ValueError("generation_max_regular_batch must be nonnegative") + if self._generation_coalescing_policy not in ("fixed", "idle_wait"): + raise ValueError("generation_coalescing_policy must be fixed or idle_wait") + if not 1 <= self._generation_min_batch_size <= self.max_num_running_reqs: + raise ValueError("generation_min_batch_size must be in [1, max_num_seqs]") + if not math.isfinite(self._generation_max_wait_s) or self._generation_max_wait_s < 0: + raise ValueError("generation_max_wait_ms must be finite and nonnegative") + if self._generation_max_wait_s and self._generation_min_batch_size > 1: + parallel = self.vllm_config.parallel_config + if not self._native_data_plane: + # CUDA-oriented profiles can inherit a platform's V1 runner. + # Their connector extras survive that override, but batching + # must revert to the fallback runner's default scheduling. + logger.warning( + "Generation batch waiting and regular-batch limit ignored: " + "requires native MRV2; using default generation scheduling." + ) + self._generation_min_batch_size = 1 + self._generation_max_wait_s = 0.0 + self._generation_max_regular_batch = 0 + elif not ( + self._retains_state_across_chunks + and parallel.tensor_parallel_size == parallel.pipeline_parallel_size == 1 + ): + raise ValueError("Generation batch waiting requires a stateful native MRV2 TP1/PP1 stage") + else: + logger.info( + "Generation input coalescing: target batch=%d, max wait=%.3f ms", + self._generation_min_batch_size, + self._generation_max_wait_s * 1000, + ) self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1" self._last_step_express = False # Streams that already had a chunk scheduled; the rest await their first. @@ -76,6 +120,101 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._express_min_slack_s, ) + def _drain_omni_connector_outputs(self): + self._generation_defer_batch = False + outputs = super()._drain_omni_connector_outputs() + max_wait = getattr(self, "_generation_max_wait_s", 0) + target = getattr(self, "_generation_min_batch_size", 1) + coordinator = self.input_coordinator + if not max_wait or target <= 1 or coordinator is None or self._pause_state != PauseState.UNPAUSED: + return outputs + + # Count distinct runnable streams, not notifications, padding rows, or + # chunks whose previous execution still owns their state. Waiting here + # precedes allocation and metadata replacement. Do not schedule a batch + # and then discard it to try to accumulate more input. + ready = set(coordinator.requests_with_ready_chunks) + terminal = set(coordinator.finished_requests) | set(coordinator.input_terminal_req_ids) + deadline = getattr(self, "_generation_batch_deadline", None) + if deadline is None: + deadline = time.monotonic() + max_wait + in_flight = any( + not request.is_finished() and request.num_in_flight_tokens > 0 for request in self.requests.values() + ) + new_outputs = outputs + waited_without_input = False + while True: + for output in new_outputs: + ready.update(output.chunk_ready_req_ids) + terminal.update(output.chunk_finished_req_ids) + terminal.update( + req_id for req_id, metadata in output.request_metadata.items() if metadata.get("input_terminal") + ) + runnable = { + req_id + for req_id in ready | terminal + if (request := self.requests.get(req_id)) is not None + and not request.is_finished() + and request.num_in_flight_tokens == 0 + } + if runnable and waited_without_input: + # Idle waiting must not eat the ready batch's coalescing budget. + # Start that budget at its first arrival, once only. The two + # bounded windows can delay control processing by at most 2W. + deadline = time.monotonic() + max_wait + waited_without_input = False + # A first/final chunk participates in the same bounded window. + # At high concurrency almost every batch contains one; bypassing + # the entire batch in that case defeats coalescing. Low concurrency + # still dispatches immediately, including single-request tails. + if not runnable: + # With no output to retire and all receivers already parked, + # a wakeable inbox wait avoids repeatedly executing empty steps. + # Admission/registration/terminal work must reach the runner + # first; waiting for a not-yet-registered receiver would stall it. + can_wait_idle = ( + getattr(self, "_generation_coalescing_policy", "fixed") == "idle_wait" + and not in_flight + and bool(self.requests) + and not getattr(coordinator, "pending_chunk_registrations", ()) + and not getattr(coordinator, "pending_input_registrations", ()) + and not getattr(self, "_pending_data_plane_terminal_req_ids", ()) + and all(request.status == RequestStatus.WAITING_FOR_CHUNK for request in self.requests.values()) + ) + if not can_wait_idle: + break + waited_without_input = True + elif len(runnable) >= target or len(self.requests) < target: + break + remaining = deadline - time.monotonic() + if remaining <= 0: + break + if in_flight: + # Let EngineCore retire the prior output before waiting for + # more input. Blocking this thread also blocks the in-flight + # fence that makes those streams eligible for the next batch. + # Emit a normal empty scheduler step (including connector and + # cancellation bookkeeping), never discard an allocated step. + self._generation_defer_batch = True + self._generation_batch_deadline = deadline + return outputs + try: + output = self._omni_connector_output_inbox.get(timeout=remaining) + except queue.Empty: + break + outputs.append(output) + new_outputs = [output] + # Drain a burst before re-evaluating, retaining notification order. + while True: + try: + output = self._omni_connector_output_inbox.get_nowait() + except queue.Empty: + break + outputs.append(output) + new_outputs.append(output) + self._generation_batch_deadline = None + return outputs + def _build_generation_scheduler_output( self, new_reqs_data: list[OmniNewRequestData], @@ -256,6 +395,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._drop_aborted_queued_requests() self._requeue_completed_native_chunks() self._process_pending_omni_inputs(model_mode="generation") + if getattr(self, "_generation_defer_batch", False): + token_budget = 0 + max_regular = getattr(self, "_generation_max_regular_batch", 0) + if max_regular: + execution_batch_size = min(execution_batch_size, max_regular) self._drop_aborted_queued_requests() self._resync_streaming_input_counter() async_chunk_transport = self._async_chunk_transport_enabled() diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 6f9f03d6c..105b79a00 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -331,13 +331,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # ready-chunk marker remains until scheduler admission succeeds. replaced_ids.discard(request_id) - def _consume_pending_connector_output(self, model_mode: str) -> None: - """Drain ``self._latest_omni_connector_output`` into the coordinator. - - Called at the top of every ``schedule()`` cycle. Identical between - AR and generation schedulers except for the ``model_mode`` argument - forwarded to ``update_request_metadata``. - """ + def _drain_omni_connector_outputs(self) -> list[OmniConnectorOutput]: + """Collect control messages in order, without changing request state.""" connector_outputs: list[OmniConnectorOutput] = [] inbox = getattr(self, "_omni_connector_output_inbox", None) if inbox is not None: @@ -350,6 +345,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase): self._latest_omni_connector_output = None if connector_output is not None: connector_outputs.append(connector_output) + return connector_outputs + + def _consume_pending_connector_output(self, model_mode: str) -> None: + """Drain input notifications into the coordinator on the scheduler thread.""" + connector_outputs = self._drain_omni_connector_outputs() input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is None: return diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index dee3a9b43..f3042b736 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -379,14 +379,17 @@ class OmniSchedulingCoordinator: ) if model_mode != "ar": - new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes")) + new_codes = metadata.get("code_predictor_codes") + new_ids = self._flatten_prompt_token_ids(new_codes) runtime_seed = None if "left_context_size" in metadata: runtime_seed = { "meta": {"left_context_size": metadata["left_context_size"]}, } request._omni_initial_model_buffer = runtime_seed - if new_ids: + # An explicit empty snapshot clears a previous codec chunk + # or prewarmed placeholder. Missing codes leave it intact. + if new_codes is not None: request.prompt_token_ids = new_ids request.num_prompt_tokens = len(new_ids) request._all_token_ids.clear() diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 77c2a42b6..c351041f5 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -18,6 +18,9 @@ connectors: connector_get_max_wait: 300 initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks + # Optional startup ramp; its first entry overrides initial chunk size. + # Extra graph shapes/calls trade codec cost and latency for less underrun. + # codec_chunk_ramp: [1, 4, 15] codec_left_context_frames: 0 # Chunk ramp (opt-in): gradual chunk-size ladder replacing the # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last diff --git a/vllm_omni/deploy/moss_tts_local_mrv2.yaml b/vllm_omni/deploy/moss_tts_local_mrv2.yaml new file mode 100644 index 000000000..4b25a41c7 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2.yaml @@ -0,0 +1,30 @@ +# Opt-in CUDA MRV2 for MOSS-TTS-Local-Transformer-v1.5. +# Keep the base profile's batching, chunk geometry and V1 default. +base_config: moss_tts_local.yaml +model_runner: v2 + +platforms: + cuda: + stages: + - stage_id: 0 + # Bound prefill work so new prompts do not monopolize streaming decode. + max_num_batched_tokens: 512 + - stage_id: 1 + compilation_config: + # The codec owns CUDA Graph capture/replay. Skip Inductor compilation + # of its stateful decoder; plain CUDA graphs remain enabled. + mode: 0 + # FULL works without Inductor; the default FULL_AND_PIECEWISE would + # normalize to NONE and erase the codec's capture sizes. + cudagraph_mode: FULL + # Platform overrides replace compilation_config as a whole. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] + cudagraph_num_of_warmups: 1 + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml new file mode 100644 index 000000000..eb155f284 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml @@ -0,0 +1,53 @@ +# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 on a large-memory GPU. +# Both stages share GPU 0. Validated with one H200 and Seed-TTS EN1088. +base_config: moss_tts_local_mrv2.yaml + +connectors: + shm: + extra: + # Coalesce ready codec streams before submitting a small execution batch. + # This is a dispatch target, not a cap: execution capacity remains 128. + generation_min_batch_size: 32 + generation_max_wait_ms: 12 + # Optional; preserves the first/steady sizes while adding a middle chunk. + # codec_chunk_ramp: [1, 4, 15] + +# Optional codec attention-ring headroom (see recipes/OpenMOSS/MOSS-TTS.md): +# sizes each ring as context + one chunk so chunked streaming matches +# whole-sequence decoding. Requires additional state/graph memory; evaluate +# throughput and audio quality before enabling for your context and capacity. +# Add codec_ring_headroom: 1 under connectors.shm.extra above. + +platforms: + cuda: + stages: + - stage_id: 0 + max_num_seqs: 128 + max_num_batched_tokens: 512 + hf_overrides: + mrv2_gpu_slot_state: true + # Throughput-oriented pair; measure first-packet latency for your load. + mrv2_batch_prefill: true + mrv2_direct_tokens: true + # Mixed prefill/decode fits the same bounded token budget as capture. + compilation_config: + mode: 3 + cudagraph_mode: FULL + max_cudagraph_capture_size: 512 + cudagraph_num_of_warmups: 1 + # Leave room for the codec state pool and its batch-128 CUDA graphs. + kv_cache_memory_bytes: 34359738368 # 32 GiB + - stage_id: 1 + max_num_seqs: 128 + hf_overrides: + codec_attention_backend: triton_slot + compilation_config: + mode: 3 + backend: inductor + # The codec owns capture/replay after Inductor compiles its decoder. + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + cudagraph_num_of_warmups: 1 + inductor_compile_config: + combo_kernels: false + benchmark_combo_kernel: false diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml new file mode 100644 index 000000000..959238ee2 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml @@ -0,0 +1,66 @@ +# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 tuned for first-packet +# latency at client concurrency 128 on one H200 (Seed-TTS EN1088, closed loop). +# Both stages share GPU 0. Use moss_tts_local_mrv2_high_concurrency.yaml for +# throughput-first serving. +# +# Optional launch-level settings (see the full command in recipes/OpenMOSS/MOSS-TTS.md): +# - CUDA MPS for the GPU (stage 0 below lowers its own MPS client priority); +# - VLLM_OMNI_EVENT_DRIVEN_ORCH=1; +# - --api-server-count 4 with VLLM_WORKER_MULTIPROC_METHOD=spawn and +# VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/ so frontends share reference codes; +# - VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 so Stage 1 picks up first chunks sooner. +base_config: moss_tts_local_mrv2.yaml + +connectors: + shm: + extra: + generation_min_batch_size: 32 + generation_max_wait_ms: 12 + # Stage 1 decodes each first chunk from its receive thread on a small + # dedicated graph set and emits it directly, so first chunks no longer + # wait for the engine loop or a regular codec step. + codec_first_chunk_fast_path: 1 + # Regular codec replays wait for an in-flight first-chunk decode. + codec_first_chunk_gate: 1 + # Short regular codec steps bound contention with first-chunk decodes. + generation_max_regular_batch: 16 + # Optional ramp. Both the first-chunk graph and regular graphs follow it. + # codec_chunk_ramp: [1, 4, 15] + +platforms: + cuda: + stages: + - stage_id: 0 + max_num_seqs: 128 + # Larger prefill budget admits bursts of new requests in fewer steps. + max_num_batched_tokens: 2048 + hf_overrides: + mrv2_gpu_slot_state: true + mrv2_batch_prefill: true + mrv2_direct_tokens: true + # Sample each frame right after the forward that produced it. + mrv2_eager_mtp: true + compilation_config: + mode: 3 + cudagraph_mode: FULL + max_cudagraph_capture_size: 2048 + cudagraph_num_of_warmups: 1 + kv_cache_memory_bytes: 34359738368 # 32 GiB + env: + # Under MPS, let codec kernels schedule ahead of the talker's. + CUDA_MPS_CLIENT_PRIORITY: "1" + - stage_id: 1 + max_num_seqs: 128 + # Keep synchronous scheduling; benchmark async scheduling separately. + async_scheduling: false + hf_overrides: + codec_attention_backend: triton_slot + compilation_config: + mode: 3 + backend: inductor + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + cudagraph_num_of_warmups: 1 + inductor_compile_config: + combo_kernels: false + benchmark_combo_kernel: false diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 7d2b16d50..d7744ffb1 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -1169,7 +1169,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._payload_finished(payload_data), ) - self._get_req_chunk[req_id] += 1 + with self._lock: + if self._async_chunk and request is not None and req_id not in self._pending_load_reqs: + # A connector get can return after cancellation removed its + # receiver. Do not recreate delivery state for that request. + return False + self._get_req_chunk[req_id] += 1 if self._async_chunk: is_finished = self._payload_finished(payload_data) @@ -1180,8 +1185,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if not self._payload_value_has_content(new_ids) and not is_finished: return False payload_consumable = self._payload_is_consumable(payload_data) + first_chunk_hook = getattr(self, "_first_chunk_hook", None) + if first_chunk_hook is not None and chunk_id == 0 and not is_finished and payload_consumable: + first_chunk_hook(req_id, request, payload_data) with self._lock: + if request is not None and req_id not in self._pending_load_reqs: + # The first-chunk hook may have overlapped cancellation. + # Its slot cleanup is ordered by the model; this payload + # must not republish readiness after receive teardown. + return False if self._model_mode == "ar": # Accumulation, staging, and model-side consume/ack share # this lock. Keeping the transition atomic prevents the @@ -1723,3 +1736,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): sink: Callable[[OmniConnectorOutput], None] | None, ) -> None: self._omni_connector_output_sink = sink + + def set_first_chunk_hook(self, hook: Callable[[str, Any, OmniPayload], bool]) -> None: + """Claim eligible first chunks before publishing them to the model loop.""" + self._first_chunk_hook = hook diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index f9fcd34b9..c84d9ed18 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -80,9 +80,12 @@ def _signal_exit_code(signum: int) -> int: def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: Any) -> bool: - """Let a TP1 in-process runner that decodes first audio deliver it as its own output.""" + """Bind in-process first audio to the generic engine output sink.""" if not isinstance(model_executor, UniProcExecutor): return False + parallel_config = model_executor.vllm_config.parallel_config + if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1: + return False worker = getattr(getattr(model_executor, "driver_worker", None), "worker", None) model_runner = getattr(worker, "model_runner", None) model_state = getattr(model_runner, "model_state", None) @@ -90,12 +93,23 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An decodes_audio = getattr(model, "first_frame_decoder", None) is not None or ( getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False)) ) - if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"): - return False from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink - assert model_state is not None - model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler)) + if decodes_audio and hasattr(model_state, "set_first_audio_sink"): + assert model_state is not None + model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler)) + return True + data_plane = getattr(model_runner, "_omni_data_plane", None) + get_model = getattr(model_runner, "get_model", None) + if data_plane is None or not callable(get_model): + return False + bind = getattr(get_model(), "bind_first_chunk_fast_path", None) + if not callable(bind): + return False + hook = bind(engine_output_queue_sink(output_queue, scheduler, upstream_first_audio=False)) + if hook is None: + return False + data_plane.set_first_chunk_hook(hook) return True diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 79217ca8b..ca86587e9 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -441,7 +441,7 @@ class OmniBase(PDDisaggregationMixin): if self.log_stats: # Emit per-request orchestrator timing (including e2e_total_ms) # before dropping request state. - req_state.metrics.build_and_log_summary() + req_state.metrics.log_timing_summary() except Exception: logger.exception( "[%s] Failed to build/log summary for req=%s", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 6ad7dcdcc..89eb463d9 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -81,13 +81,24 @@ def conditioning_cache_salt( This replaces the uploaded data URI in the hash without changing the request. """ h = hashlib.sha256() + ref_audio: str | list[str] | tuple[str, int] | None = request.ref_audio + if ( + isinstance(ref_audio, str) + and len(ref_audio) > 256 + and tts_params is not None + and tts_params.get("ref_audio_cache_key") is not None + ): + # The resolve key below is already a content-aware digest of this + # locator. Hashing a repr of a large inline data URI again dominates + # the per-request serving cost without adding information. + ref_audio = ("ref_audio_cache_key", len(ref_audio)) for part in ( request.input, request.task_type, request.language, request.voice, request.ref_text, - registered_voice if registered_voice is not None else request.ref_audio, + registered_voice if registered_voice is not None else ref_audio, request.instructions, request.x_vector_only_mode, request.speaker_embedding, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index db2833f42..ed7865789 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -32,6 +32,59 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _local_continuation_prompt(proc: Any, user_kwargs: dict[str, Any], reference: Any) -> tuple[list[int], Any] | None: + """Build the Local-v1.5 continuation prompt without the generic processor call. + + Mirrors ``_build_continuation_codes`` of the model's processor: every + template segment is encoded separately there, so the constant segments + (keyed by ``language``) are cached and only the request text is + tokenized. Returns ``(text_ids, audio_codes)`` equal to + ``unified[:, 0].tolist()`` / ``unified[:, 1:]``, or ``None`` when the + request or processor does not match that exact layout. + """ + import torch + + if set(user_kwargs) - {"text", "language"}: + return None + if not isinstance(reference, list) or len(reference) != 1 or not isinstance(reference[0], torch.Tensor): + return None + cfg = getattr(proc, "model_config", None) + n_vq = getattr(cfg, "n_vq", None) + if cfg is None or not isinstance(n_vq, int) or n_vq <= 0: + return None + codes = reference[0] + if codes.ndim != 2 or int(codes.shape[1]) != n_vq or codes.shape[0] == 0: + return None + required = ( + "_user_prompt_prefix_ids", + "_encode_text", + "_user_prompt_after_reference_ids", + "_assistant_prompt_prefix_ids", + ) + if not all(callable(getattr(proc, name, None)) for name in required): + return None + language = user_kwargs.get("language") + cache = proc.__dict__.setdefault("_vllm_omni_continuation_segments", {}) + segments = cache.get(language) + if segments is None: + fields = {} if language is None else {"language": language} + head = ( + proc._user_prompt_prefix_ids() + + proc._encode_text("None") + + proc._user_prompt_after_reference_ids(language, fields) + ) + tail = proc._assistant_prompt_prefix_ids() + [int(cfg.audio_start_token_id)] + segments = cache[language] = ([int(t) for t in head], [int(t) for t in tail]) + head, tail = segments + text = user_kwargs.get("text") + prompt_ids = head + [int(t) for t in proc._encode_text("" if text is None else str(text))] + tail + num_audio = int(codes.shape[0]) + text_ids = prompt_ids + [int(cfg.audio_assistant_slot_token_id)] * num_audio + audio_codes = torch.full((len(text_ids), n_vq), int(cfg.audio_pad_token_id), dtype=torch.int64) + audio_codes[len(prompt_ids) :] = codes.to(dtype=torch.int64, device="cpu") + return text_ids, audio_codes + + class _MossTTSAdapterBase(ARTTSAdapter): accumulate_nonstreaming: bool = False @@ -82,6 +135,28 @@ class _MossTTSAdapterBase(ARTTSAdapter): self._moss_ref_encoder = encoder return encoder + async def warmup(self) -> None: + """Set up the reference encoder at startup, not on the first cloning request. + + Its CUDA graph capture takes seconds per API process (about a minute for + a shared encoder's workers). + """ + if self._moss_variant not in ("tts", "local", "ttsd"): + return + from vllm_omni.model_executor.models.moss_tts.reference_encoder import reference_graphs_enabled + from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError + + if not reference_graphs_enabled(): + return + try: + await asyncio.to_thread(self._get_moss_ref_encoder().prepare) + except SharedReferenceEncoderStartupError: + # Do not advertise readiness and trigger concurrent local captures + # while the shared host is still compiling or has failed startup. + raise + except Exception: # noqa: BLE001 - requests still encode (and capture) on demand + logger.warning("MOSS reference encoder warmup failed", exc_info=True) + async def _encode_moss_references( self, request: "OpenAICreateSpeechRequest", @@ -269,8 +344,16 @@ class _MossTTSAdapterBase(ARTTSAdapter): proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) if hasattr(proc, "audio_tokenizer"): device = self._resolve_ref_encoder_device() - proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval() - logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device) + proc._vllm_omni_ref_encoder_device = device + from vllm_omni.model_executor.models.moss_tts.reference_encoder import shared_encoder_role + + if shared_encoder_role() == "client": + # Another API process hosts the shared encoder on the GPU. + proc.audio_tokenizer = proc.audio_tokenizer.eval() + logger.info("MOSS reference-audio encoder: using the shared encoder of another API process") + else: + proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval() + logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device) self._moss_processor_cache = proc return proc @@ -405,18 +488,24 @@ 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). + fast = None 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") + fast = _local_continuation_prompt(proc, user_kwargs, reference) + if fast is None: + 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) + if fast is not None: + text_ids, audio_codes = fast + else: + unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq) + text_ids = unified[:, 0].tolist() + audio_codes = unified[:, 1:].contiguous().to(torch.int64) params = { "prompt_token_ids": text_ids, diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index df5183f45..edb9a87db 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy +import logging import time from collections import defaultdict from collections.abc import Callable, Sequence @@ -851,9 +852,62 @@ class OrchestratorAggregator: for cached_key in [k for k in self._replica_cache if k[1] == rid_key]: self._replica_cache.pop(cached_key, None) + def log_timing_summary(self) -> None: + """Emit only the concise per-request ``[OmniTiming]`` lines. + + The serving path uses this unless DEBUG tables are requested; it skips + the table rows that ``build_and_log_summary`` builds for its result. + """ + if not self.log_stats: + return + if logger.isEnabledFor(logging.DEBUG): + self.build_and_log_summary() + return + for rid in sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events}): + e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None) + stage_evts = sorted( + self.stage_events.get(rid, []), + key=lambda e: e.stage_id if e.stage_id is not None else -1, + ) + self._log_omni_timing(rid, e2e_evt, stage_evts) + + def _log_omni_timing(self, rid: str, e2e_evt: RequestE2EStats | None, stage_evts: list[StageRequestStats]) -> None: + pt: dict[str, float] = {} + if stage_evts: + pt = stage_evts[-1].pipeline_timings or {} + if pt or e2e_evt: + parts = [f"req={rid}"] + if e2e_evt: + parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s") + if "preprocess_ms" in pt: + parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s") + if e2e_evt: + engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0) + parts.append(f"engine={engine_ms / 1000.0:.2f}s") + stage_parts = [] + for evt in stage_evts: + sid = evt.stage_id if evt.stage_id is not None else "?" + t = evt.stage_gen_time_ms / 1000.0 + stage_parts.append(f"{sid}:{t:.2f}s") + if stage_parts: + parts.append(f"stages=[{','.join(stage_parts)}]") + transfer_parts = [] + for te in self.transfer_events.values(): + if te.request_id == rid: + transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms") + if transfer_parts: + parts.append(f"transfers=[{','.join(transfer_parts)}]") + if "ar2diffusion_ms" in pt: + parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms") + logger.info("[OmniTiming] %s", " ".join(parts)) + def build_and_log_summary(self) -> dict[str, Any]: if not self.log_stats: return {} + # Per-request tables are diagnostic. Formatting them on the serving + # event loop costs more than the request's own bookkeeping, so they + # are emitted at DEBUG; the concise [OmniTiming] line stays at INFO. + log_tables = logger.isEnabledFor(logging.DEBUG) wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0 @@ -895,8 +949,8 @@ class OrchestratorAggregator: v = overall_summary.get(k, None) if v not in (0, 0.0, 0.000, None, ""): overall_fields.append(k) - if overall_fields: - logger.info( + if log_tables and overall_fields: + logger.debug( "\n%s", _format_table("Overall Summary", overall_summary, overall_fields), ) @@ -921,8 +975,8 @@ class OrchestratorAggregator: nonzero_e2e_fields.add(k) value_fields_e2e = sorted(nonzero_e2e_fields) - if value_fields_e2e: - logger.info( + if log_tables and value_fields_e2e: + logger.debug( "\n%s", _format_table( f"RequestE2EStats [request_id={rid}]", @@ -936,34 +990,7 @@ class OrchestratorAggregator: self.stage_events.get(rid, []), key=lambda e: e.stage_id if e.stage_id is not None else -1, ) - pt = {} - if stage_evts: - pt = stage_evts[-1].pipeline_timings or {} - if pt or e2e_evt: - parts = [f"req={rid}"] - if e2e_evt: - parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s") - if "preprocess_ms" in pt: - parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s") - if e2e_evt: - engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0) - parts.append(f"engine={engine_ms / 1000.0:.2f}s") - stage_parts = [] - for evt in stage_evts: - sid = evt.stage_id if evt.stage_id is not None else "?" - t = evt.stage_gen_time_ms / 1000.0 - stage_parts.append(f"{sid}:{t:.2f}s") - if stage_parts: - parts.append(f"stages=[{','.join(stage_parts)}]") - transfer_parts = [] - for te in self.transfer_events.values(): - if te.request_id == rid: - transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms") - if transfer_parts: - parts.append(f"transfers=[{','.join(transfer_parts)}]") - if "ar2diffusion_ms" in pt: - parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms") - logger.info("[OmniTiming] %s", " ".join(parts)) + self._log_omni_timing(rid, e2e_evt, stage_evts) # === Stage table (columns = stage_id) === # if any stage has diffusion_metrics, remove postprocess_time_ms field @@ -987,7 +1014,7 @@ class OrchestratorAggregator: result_stage_table.append({"request_id": rid, "stages": stage_rows}) - if stage_rows: + if log_tables and stage_rows: # filter out all-zero fields for logging all_value_fields = set() for row in stage_rows: @@ -1006,7 +1033,7 @@ class OrchestratorAggregator: value_fields_list.append(field) if value_fields_list: - logger.info( + logger.debug( "\n%s", _format_table( f"StageRequestStats [request_id={rid}]", @@ -1027,7 +1054,7 @@ class OrchestratorAggregator: ] result_trans_table.append({"request_id": rid, "transfers": transfer_rows}) - if transfer_rows: + if log_tables and transfer_rows: # filter out all-zero fields for logging all_value_fields = set() for row in transfer_rows: @@ -1046,7 +1073,7 @@ class OrchestratorAggregator: value_fields_list.append(field) if value_fields_list: - logger.info( + logger.debug( "\n%s", _format_table( f"TransferEdgeStats [request_id={rid}]", 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 4dbedbd19..89f32ee55 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 @@ -539,11 +539,15 @@ class RingKVCache: end_offset = self.end_offset.index_select(0, slots) row_cache = self.cache.index_select(1, slots) - indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype) + # Upsampling can make a chunk longer than the ring. Only the last + # capacity tokens survive; writing all T tokens gives scatter_ + # duplicate destinations and nondeterministic cache contents. + write_length = min(T, self.capacity) + indexes = torch.arange(T - write_length, T, device=end_offset.device, dtype=end_offset.dtype) indexes = (indexes + end_offset.view(-1, 1)) % self.capacity - scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) - row_cache[0].scatter_(2, scatter_indexes, k) - row_cache[1].scatter_(2, scatter_indexes, v) + scatter_indexes = indexes.view(B, 1, write_length, 1).expand(-1, H, write_length, D) + row_cache[0].scatter_(2, scatter_indexes, k[:, :, -write_length:]) + row_cache[1].scatter_(2, scatter_indexes, v[:, :, -write_length:]) # Live and graph-padding rows always map to distinct slots. The # latter map only to scratch state, so this write cannot corrupt a # request even though dense graph operators still execute it. @@ -729,7 +733,13 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.context is None: capacity = self.weights_per_step if self.weights_per_step else 1024 else: - capacity = self.context + # RingKVCache.complete writes a whole chunk before attending, so a + # ring of exactly ``context`` entries lets a T-token chunk evict up + # to T tokens of history that the chunk's own queries should still + # see (and, for T > context, evict part of the chunk itself). + # ``_ring_headroom`` (max chunk tokens at this layer) keeps every + # in-window key resident; 0 preserves the legacy truncation. + capacity = self.context + int(getattr(self, "_ring_headroom", 0)) kv_cache = RingKVCache( batch_size, @@ -791,29 +801,63 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.rope: q, k = self.rope(q, k, offset, time_before_heads=False) - k, v, pos_k = self._complete_kv(k, v, execution_context) - pos_k = pos_k[:, None] - - if self.causal: - pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) - delta = pos_q - pos_k - attn_bias = (pos_k >= 0) & (delta >= 0) - if self.context is not None: - attn_bias = attn_bias & (delta < self.context) - attn_bias = attn_bias[:, None] + slot_attention = getattr(self, "_slot_attention", None) + slot_attention_rows = getattr(self, "_slot_attention_rows", None) + advance_kv_offset = False + slot_path = state is not None and execution_context is not None and self.causal + if slot_path and slot_attention_rows is not None: + # One launch: the gathered MHA offset equals the ring's end_offset + # for every slot, and both advance together below. + valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T + x = slot_attention_rows( + q, + k, + v, + state.kv_cache.cache, + offset, + execution_context.state_slot_ids, + valid_lengths, + self.context if self.context is not None else -1, + ) + advance_kv_offset = True + elif slot_attention is not None and slot_path: + # Current streaming batches use one exact T. Padding rows own + # scratch slots and must not advance persistent attention state. + valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T + x = slot_attention( + q, + k, + v, + state.kv_cache.cache, + state.kv_cache.end_offset, + execution_context.state_slot_ids, + valid_lengths, + self.context if self.context is not None else -1, + ) else: - attn_bias = None + k, v, pos_k = self._complete_kv(k, v, execution_context) + pos_k = pos_k[:, None] - 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) + if self.causal: + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) + delta = pos_q - pos_k + attn_bias = (pos_k >= 0) & (delta >= 0) + if self.context is not None: + attn_bias = attn_bias & (delta < self.context) + attn_bias = attn_bias[:, None] + else: + attn_bias = None + + 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) @@ -824,6 +868,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): else: next_offset = torch.where(execution_context.valid_rows, offset + T, offset) state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset) + if advance_kv_offset: + state.kv_cache.end_offset.index_copy_(0, execution_context.state_slot_ids, next_offset) return x @@ -1624,6 +1670,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs) self.decoder = nn.ModuleList() + decoder_input_frame_rate = current_frame_rate for decoder_kwargs_i in decoder_kwargs_list: decoder_kwargs_i = dict(decoder_kwargs_i) if decoder_kwargs_i["module_type"] == "PatchedPretransform": @@ -1638,6 +1685,11 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): context=int(round(current_frame_rate * context_duration)), ) ) + # Tokens this transformer sees per input code frame. Streaming + # ring headroom is sized from it (see initialize_decoder_state_pool). + self.decoder[-1].tokens_per_input_frame = max( + 1, int(round(current_frame_rate / decoder_input_frame_rate)) + ) current_frame_rate *= self.decoder[-1].downsample_ratio expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor) @@ -1731,10 +1783,29 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): self._decoder_state_capacity = 0 self._decoder_slot_offsets = None - def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None: - """Allocate persistent decoder state independently of execution B.""" - if state_capacity <= 0 or scratch_capacity < 0: - raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.") + def initialize_decoder_state_pool( + self, state_capacity: int, scratch_capacity: int = 0, chunk_frames: int = 0 + ) -> None: + """Allocate persistent decoder state independently of execution B. + + ``chunk_frames`` > 0 sizes every decoder attention ring as + ``context + chunk_frames * tokens_per_input_frame`` so that a chunk of + up to ``chunk_frames`` code frames never evicts keys inside its own + causal window; chunked streaming then matches whole-sequence decoding. + ``chunk_frames`` = 0 keeps the legacy ``context``-sized rings. + """ + if state_capacity <= 0 or scratch_capacity < 0 or chunk_frames < 0: + raise ValueError( + "Invalid decoder state capacities: " + f"state={state_capacity}, scratch={scratch_capacity}, chunk_frames={chunk_frames}." + ) + for module in self.decoder: + if not isinstance(module, MossAudioTokenizerProjectedTransformer): + continue + tokens_per_frame = int(getattr(module, "tokens_per_input_frame", 1)) + for attention in module.modules(): + if isinstance(attention, MossAudioTokenizerMultiheadAttention): + attention._ring_headroom = chunk_frames * tokens_per_frame self._start_streaming(state_capacity + scratch_capacity, decoder_only=True) self._decoder_state_capacity = state_capacity + scratch_capacity diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py index 4d277f44e..26e9f2fde 100644 --- a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py @@ -81,15 +81,31 @@ class CUDAGraphStreamingDecoderWrapper: frame_sizes: list[int], num_quantizers: int, vllm_config: VllmConfig, + scratch_base: int | None = None, + private_pool: bool = False, + compiled_decode: nn.Module | None = None, ) -> None: self.codec = codec self.state_capacity = int(state_capacity) + # Graph padding rows use scratch slots [scratch_base, scratch_base+B). + # A wrapper replaying concurrently with another needs its own scratch + # range and graph memory pool. + self.scratch_base = self.state_capacity if scratch_base is None else int(scratch_base) + self._private_pool = bool(private_pool) + # Graphs bake in the cuBLAS workspace of the stream they were captured + # on, and torch.cuda.graph shares one default capture stream. Split-K + # GEMMs synchronize through that workspace, so two graphs replaying + # concurrently from one capture stream can spin on each other forever. + self._capture_stream: torch.cuda.Stream | None = None self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity}) self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0}) self.num_quantizers = int(num_quantizers) self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {} self._pool = None self._warmed_up = False + if compiled_decode is not None: + self._compiled_decode: nn.Module | None = compiled_decode + return # vLLM owns Inductor compilation; this wrapper remains the sole owner # of CUDA Graph capture/replay because it understands persistent codec # state slots. Disable vLLM's CUDA Graph layer to avoid nested graphs. @@ -98,7 +114,7 @@ class CUDAGraphStreamingDecoderWrapper: compile_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE compile_config.compilation_config.static_forward_context = {} with set_current_vllm_config(compile_config): - self._compiled_decode: nn.Module | None = _MossStreamingDecodeCompileAdapter( + self._compiled_decode = _MossStreamingDecodeCompileAdapter( codec, vllm_config=compile_config, ) @@ -175,7 +191,7 @@ class CUDAGraphStreamingDecoderWrapper: frame_size, exc_info=True, ) - scratch_slots = self.state_capacity + torch.arange( + scratch_slots = self.scratch_base + torch.arange( batch_size, dtype=torch.long, device=device, @@ -206,10 +222,13 @@ class CUDAGraphStreamingDecoderWrapper: device=device, ) lengths = torch.zeros(batch_size, dtype=torch.long, device=device) - scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + scratch_slots = self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=device) valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) - stream = torch.cuda.Stream() + if self._private_pool and self._capture_stream is None: + self._capture_stream = torch.cuda.Stream(device=device) + # Warm up on the capture stream so its workspace exists before capture. + stream = self._capture_stream or torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(3): @@ -219,9 +238,11 @@ class CUDAGraphStreamingDecoderWrapper: self.codec.reset_decoder_state_slots(scratch_slots) if self._pool is None: - self._pool = current_platform.get_global_graph_pool() + self._pool = ( + torch.cuda.graph_pool_handle() if self._private_pool else current_platform.get_global_graph_pool() + ) graph = CUDAGraph() - with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"): + with torch.cuda.graph(graph, pool=self._pool, stream=self._capture_stream, capture_error_mode="thread_local"): audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows) self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph( @@ -272,7 +293,7 @@ class CUDAGraphStreamingDecoderWrapper: entry.static_lengths.zero_() entry.static_lengths[:actual_batch_size].fill_(int(frame_size)) entry.static_state_slot_ids.copy_( - self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) + self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) ) entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True) entry.static_valid_rows.zero_() diff --git a/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py new file mode 100644 index 000000000..34361e547 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Decode a stream's first codec chunk as soon as Stage 1 receives it. + +The Stage-1 engine loop only sees a new chunk between steps, and a step waits +for its codec graph to finish. A first chunk that arrives mid-step therefore +waits for the rest of that step, then for its own replay. This path takes the +first chunk from the receive thread, replays a small dedicated graph on its +own stream while the main step runs, and emits the audio straight to the +engine output queue. The main path later receives the same chunk, finds it +already decoded, and continues the stream from the state written here. +""" + +from __future__ import annotations + +import math +import queue +import threading +import time +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink, _PreparedDelivery + +logger = init_logger(__name__) + + +class _SlotHandoff: + """Orders the main path's first use of a slot after the fast decode.""" + + __slots__ = ("done", "event", "error") + + def __init__(self) -> None: + self.done = threading.Event() + self.event: torch.cuda.Event | None = None + self.error: Exception | None = None + + +@dataclass +class _Job: + request_id: str + request_key: str + delivery: _PreparedDelivery + slot: int + codes: torch.Tensor # (n_vq, T) int64 on host + handoff: _SlotHandoff + + +class MossFirstChunkFastPath: + def __init__( + self, + session: Any, + wrapper: Any, + *, + n_vq: int, + frames: int, + codebook_size: int, + samples_per_frame: int, + n_channels: int, + sample_rate: torch.Tensor, + device: torch.device, + gate_main: bool = False, + handoff_timeout_s: float = 30.0, + ) -> None: + if not math.isfinite(handoff_timeout_s) or handoff_timeout_s <= 0: + raise ValueError("handoff_timeout_s must be positive and finite") + self._handoff_timeout_s = float(handoff_timeout_s) + self._session = session + # When set, the main codec stream waits for an in-flight fast decode + # before its next replay, so the small first-chunk graph runs mostly + # uncontended instead of interleaving with a large regular replay. + self._gate_main = bool(gate_main) + self._inflight: torch.cuda.Event | None = None + self._wrapper = wrapper + self._n_vq = int(n_vq) + self.frames = int(frames) + self._codebook_size = int(codebook_size) + self._samples = int(frames) * int(samples_per_frame) + self._n_channels = int(n_channels) + self._sr = sample_rate + if device.index is None: + device = torch.device(device.type, torch.accelerator.current_device_index()) + self._device = device + self._max_batch = max(wrapper.batch_sizes) + self._stream = torch.cuda.Stream(device=device, priority=-1) + self._host_codes = torch.zeros((self._n_vq, self._max_batch, self.frames), dtype=torch.long).pin_memory() + self._host_slots = torch.zeros(self._max_batch, dtype=torch.long).pin_memory() + audio_shape = (self._max_batch, self._n_channels, self._samples) + self._host_audio = torch.empty(audio_shape, dtype=torch.float32).pin_memory() + self._copied = torch.cuda.Event() + self._sink: FirstAudioSink | None = None + # Request key -> slot for chunks decoded here and not yet seen by the main path. + self._decoded: dict[str, int] = {} + # A failed output-queue handoff must be replayed from the already + # decoded samples. Re-decoding would advance the streaming state twice. + self._unsent_audio: dict[str, torch.Tensor] = {} + self._handoffs: dict[int, _SlotHandoff] = {} + self._lock = threading.Lock() + self._jobs: queue.SimpleQueue[_Job | None] = queue.SimpleQueue() + self._thread: threading.Thread | None = None + self._closed = False + self._worker_error: Exception | None = None + self._emitted = 0 + self._batches = 0 + self._decode_s = 0.0 + + def bind(self, sink: FirstAudioSink) -> None: + if not callable(getattr(sink, "prepare", None)): + raise TypeError("The codec first-chunk sink requires prepare(request_ids)") + with self._lock: + if self._closed: + raise RuntimeError("The codec first-chunk decoder is closed") + self._sink = sink + if self._thread is None: + self._thread = threading.Thread(target=self._run, name="moss-codec-first-chunk", daemon=True) + self._thread.start() + + def close(self) -> None: + with self._lock: + if not self._closed: + self._closed = True + self._jobs.put(None) + thread = self._thread + if thread is not None: + thread.join(timeout=5.0) + # A stalled decode may still own GPU state: keep its thread and + # handoff, rather than permitting a restart or slot reuse. + if thread.is_alive(): + logger.warning("Codec first-chunk decoder is still stopping; slot waits remain bounded") + + # ------------------------------------------------------------------ + # Receive thread + # ------------------------------------------------------------------ + def submit( + self, + request_id: str, + request_key: str, + codes_flat: Any, + req_slots: dict[str, int], + ) -> bool: + """Claim a first chunk for fast decode; False leaves it to the main path.""" + # Own the codes: the job outlives this call and the payload it came from. + codes = torch.as_tensor(codes_flat, dtype=torch.long).reshape(-1).clone() + if codes.numel() != self._n_vq * self.frames: + return False + with self._lock: + if ( + self._closed + or self._sink is None + or self._thread is None + or not self._thread.is_alive() + or request_key in req_slots + ): + return False + delivery = self._sink.prepare([request_id]) + if request_id not in delivery.routes: + return False + slot = self._session.acquire(oldest=True) + if slot is None: + return False + handoff = _SlotHandoff() + self._handoffs[slot] = handoff + self._decoded[request_key] = slot + req_slots[request_key] = slot + codes = codes.reshape(self._n_vq, self.frames) + # Enqueue under the admission lock: close's sentinel cannot + # overtake an accepted job. + self._jobs.put(_Job(request_id, request_key, delivery, slot, codes, handoff)) + return True + + # ------------------------------------------------------------------ + # Main (engine) thread + # ------------------------------------------------------------------ + def get_request_slot(self, request_key: str, req_slots: dict[str, int]) -> int | None: + """Observe a completed admission before deciding whether cleanup is needed. + + The scheduler removes cancelled routes before the runner's finish hook. + A submit that already froze its route must finish leasing its slot + before that hook checks it. Waiting for decode happens outside this + lock, so the worker can still complete the slot handoff. + """ + with self._lock: + return req_slots.get(request_key) + + def take_decoded(self, request_key: str) -> bool: + """True once per request whose first chunk this path already decoded.""" + with self._lock: + return self._decoded.pop(request_key, None) is not None + + def take_unsent_audio(self, request_key: str) -> torch.Tensor | None: + """Return first-chunk samples when the out-of-band sink rejected them.""" + with self._lock: + return self._unsent_audio.pop(request_key, None) + + def forget(self, request_key: str) -> None: + with self._lock: + self._decoded.pop(request_key, None) + self._unsent_audio.pop(request_key, None) + + def gate(self) -> None: + """Queue the current stream behind the latest fast decode, if still running.""" + event = self._inflight + if self._gate_main and event is not None and not event.query(): + torch.cuda.current_stream(self._device).wait_event(event) + + def order_after(self, slot: int) -> None: + """Make the current stream's next use of ``slot`` follow the fast decode.""" + with self._lock: + handoff = self._handoffs.get(slot) + if handoff is None: + return + if not handoff.done.wait(timeout=self._handoff_timeout_s): + # Keep the handoff. The worker may still be writing the slot, so + # neither cancellation nor a retry may release/re-decode it. + raise TimeoutError(f"Codec first-chunk slot {slot} did not complete within {self._handoff_timeout_s:g}s") + if handoff.error is not None: + raise RuntimeError(f"Codec first-chunk decode failed for slot {slot}") from handoff.error + if handoff.event is not None: + torch.cuda.current_stream(self._device).wait_event(handoff.event) + with self._lock: + self._handoffs.pop(slot, None) + + # ------------------------------------------------------------------ + # Decode thread + # ------------------------------------------------------------------ + def _run(self) -> None: + try: + current_omni_platform.set_device(self._device) + self._run_jobs() + except Exception as error: + logger.exception("MOSS codec first-chunk worker failed") + self._worker_error = error + finally: + with self._lock: + self._closed = True + for handoff in self._handoffs.values(): + if not handoff.done.is_set(): + handoff.error = self._worker_error or RuntimeError("Codec first-chunk worker stopped") + handoff.done.set() + # Accepted, unprocessed routes must not be left waiting for PCM. + while True: + try: + job = self._jobs.get_nowait() + except queue.Empty: + break + if job is not None: + self._fail_delivery(job) + + @staticmethod + def _fail_delivery(job: _Job) -> None: + try: + job.delivery.fail([job.request_id]) + except Exception: + logger.exception("Codec first-chunk error delivery failed for %s", job.request_id) + + def _run_jobs(self) -> None: + while True: + job = self._jobs.get() + if job is None: + return + jobs = [job] + while len(jobs) < self._max_batch: + try: + extra = self._jobs.get_nowait() + except queue.Empty: + break + if extra is None: + self._jobs.put(None) + break + jobs.append(extra) + try: + start = time.perf_counter() + self._decode(jobs) + self._decode_s += time.perf_counter() - start + self._batches += 1 + before, self._emitted = self._emitted, self._emitted + len(jobs) + if self._emitted // 1000 != before // 1000: + logger.info( + "MOSS codec first-chunk fast path: %d chunks in %d batches, %.2f ms/batch", + self._emitted, + self._batches, + 1e3 * self._decode_s / self._batches, + ) + except Exception as error: + logger.exception("MOSS codec first-chunk fast decode failed for %s", [j.request_id for j in jobs]) + # A partial graph replay may have advanced the slot. Surface + # the failure instead of re-decoding and advancing it twice. + for item in jobs: + item.handoff.error = error + item.handoff.done.set() + self._fail_delivery(item) + raise + + @torch.no_grad() + def _decode(self, jobs: list[_Job]) -> None: + n = len(jobs) + for row, item in enumerate(jobs): + self._host_codes[:, row].copy_(item.codes) + self._host_slots[row] = item.slot + stream = self._stream + with torch.cuda.stream(stream): + for item in jobs: + self._session.order_after_reset(item.slot, stream) + codes = self._host_codes[:, :n].to(self._device, non_blocking=True).clamp_(0, self._codebook_size - 1) + slots = self._host_slots[:n].to(self._device, non_blocking=True) + result = self._wrapper.decode(codes, slots) + if result is None: + raise RuntimeError(f"no first-chunk codec graph for B={n}, T={self.frames}") + audio = result[0][:n, ..., : self._samples].float() + state_written = torch.cuda.Event() + state_written.record(stream) + self._inflight = state_written + self._host_audio[:n].view_as(audio).copy_(audio, non_blocking=True) + self._copied.record(stream) + for item in jobs: + item.handoff.event = state_written + self._copied.synchronize() + for row, item in enumerate(jobs): + wav = self._host_audio[row].clone() + if self._n_channels == 1: + wav = wav.reshape(-1) + try: + item.delivery([item.request_id], [wav], self._sr) + except Exception: + logger.exception("MOSS codec first-chunk output failed for %s", item.request_id) + with self._lock: + self._unsent_audio[item.request_key] = wav + finally: + item.handoff.done.set() diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py new file mode 100644 index 000000000..c382d4439 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/local_model_state.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MOSS Local request-slot state, independent of each step's batch order. + +State writes, gathers and MTP replay run on the runner stream. Output gathers +own their storage; the runner retains its normal asynchronous output snapshot. +No request dictionary holds a view into a mutable slot or CUDA graph buffer. +""" + +import numpy as np +import torch +from vllm.config import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.sampler import Sampler + +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +logger = init_logger(__name__) + + +def _metadata_to_device(data: np.ndarray, device: torch.device) -> torch.Tensor: + # CPU state tests need no CUDA allocator. The serving CUDA path retains + # vLLM's asynchronous pinned upload and its ownership/lifetime handling. + if device.type == "cpu": + return torch.from_numpy(data) + return async_tensor_h2d(data, device=device) + + +class _CodeRowsSnapshot(PackedOutputSnapshot): + """Owned per-request code rows backed by one gathered device tensor. + + The generic snapshot flattens every per-request leaf, concatenates them + into a slab and slices each back, then slices each again after the D2H. + Here the gathered ``[decode_rows, n_vq]`` tensor is already owned, so the + copy is a single D2H and the per-request views come from one ``split``. + """ + + def __init__(self, codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> None: + dict.__init__(self, {"codes": {"audio": self._rows(codes, decode_rows, num_reqs)}}) + self._codes = codes + self._decode_rows = decode_rows + self._num_reqs = num_reqs + self.producer_event = None + + @staticmethod + def _rows(codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> list[torch.Tensor]: + outputs = [codes.new_empty((0, codes.shape[-1]))] * num_reqs + for batch_row, row in zip(decode_rows, codes.split(1), strict=True): + outputs[batch_row] = row + return outputs + + def copy_to_cpu(self, copy_tensor): + cpu = copy_tensor(self._codes) + return {"codes": {"audio": self._rows(cpu, self._decode_rows, self._num_reqs)}} + + +class MossLocalModelState(OmniModelState): + def __init__(self, vllm_config, model, encoder_cache, device): + super().__init__(vllm_config, model, encoder_cache, device) + self._init_slot_buffers(self.scheduler_config.max_num_seqs, model.hidden_size, device, self.dtype) + self._batch_prefill = bool(getattr(model.config, "mrv2_batch_prefill", False)) + self._direct_tokens = bool(getattr(model.config, "mrv2_direct_tokens", False)) + self._local_eager_mtp = bool(getattr(model.config, "mrv2_eager_mtp", False)) + logger.info( + "MOSS Local MRV2 batch prefill=%s, determined tokens=%s, eager MTP=%s", + self._batch_prefill, + self._direct_tokens, + self._local_eager_mtp, + ) + logger.info("MOSS Local MRV2 GPU slot state enabled: capacity=%d", self.scheduler_config.max_num_seqs) + + def _init_slot_buffers(self, capacity, hidden_size, device, dtype): + self._hidden_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device) + self._codes_pool = torch.full( + (capacity, self.model.n_vq), self.model.audio_pad_token_id, dtype=torch.long, device=device + ) + # Execution eligibility mirrors audio_state.is_stopping. Binary MTP + # continuation is a per-step output, not a sticky request-state update. + self._active_pool = torch.zeros(capacity, dtype=torch.bool, device=device) + self._active_host = [False] * capacity + self._decode_rows = [] + self._prefill_staging = [] + self._batch_prefill = False + self._direct_tokens = False + # Eager MTP: a frame is sampled right after the forward that produced + # its hidden state, one step earlier than the canonical preprocess + # path. Its input-embedding contribution (audio embedding times emit + # mask) is added to the next step's text embedding; ``_keep_pool`` + # chains emission so a stopped stream cannot emit again. + self._local_eager_mtp = False + self._eager_emb_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device) + self._keep_pool = torch.zeros(capacity, dtype=torch.bool, device=device) + self._completing_rows = [] + self._local_eager_rows = [] + # A slot takes the eager path only after its final prefill chunk ran + # an eager MTP in this state; others keep the canonical path, so a + # reused slot never reads a previous request's frame state. + self._eager_host = [False] * capacity + self._mtp_dispatcher = None + + def sample_determined_tokens(self, batch, sampler): + """Skip full-vocabulary sampling only for unconstrained Local tokens. + + The audio/binary MTP sampler has already made the random decision. The + outer text distribution has exactly one finite logit. Keep the original + path for processors that can mask it or request distribution metadata. + """ + if not self._direct_tokens or type(sampler) is not Sampler: + return None + rows = batch.idx_mapping_np + if ( + sampler.compute_nans + or sampler.return_sampling_mask + or sampler.trace_replay_state is not None + or sampler.get_logprobs_dims(rows) is not None + or np.any(sampler.logit_bias_state.use_logit_bias[rows]) + or np.any(sampler.penalties_state.use_penalty[rows]) + or np.any(sampler.bad_words_state.num_bad_words.np[rows]) + or ( + sampler.thinking_budget_state.enabled + and np.any(sampler.thinking_budget_state.use_thinking_budget[rows]) + ) + ): + return None + keep = self.model._batch_should_continue + if keep is None or keep.numel() != batch.num_reqs: + return None + # Owned storage: subsequent forwards can replace the continue mask + # while these token IDs are being copied or delivered asynchronously. + tokens = torch.where(keep, self.model.audio_assistant_slot_token_id, self.model.im_end_token_id) + num_sampled, num_rejected = get_num_sampled_and_rejected( + batch.seq_lens.new_ones(batch.num_reqs), + batch.seq_lens, + batch.cu_num_logits, + batch.idx_mapping, + sampler.req_states.prefill_len.gpu, + ) + return SamplerOutput(tokens.reshape(-1, 1), None, None, num_sampled, num_rejected) + + def _apply_reference_batch(self, embeds, references): + """One owned pinned upload for reference codes and their token positions. + + Host storage may be reused only after its H2D event completes. Device + storage is consumed and overwritten in runner-stream order. If all + cached slots are busy, allocate a transient snapshot instead of waiting. + """ + if not references: + return + count = sum(codes.shape[0] for _, codes in references) + capacity = max(count, self._static_inputs_embeds.shape[0]) + staging = None + for item in self._prefill_staging: + if item[0].shape[0] >= count and (item[2] is None or item[2].query()): + staging = item + break + if staging is None: + host = torch.empty((capacity, self.model.n_vq + 1), dtype=torch.long, pin_memory=embeds.is_cuda) + device = torch.empty_like(host, device=embeds.device) if embeds.is_cuda else host + staging = [host, device, torch.cuda.Event() if embeds.is_cuda else None] + if len(self._prefill_staging) < 3: + self._prefill_staging.append(staging) + host, device, event = staging + offset = 0 + for start, codes in references: + end = offset + codes.shape[0] + host[offset:end, 0].numpy()[:] = np.arange(start, start + codes.shape[0]) + host[offset:end, 1:].copy_(codes) + offset = end + if embeds.is_cuda: + device[:count].copy_(host[:count], non_blocking=True) + event.record(torch.cuda.current_stream(embeds.device)) + packed = device[:count] + positions = packed[:, 0] + audio = self.model._audio_embed(packed[:, 1:]) + embeds.index_copy_(0, positions, embeds.index_select(0, positions) + audio) + + def add_request(self, req_index, new_req_data): + super().add_request(req_index, new_req_data) + self._hidden_pool[req_index].zero_() + self._codes_pool[req_index].fill_(self.model.audio_pad_token_id) + self._active_pool[req_index].fill_(True) + self._active_host[req_index] = True + self._eager_host[req_index] = False + self._keep_pool[req_index].fill_(False) + self._eager_emb_pool[req_index].zero_() + + def remove_request(self, req_index_or_id): + # Freeing is CPU bookkeeping. Reinitialization happens on admission, + # in stream order, so pending output snapshots remain independent. + super().remove_request(req_index_or_id) + + def _select_rows(self, tensor, rows): + if rows == list(range(len(rows))): + return tensor[: len(rows)] + indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device) + return tensor.index_select(0, indices) + + def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None): + input_ids = model_inputs.get("input_ids") + if input_ids is None: + input_ids = input_batch.input_ids + embeds = model_inputs.get("inputs_embeds") + if embeds is None: + embeds = self.model.embed_input_ids(input_ids[: input_batch.num_tokens]) + model_inputs["inputs_embeds"] = embeds + elif self._static_inputs_embeds is not None and embeds.data_ptr() == self._static_inputs_embeds.data_ptr(): + embeds[: input_batch.num_tokens].copy_(self.model.embed_input_ids(input_ids[: input_batch.num_tokens])) + + decode_rows = [] + completing_rows = [] + references = [] + for row in range(input_batch.num_reqs): + slot = int(input_batch.idx_mapping_np[row]) + buf = self.intermediate_buffer.buffers[slot] + if not buf or str(buf.get("req_id", "")).startswith("_warmup_"): + continue + start = int(input_batch.query_start_loc_np[row]) + count = int(input_batch.num_scheduled_tokens[row]) + prompt_len = ( + None if req_states is None else self._get_req_state_value(getattr(req_states, "prompt_len", None), slot) + ) + computed = ( + None + if req_states is None + else self._get_req_state_value(getattr(req_states, "num_computed_tokens", None), slot) + ) + if computed is None: + computed = self._get_input_batch_num_computed(input_batch, slot, row) + prefill = computed < prompt_len if computed is not None and prompt_len is not None else count > 1 + if prefill and computed is not None and prompt_len is not None and computed + count >= prompt_len: + completing_rows.append(row) + if count == 1 and not prefill and isinstance(buf.get("audio_state"), dict): + active = not bool(buf["audio_state"].get("is_stopping")) + if active != self._active_host[slot]: + self._active_pool[slot].fill_(active) + self._active_host[slot] = active + decode_rows.append(row) + continue + if self._batch_prefill: + ref = (buf.get("codes", {}) or {}).get("ref") + ref_offset = int(buf.get("ref_offset", 0)) + # Device-side references keep the canonical path; do not add + # an implicit D2H just to place them in a host staging buffer. + if not isinstance(ref, torch.Tensor) or ref.device.type == "cpu": + if isinstance(ref, torch.Tensor) and ref.numel(): + if ref.dim() == 1 and ref.numel() % self.model.n_vq == 0: + ref = ref.view(-1, self.model.n_vq) + if ref.dim() == 2: + chunk = ref[ref_offset : ref_offset + count] + if chunk.numel() and chunk.shape[0] == count: + references.append((start, chunk)) + # Text embeddings have already been computed for all input + # tokens. Match the canonical prefill state transition. + self.intermediate_buffer.update( + slot, + { + "audio_state": {"is_stopping": False}, + "ref_offset": ref_offset + count, + }, + self.model.gpu_resident_buffer_keys, + ) + if not self._active_host[slot]: + self._active_pool[slot].fill_(True) + self._active_host[slot] = True + continue + # Reference conditioning and chunked-prefill coordinates retain + # the model's canonical scalar implementation. + info = dict(buf) + info["_omni_is_prefill"] = prefill + if prompt_len is not None: + info["_omni_prompt_len"] = prompt_len + if computed is not None: + info["_omni_num_computed_tokens"] = computed + ids = input_ids[start : start + count] + previous = embeds[start : start + count] + new_ids, new_embeds, updates = self.model.preprocess(ids, previous, **info) + if "mtp_inputs" in updates: + raise RuntimeError("MOSS slot state requires an admitted prefill before decode") + if self._preprocess_result_needs_writeback(ids, new_ids): + ids.copy_(new_ids) + if self._preprocess_result_needs_writeback(previous, new_embeds): + previous.copy_(new_embeds) + self.intermediate_buffer.update(slot, updates, self.model.gpu_resident_buffer_keys) + active = not bool(updates.get("audio_state", {}).get("is_stopping", False)) + if active != self._active_host[slot]: + self._active_pool[slot].fill_(active) + self._active_host[slot] = active + + self._apply_reference_batch(embeds, references) + self._decode_rows = decode_rows + if self._local_eager_mtp: + self._completing_rows = completing_rows + self._mtp_dispatcher = mtp_batch_descriptor_dispatcher + eager_rows = [row for row in decode_rows if self._eager_host[int(input_batch.idx_mapping_np[row])]] + canonical_rows = [row for row in decode_rows if not self._eager_host[int(input_batch.idx_mapping_np[row])]] + self._local_eager_rows = eager_rows + if eager_rows: + # Their frame was sampled after the previous forward. + slots = self._select_rows(input_batch.idx_mapping, eager_rows) + offsets = self._mtp_offsets[: len(eager_rows)] + offsets.copy_(self._select_rows(input_batch.query_start_loc, eager_rows)) + embeds.index_copy_( + 0, offsets, embeds.index_select(0, offsets) + self._eager_emb_pool.index_select(0, slots) + ) + if canonical_rows: + self._run_slot_mtp(input_batch, input_ids, embeds, canonical_rows, mtp_batch_descriptor_dispatcher) + elif decode_rows: + self._run_slot_mtp(input_batch, input_ids, embeds, decode_rows, mtp_batch_descriptor_dispatcher) + + def _run_slot_mtp(self, batch, input_ids, embeds, rows, dispatcher): + bsz = len(rows) + slots = self._select_rows(batch.idx_mapping, rows) + # Upstream query_start_loc is int32; index_copy_ requires int64. + offsets = self._mtp_offsets[:bsz] + offsets.copy_(self._select_rows(batch.query_start_loc, rows)) + ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz] + hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz] + ids.copy_(input_ids.index_select(0, offsets)) # MRV2 token inputs are int32. + torch.index_select(embeds, 0, offsets, out=emb) + torch.index_select(self._hidden_pool, 0, slots, out=hidden) + step[:, 0].copy_(self._active_pool.index_select(0, slots)) + + buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows] + req_ids = [str(buf["req_id"]) for buf in buffers] + generators = [ + self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device) + for req_id, buf in zip(req_ids, buffers, strict=True) + ] + can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators) + desc = dispatcher(bsz) if can_graph and dispatcher is not None else None + mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE + num_tokens = int(desc.num_tokens) if desc is not None else bsz + with set_forward_context( + None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc + ): + new_emb, codes = self._call_mtp_with_sampling( + self._mtp_input_ids[:num_tokens], + self._mtp_input_embeds[:num_tokens], + self._mtp_hidden[:num_tokens], + self._mtp_text_step[:num_tokens], + buffers=buffers, + req_ids=req_ids, + generators=generators, + ) + embeds.index_copy_(0, offsets, new_emb[:bsz]) + self._codes_pool.index_copy_(0, slots, codes[:bsz]) + + def _run_local_eager_mtp(self, batch, rows, completing): + """Sample the next frame from this forward's hidden state. + + Same MTP inputs as the canonical next-step preprocess (hidden state + and active flag); the input embedding is zero, so the returned + embedding is exactly the audio contribution the next step adds to its + text embedding. A decode row emits only if its previous frame did. + """ + bsz = len(rows) + slots = self._select_rows(batch.idx_mapping, rows) + ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz] + hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz] + ids.zero_() # the MTP does not read token ids + emb.zero_() + torch.index_select(self._hidden_pool, 0, slots, out=hidden) + gate = self._keep_pool.index_select(0, slots) + if completing: + first = _metadata_to_device( + np.asarray([row in completing for row in rows], dtype=np.bool_), + gate.device, + ) + gate = gate | first + step[:, 0].copy_(self._active_pool.index_select(0, slots) & gate) + + buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows] + req_ids = [str(buf["req_id"]) for buf in buffers] + generators = [ + self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device) + for req_id, buf in zip(req_ids, buffers, strict=True) + ] + dispatcher = self._mtp_dispatcher + can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators) + desc = dispatcher(bsz) if can_graph and dispatcher is not None else None + mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE + num_tokens = int(desc.num_tokens) if desc is not None else bsz + with set_forward_context( + None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc + ): + contribution, codes = self._call_mtp_with_sampling( + self._mtp_input_ids[:num_tokens], + self._mtp_input_embeds[:num_tokens], + self._mtp_hidden[:num_tokens], + self._mtp_text_step[:num_tokens], + buffers=buffers, + req_ids=req_ids, + generators=generators, + ) + codes = codes[:bsz] + self._eager_emb_pool.index_copy_(0, slots, contribution[:bsz]) + self._codes_pool.index_copy_(0, slots, codes) + self._keep_pool.index_copy_(0, slots, codes.ne(self.model.audio_pad_token_id).any(dim=-1)) + + def run_postprocess(self, hidden_states, input_batch): + # query_start_loc is already device-resident and authoritative for + # mixed/chunked-prefill batches. index_copy snapshots graph outputs. + last = input_batch.query_start_loc[1 : input_batch.num_reqs + 1] - 1 + self._hidden_pool.index_copy_( + 0, input_batch.idx_mapping[: input_batch.num_reqs], hidden_states.index_select(0, last) + ) + if self._local_eager_mtp: + completing = self._completing_rows + rows = sorted(self._local_eager_rows + completing) + # Canonical decode rows emit their preprocess frame; eager rows + # and completed prefills emit the frame sampled below. + self._decode_rows = sorted(set(self._decode_rows) | set(completing)) + if rows: + self._run_local_eager_mtp(input_batch, rows, set(completing)) + for row in completing: + self._eager_host[int(input_batch.idx_mapping_np[row])] = True + + def postprocess_model_output(self, model_output, input_batch, req_states): + # Local backbone returns a tensor; keep the standard runner's output + # snapshot/materializer contract and omit hidden payload transport. + if not isinstance(model_output, torch.Tensor): + raise TypeError("MOSS Local slot state requires a tensor backbone output") + model = self.model + model._batch_state = None + model._batch_state_spans = None + model._batch_should_continue = torch.ones(input_batch.num_reqs, device=model_output.device, dtype=torch.bool) + if not self._decode_rows: + return model_output, {} + # This gather is an owned snapshot. Later slot writes cannot change + # data already handed to the runner, even before its D2H completes. + slots = self._select_rows(input_batch.idx_mapping, self._decode_rows) + codes = self._codes_pool.index_select(0, slots) + keep = codes.ne(model.audio_pad_token_id).any(dim=-1) + if self._decode_rows == list(range(input_batch.num_reqs)): + model._batch_should_continue = keep + else: + rows = _metadata_to_device( + np.asarray(self._decode_rows, dtype=np.int64), + codes.device, + ) + model._batch_should_continue.index_copy_(0, rows, keep) + return model_output, _CodeRowsSnapshot(codes, list(self._decode_rows), input_batch.num_reqs) diff --git a/vllm_omni/model_executor/models/moss_tts/local_short_attention.py b/vllm_omni/model_executor/models/moss_tts/local_short_attention.py new file mode 100644 index 000000000..522d74c76 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/local_short_attention.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Single-query attention for the frame-local MOSS depth transformer. + +The caller has already written this token's K/V and provides only its valid +prefix. There is no persistent state here. BF16 probability rounding matches +the FlashAttention P/V multiply, but reduction order is not bit-identical. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _short_attention( + q_ptr, + k_ptr, + v_ptr, + out_ptr, + qs0: tl.constexpr, + qs1: tl.constexpr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + heads: tl.constexpr, + dim: tl.constexpr, + tokens: tl.constexpr, + block_d: tl.constexpr, + block_t: tl.constexpr, +): + row = tl.program_id(0) + batch, head = row // heads, row % heads + d, t = tl.arange(0, block_d), tl.arange(0, block_t) + query = tl.load(q_ptr + batch * qs0 + head * qs1 + d, d < dim, 0).to(tl.float32) + keys = tl.load( + k_ptr + batch * ks0 + head * ks1 + t[:, None] * ks2 + d[None, :], + (t[:, None] < tokens) & (d[None, :] < dim), + 0, + ).to(tl.float32) + logits = tl.sum(keys * query[None, :], axis=1) * (dim**-0.5 * 1.4426950408889634) + logits = tl.where(t < tokens, logits, float("-inf")) + probs = tl.exp2(logits - tl.max(logits, axis=0)) + denominator = tl.sum(probs, axis=0) + values = tl.load( + v_ptr + batch * vs0 + head * vs1 + t[:, None] * vs2 + d[None, :], + (t[:, None] < tokens) & (d[None, :] < dim), + 0, + ) + probs = probs.to(values.dtype).to(tl.float32) + output = tl.sum(values.to(tl.float32) * probs[:, None], axis=0) / denominator + tl.store(out_ptr + row * dim + d, output, d < dim) + + +@torch.library.custom_op("vllm_omni::moss_local_short_attention", mutates_args=()) +def local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + batch, heads, seq_len, dim = query.shape + if not ( + query.is_cuda + and key.device == query.device + and value.device == query.device + and query.dtype == key.dtype == value.dtype == torch.bfloat16 + and seq_len == 1 + and 0 < dim <= 128 + and key.shape == value.shape + and key.shape[:2] == query.shape[:2] + and key.shape[-1] == dim + and 0 < key.shape[2] <= 16 + and query.stride(-1) == key.stride(-1) == value.stride(-1) == 1 + ): + raise ValueError("MOSS short attention requires CUDA BF16 single-query inputs and 1..16 valid keys") + output = query.new_empty((batch, heads, 1, dim)) + _short_attention[(batch * heads,)]( + query, + key, + value, + output, + query.stride(0), + query.stride(1), + key.stride(0), + key.stride(1), + key.stride(2), + value.stride(0), + value.stride(1), + value.stride(2), + heads, + dim, + key.shape[2], + triton.next_power_of_2(dim), + triton.next_power_of_2(key.shape[2]), + num_warps=1, + ) + return output + + +@local_short_attention.register_fake +def _fake_local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + return query.new_empty(query.shape) 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 d87f01b4b..a526d8e33 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 @@ -8,7 +8,10 @@ from __future__ import annotations +import threading from collections.abc import Iterable +from contextlib import nullcontext +from functools import partial from typing import Any import torch @@ -32,11 +35,16 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import ( CUDAGraphStreamingDecoderWrapper, ) +from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import ( + MossFirstChunkFastPath, +) from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, pack_output_snapshot from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp +from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink logger = init_logger(__name__) @@ -53,18 +61,46 @@ class _MossCodecStreamSession: vllm_config: VllmConfig, graph_batch_sizes: list[int] | None = None, graph_frame_sizes: list[int] | None = None, + gpu_output: bool = False, + chunk_frames: int = 0, + ring_headroom: bool = False, + fast_batch_sizes: list[int] | None = None, + fast_frames: int = 0, ) -> None: self._codec = codec self._state_capacity = int(state_capacity) + self._chunk_frames = int(chunk_frames) + self._ring_headroom = bool(ring_headroom) self._n_vq = int(n_vq) self._device = next(codec.parameters()).device + self._gpu_output = gpu_output and self._device.type == "cuda" self._free_stream_slots = list(reversed(range(self._state_capacity))) self._leased_slots: set[int] = set() + # The first-chunk fast path leases slots from its decode thread. + self._slot_lock = threading.Lock() + self._slot_reset_events: dict[int, torch.cuda.Event] = {} self._closed = False + self._metadata_index = 0 + self._metadata_ring = [] + if self._gpu_output: + # The metadata DMA is tiny, but a short ring forces the producer + # CPU thread to synchronize with the codec stream while it waits + # for a host slab to become reusable. Keep enough entries for a + # scheduler burst so staging remains enqueue-only in practice. + for _ in range(64): + host = torch.empty(2 * self._state_capacity, dtype=torch.long, pin_memory=True) + gpu = torch.empty_like(host, device=self._device) + self._metadata_ring.append((host, host.numpy(), gpu, torch.cuda.Event(), False)) self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None + self._fast_wrapper: CUDAGraphStreamingDecoderWrapper | None = None batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity}) frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0}) scratch_capacity = max(batch_sizes, default=0) if self._device.type in ("cuda", "npu") else 0 + fast_batch_sizes = sorted({int(size) for size in (fast_batch_sizes or []) if 0 < int(size)}) + if not (fast_frames > 0 and fast_batch_sizes and batch_sizes and frame_sizes and self._device.type == "cuda"): + fast_batch_sizes = [] + fast_scratch_base = self._state_capacity + scratch_capacity + scratch_capacity += max(fast_batch_sizes, default=0) self._total_state_capacity = self._state_capacity + scratch_capacity self._state_slot_ids = torch.arange( self._total_state_capacity, @@ -76,7 +112,14 @@ class _MossCodecStreamSession: if not callable(initialize_state_pool): raise RuntimeError("The streaming codec does not implement a decoder state pool.") with torch.no_grad(): - initialize_state_pool(self._state_capacity, scratch_capacity) + if self._ring_headroom: + if self._chunk_frames <= 0: + raise ValueError("codec_ring_headroom requires a positive codec chunk frame count.") + # Size every attention ring for context + one full chunk so a + # chunk write never evicts keys inside its own causal window. + initialize_state_pool(self._state_capacity, scratch_capacity, chunk_frames=self._chunk_frames) + else: + initialize_state_pool(self._state_capacity, scratch_capacity) if batch_sizes and frame_sizes and self._device.type in ("cuda", "npu"): if self._device.type == "npu": from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import ( @@ -101,16 +144,34 @@ class _MossCodecStreamSession: vllm_config=vllm_config, ) self._cudagraph_wrapper.warmup(self._device) + if fast_batch_sizes and self._cudagraph_wrapper.is_ready: + main_compiled = self._cudagraph_wrapper._compiled_decode + self._fast_wrapper = CUDAGraphStreamingDecoderWrapper( + codec, + state_capacity=self._state_capacity, + batch_sizes=fast_batch_sizes, + frame_sizes=[int(fast_frames)], + num_quantizers=self._n_vq, + vllm_config=vllm_config, + scratch_base=fast_scratch_base, + private_pool=True, + compiled_decode=main_compiled if main_compiled is not None else codec.decode_streaming_tensors, + ) + self._fast_wrapper.warmup(self._device) + if not self._fast_wrapper.is_ready: + self._fast_wrapper = None self.reset_slots(list(range(self._state_capacity + scratch_capacity))) if not self._cudagraph_wrapper.is_ready: self._cudagraph_wrapper = None - def acquire(self) -> int | None: - if not self._free_stream_slots: - return None - slot = self._free_stream_slots.pop() - self._leased_slots.add(slot) - return slot + def acquire(self, *, oldest: bool = False) -> int | None: + """Lease a slot; ``oldest`` takes the one whose reset was queued first.""" + with self._slot_lock: + if not self._free_stream_slots: + return None + slot = self._free_stream_slots.pop(0 if oldest else -1) + self._leased_slots.add(slot) + return slot def release(self, slot: int, *, state_already_reset: bool = False) -> None: if self._closed: @@ -119,8 +180,23 @@ class _MossCodecStreamSession: return if not state_already_reset: self.reset_slots([slot]) - self._leased_slots.remove(slot) - self._free_stream_slots.append(slot) + reset_event = None + if self._fast_wrapper is not None: + # The fast path decodes on another stream; it must not read this + # slot before the reset queued on the current stream has run. + reset_event = torch.cuda.Event() + reset_event.record(torch.cuda.current_stream(self._device)) + with self._slot_lock: + if reset_event is not None: + self._slot_reset_events[slot] = reset_event + self._leased_slots.remove(slot) + self._free_stream_slots.append(slot) + + def order_after_reset(self, slot: int, stream: torch.cuda.Stream) -> None: + with self._slot_lock: + event = self._slot_reset_events.pop(slot, None) + if event is not None and not event.query(): + stream.wait_event(event) def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None: reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) @@ -139,9 +215,29 @@ class _MossCodecStreamSession: raise ValueError(f"Decoder state slots must be contiguous, got {slots}") self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)]) + def _stage_slot_ids(self, slots: list[int], terminal_slots: set[int]) -> tuple[torch.Tensor, torch.Tensor]: + # Host slabs cannot be overwritten until their DMA completes. Device + # slabs are reused in producer-stream order, after all previous reads. + index = self._metadata_index + host, view, gpu, event, recorded = self._metadata_ring[index] + if recorded and not event.query(): + event.synchronize() + terminal_ids = [slot for slot in slots if slot in terminal_slots] + n, end = len(slots), len(slots) + len(terminal_ids) + view[:n] = slots + view[n:end] = terminal_ids + gpu[:end].copy_(host[:end], non_blocking=True) + event.record(torch.cuda.current_stream(self._device)) + self._metadata_ring[index] = (host, view, gpu, event, True) + self._metadata_index = (index + 1) % len(self._metadata_ring) + return gpu[:n], gpu[n:end] + def close(self) -> None: if self._closed: return + for _, _, _, event, recorded in self._metadata_ring: + if recorded: + event.synchronize() close_state_pool = getattr(self._codec, "close_decoder_state_pool", None) with torch.no_grad(): if callable(close_state_pool): @@ -154,6 +250,7 @@ class _MossCodecStreamSession: slot_codes: dict[int, torch.Tensor], *, terminal_slots: set[int] | None = None, + output_buffers: dict[int, torch.Tensor] | None = None, ) -> dict[int, torch.Tensor]: if not slot_codes: return {} @@ -178,17 +275,22 @@ class _MossCodecStreamSession: ], dim=1, ) - state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) + staged_terminal_ids = None + if self._metadata_ring: + state_slot_ids, staged_terminal_ids = self._stage_slot_ids(slots, terminal_slots) + else: + state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None if self._cudagraph_wrapper is not None: - # A terminal tail can safely pad T to a larger graph: decoder - # attention is causal, audio_lengths trims the padding, and every - # affected state slot is reset immediately after this step. + # An all-terminal group may borrow the next larger frame graph. + # The decoder is causal, so cropping the padded tail on the time + # axis recovers the exact-T waveform up to attention-ring + # eviction (see codec_ring_headroom); the state is reset below. graph_output = self._cudagraph_wrapper.decode( codes_step, state_slot_ids, - allow_frame_padding=len(terminal_slots) == len(slots), + allow_frame_padding=bool(terminal_slots) and len(terminal_slots) == len(slots), ) used_cudagraph = graph_output is not None @@ -214,11 +316,69 @@ class _MossCodecStreamSession: audio_tensor = result.audio if terminal_slots: - terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots] - terminal_slot_ids = state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows] + if staged_terminal_ids is not None: + terminal_slot_ids = staged_terminal_ids + else: + terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots] + terminal_slot_ids = ( + state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows] + ) self._reset_slot_ids(terminal_slot_ids) - audio = audio_tensor.detach().to("cpu", torch.float32) + if self._gpu_output and output_buffers is not None: + # Write directly into the caller-owned per-forward slab. This + # preserves graph output before the next replay without creating + # a temporary waveform tensor that would be packed again later. + out: dict[int, torch.Tensor] = {} + row = 0 + while row < len(slots): + destination = output_buffers.get(slots[row]) + if destination is None: + row += 1 + continue + end = row + 1 + if end < len(slots) and destination.is_contiguous(): + storage = destination.untyped_storage().data_ptr() + offset = destination.storage_offset() + row_elements = destination.numel() + while end < len(slots): + following = output_buffers.get(slots[end]) + if not ( + following is not None + and following.shape == destination.shape + and following.dtype == destination.dtype + and following.device == destination.device + and following.is_contiguous() + and following.untyped_storage().data_ptr() == storage + and following.storage_offset() == offset + (end - row) * row_elements + ): + break + end += 1 + if end > row + 1: + # A mixed-frame forward can leave gaps between equal-sized + # rows in its slab. Copy each adjacent run as one batch, + # without overwriting the intervening requests' output. + destination = destination.as_strided( + (end - row, *destination.shape), (destination.numel(), *destination.stride()) + ) + source = audio_tensor[row:end] + else: + source = audio_tensor[row] + # Crop time before reshaping: flattening first mixes stereo + # channels whenever a row is shorter than the decode batch. + destination.copy_(source[..., : destination.shape[-1]].view_as(destination)) + out.update((slot, output_buffers[slot]) for slot in slots[row:end]) + row = end + return out + + if self._gpu_output: + # A graph's static output is overwritten on every replay, including + # multiple replays within one forward. Snapshot on the producer + # stream before returning; record_stream alone cannot prevent this + # overwrite. MRV2 owns the subsequent pinned, asynchronous D2H. + audio = audio_tensor.detach().to(dtype=torch.float32, copy=True) + else: + audio = audio_tensor.detach().to("cpu", torch.float32) out: dict[int, torch.Tensor] = {} for row, slot in enumerate(slots): audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame @@ -270,6 +430,7 @@ class MossTTSCodecDecoder(nn.Module): requires_exact_input_shape: bool = True _OUTPUT_SAMPLE_RATE: int = 24_000 + _first_chunk_fast_path: MossFirstChunkFastPath | None = None def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -292,20 +453,38 @@ class MossTTSCodecDecoder(nn.Module): self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) self._stream_session: _MossCodecStreamSession | None = None + self._codec_stream: torch.cuda.Stream | None = None scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)) self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", default=0) self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0) self._stream_max_step_frames: int = self._stream_chunk_frames or 100 + # Opt-in: size codec attention rings as context + one chunk so chunked + # streaming reproduces whole-sequence decoding (costs ring memory). + self._stream_ring_headroom: bool = bool(self._connector_int("codec_ring_headroom", default=0)) self._stream_req_slots: dict[str, int] = {} + # Opt-in: decode each stream's first chunk from the receive thread. + self._first_chunk_fast = bool(self._connector_int("codec_first_chunk_fast_path", default=0)) + self._first_chunk_fast_path: MossFirstChunkFastPath | None = None self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) + self._gpu_stream_output = ( + self._async_chunk + and bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + and getattr(cfg, "model_type", None) == "moss_tts_local" + # Diagnostic switch: 0 copies each decode group's batch to the host synchronously. + and bool(self._connector_int("codec_gpu_stream_output", default=1)) + ) self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config() - # codec_chunk_ramp ladder sizes join the streaming graph frame-size - # set so every ladder (B, T) combination is captured during warmup. + ramp = parse_chunk_ramp(self._connector_extra(), steady=self._stream_chunk_frames) + if ramp: + # The sender gives the ramp precedence over initial chunk size. + # The dedicated first-chunk graph must use that same first shape. + self._initial_stream_chunk_frames = ramp[0] + self._stream_max_step_frames = max( + self._stream_max_step_frames, self._initial_stream_chunk_frames, *(ramp or []) + ) self._streaming_graph_frame_sizes = _resolve_streaming_graph_frame_sizes( - self._initial_stream_chunk_frames, - self._stream_chunk_frames, - self._connector_extra_cfg(), + self._initial_stream_chunk_frames, self._stream_chunk_frames, self._connector_extra() ) # ------------------------------------------------------------------ @@ -337,6 +516,41 @@ class MossTTSCodecDecoder(nn.Module): inputs_embeds: torch.Tensor | None = None, runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: Any, + ) -> OmniOutput: + if self._gpu_stream_output and input_ids is not None and input_ids.is_cuda: + if self._codec_stream is None: + self._codec_stream = torch.cuda.Stream(device=input_ids.device) + main_stream = torch.cuda.current_stream(input_ids.device) + # Runner input storage is reused by the next batch on main_stream. + # Own it before allowing that preparation to overlap codec decode. + owned_input_ids = input_ids.clone() + with torch.cuda.stream(self._codec_stream): + self._codec_stream.wait_stream(main_stream) + owned_input_ids.record_stream(self._codec_stream) + output = self._forward( + owned_input_ids, + positions, + intermediate_tensors, + inputs_embeds, + runtime_additional_information, + **kwargs, + ) + if isinstance(output.multimodal_outputs, PackedOutputSnapshot): + output.multimodal_outputs.record_producer_event(self._codec_stream) + return output + return self._forward( + input_ids, positions, intermediate_tensors, inputs_embeds, runtime_additional_information, **kwargs + ) + + @torch.no_grad() + def _forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, ) -> OmniOutput: """Decode audio VQ codes to waveform. @@ -493,13 +707,37 @@ class MossTTSCodecDecoder(nn.Module): audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav if streaming_work: - for i, wav in self._decode_streaming_batch(streaming_work).items(): + stream_output_slab = None + stream_output_buffers: dict[int, torch.Tensor] = {} + if self._gpu_stream_output and device.type == "cuda": + total_elements = 0 + for output_index, _, codes_nq_t, _ in streaming_work: + samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate + total_elements += samples * (self._n_channels if self._n_channels > 1 else 1) + if total_elements: + stream_output_slab = torch.empty(total_elements, dtype=torch.float32, device=device) + offset = 0 + for output_index, _, codes_nq_t, _ in streaming_work: + samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate + elements = samples * (self._n_channels if self._n_channels > 1 else 1) + shape = (self._n_channels, samples) if self._n_channels > 1 else (samples,) + stream_output_buffers[output_index] = stream_output_slab.narrow(0, offset, elements).view(shape) + offset += elements + for i, wav in self._decode_streaming_batch( + streaming_work, + output_buffers=stream_output_buffers or None, + ).items(): audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": audios, "sr": srs}, - ) + payload = {"model_outputs": audios, "sr": srs} + if self._gpu_stream_output and device.type == "cuda": + # One transfer per dtype/device preserves ragged request lengths + # without issuing one D2H per waveform. A fresh slot owns each + # output until the runner completes its copy; no reusable ring. + packed = pack_output_snapshot(payload, {}, max_buckets=0, reuse_existing_storage=True) + if packed is not None: + payload = packed + return OmniOutput(text_hidden_states=None, multimodal_outputs=payload) @staticmethod def _normalize_seq_token_counts(value: Any) -> list[int] | None: @@ -546,25 +784,104 @@ class MossTTSCodecDecoder(nn.Module): vllm_config=self.vllm_config, graph_batch_sizes=self._streaming_graph_batch_sizes, graph_frame_sizes=self._streaming_graph_frame_sizes, + gpu_output=self._gpu_stream_output, + chunk_frames=self._stream_max_step_frames, + ring_headroom=self._stream_ring_headroom, + fast_batch_sizes=[1, 2, 4, 8] if self._first_chunk_fast else None, + fast_frames=self._initial_stream_chunk_frames if self._first_chunk_fast else 0, ) + if self._stream_session._fast_wrapper is not None: + self._first_chunk_fast_path = MossFirstChunkFastPath( + self._stream_session, + self._stream_session._fast_wrapper, + n_vq=self._n_vq, + frames=self._initial_stream_chunk_frames, + codebook_size=int(self._codec.config.codebook_size), + samples_per_frame=int(self._codec.downsample_rate), + n_channels=self._n_channels, + sample_rate=self._sr_tensor, + device=next(self._codec.parameters()).device, + gate_main=bool(self._connector_int("codec_first_chunk_gate", default=0)), + ) + logger.info( + "MOSS-TTS codec first-chunk fast path ready: B=%s T=%d", + self._stream_session._fast_wrapper.capture_sizes, + self._initial_stream_chunk_frames, + ) + elif self._first_chunk_fast: + logger.warning("MOSS-TTS codec first-chunk fast path requested but unavailable; using the main path.") + if self._stream_ring_headroom: + logger.info( + "MOSS-TTS codec attention rings sized with headroom for %d-frame chunks", + self._stream_max_step_frames, + ) return self._stream_session + def bind_first_chunk_fast_path(self, sink: FirstAudioSink) -> Any: + """Start the fast path with an output sink; returns the receive hook.""" + fast = self._first_chunk_fast_path + if fast is None: + return None + fast.bind(sink) + return self._submit_first_chunk + + def _submit_first_chunk(self, request_id: str, handle: Any, payload: Any) -> bool: + fast = self._first_chunk_fast_path + if fast is None or not isinstance(payload, dict): + return False + codes = payload.get("codes") + codes = codes.get("audio") if isinstance(codes, dict) else None + if codes is None: + return False + meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} + key = meta.get("req_id") + if isinstance(key, (list, tuple)): + key = key[0] if key else None + if key is None: + key = getattr(handle, "external_req_id", None) or request_id + return fast.submit(request_id, str(key), codes, self._stream_req_slots) + def _decode_streaming_batch( self, items: list[tuple[int, str, torch.Tensor, bool]], + *, + output_buffers: dict[int, torch.Tensor] | None = None, ) -> dict[int, torch.Tensor]: session = self._ensure_stream_session() if session is None: return {} outputs: dict[int, torch.Tensor] = {} + unsent_first_audio: 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 + # Terminal tails join the regular padded frame bucket when every + # execution batch width has a graph, so a step stays one replay + # instead of one eager or plain-graph decode per distinct tail length. coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity + fast = self._first_chunk_fast_path for output_index, request_id, codes_nq_t, finished in items: slot = self._stream_req_slots.get(request_id) + if fast is not None and slot is not None: + fast.order_after(slot) + if fast.take_decoded(request_id): + unsent = fast.take_unsent_audio(request_id) + if unsent is not None: + unsent_first_audio[output_index] = unsent + # The fast path already emitted this chunk's audio and + # advanced the slot state past it. + codes_nq_t = codes_nq_t[:, fast.frames :] + if codes_nq_t.shape[1] == 0: + if unsent is not None: + outputs[output_index] = unsent + if finished: + self._finish_stream_request(request_id, session, slot) + continue + if output_buffers is not None: + # Sized for the untrimmed chunk; decode into a fresh tensor. + output_buffers.pop(output_index, None) if slot is None: slot = session.acquire() if slot is None: @@ -580,9 +897,11 @@ class MossTTSCodecDecoder(nn.Module): slot, codes_nq_t, terminal=finished, + output_buffer=output_buffers.get(output_index) if output_buffers is not None else None, ) if wav is not None: - outputs[output_index] = wav + prefix = unsent_first_audio.pop(output_index, None) + outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav if finished: self._finish_stream_request( request_id, @@ -593,21 +912,33 @@ class MossTTSCodecDecoder(nn.Module): continue 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: + # Only terminal rows may be padded: their state is reset after + # the step, and the output is cropped to the exact frame count. 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)) + if fast is not None and grouped: + fast.gate() for group in grouped.values(): plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} terminal_slots = {slot for _, _, slot, _, finished in group if finished} - decoded = session.step(plan, terminal_slots=terminal_slots) + step_buffers = ( + { + slot: output_buffers[output_index] + for output_index, _, slot, _, _ in group + if output_buffers is not None and output_index in output_buffers + } + if output_buffers is not None + else None + ) + decoded = session.step(plan, terminal_slots=terminal_slots, output_buffers=step_buffers) for output_index, request_id, slot, _, finished in group: wav = decoded.get(slot) if wav is not None: - outputs[output_index] = wav + prefix = unsent_first_audio.pop(output_index, None) + outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav if finished: self._finish_stream_request( request_id, @@ -625,6 +956,7 @@ class MossTTSCodecDecoder(nn.Module): codes_nq_t: torch.Tensor, *, terminal: bool = False, + output_buffer: torch.Tensor | None = None, ) -> torch.Tensor | None: if codes_nq_t.numel() == 0: return None @@ -633,13 +965,24 @@ class MossTTSCodecDecoder(nn.Module): for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): chunk = codes_nq_t[:, start : start + max_step_frames] is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1]) + chunk_samples = int(chunk.shape[1]) * self._codec.downsample_rate + chunk_buffer = ( + output_buffer[ + ..., start * self._codec.downsample_rate : start * self._codec.downsample_rate + chunk_samples + ] + if output_buffer is not None + else None + ) decoded = session.step( {slot: chunk}, terminal_slots={slot} if is_terminal_chunk else None, + output_buffers={slot: chunk_buffer} if chunk_buffer is not None else None, ) wav = decoded.get(slot) if wav is not None: parts.append(wav) + if output_buffer is not None: + return output_buffer if not parts: return None return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] @@ -665,29 +1008,30 @@ class MossTTSCodecDecoder(nn.Module): terminal payload, so the runner calls this hook from its finished-request path to avoid leaking stream slots and buffered codes. """ - session = self._stream_session - for req_id in finished_req_ids: - request_id = str(req_id) - slot = self._stream_req_slots.get(request_id) - if slot is None: - continue - if session is not None: - self._finish_stream_request(request_id, session, slot) - else: - self._stream_req_slots.pop(request_id, None) + # Cancellation and terminal cleanup share the decoder's stream order; + # a reused slot must not be reset ahead of an outstanding decode. + context = torch.cuda.stream(self._codec_stream) if self._codec_stream is not None else nullcontext() + with context: + session = self._stream_session + for req_id in finished_req_ids: + request_id = str(req_id) + fast = self._first_chunk_fast_path + slot = ( + fast.get_request_slot(request_id, self._stream_req_slots) + if fast is not None + else self._stream_req_slots.get(request_id) + ) + if slot is None: + continue + if fast is not None: + fast.order_after(slot) + fast.forget(request_id) + if session is not None: + self._finish_stream_request(request_id, session, slot) + else: + self._stream_req_slots.pop(request_id, None) - def _connector_int(self, name: str, default: int = 0) -> int: - model_cfg = getattr(self.vllm_config, "model_config", None) - connector_cfg = getattr(model_cfg, "stage_connector_config", None) - if isinstance(connector_cfg, dict): - extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) - else: - extra_cfg = getattr(connector_cfg, "extra", None) - if isinstance(extra_cfg, dict) and name in extra_cfg: - return int(extra_cfg[name]) - return default - - def _connector_extra_cfg(self) -> dict: + def _connector_extra(self) -> dict[str, Any]: model_cfg = getattr(self.vllm_config, "model_config", None) connector_cfg = getattr(model_cfg, "stage_connector_config", None) if isinstance(connector_cfg, dict): @@ -696,6 +1040,9 @@ class MossTTSCodecDecoder(nn.Module): extra_cfg = getattr(connector_cfg, "extra", None) return extra_cfg if isinstance(extra_cfg, dict) else {} + def _connector_int(self, name: str, default: int = 0) -> int: + return int(self._connector_extra().get(name, default)) + def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]: if getattr(self.vllm_config.model_config, "enforce_eager", True): return [] @@ -816,15 +1163,33 @@ class MossTTSCodecDecoder(nn.Module): 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": + if attention_backend not in {"triton", "triton_slot"} 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 + if attention_backend == "triton_slot": + from vllm_omni.model_executor.models.moss_tts.slot_attention import ( + slot_ring_attention, + slot_ring_attention_rows, + ) + + skip_empty_tiles = bool(self._connector_int("codec_skip_empty_attention_tiles", default=0)) + fused_slots = bool(self._connector_int("codec_fused_slot_attention", default=0)) + if skip_empty_tiles and fused_slots: + raise ValueError("codec_skip_empty_attention_tiles requires unfused slot attention") + if skip_empty_tiles: + slot_ring_attention = partial(slot_ring_attention, skip_empty_tiles=True) + logger.info("MOSS codec slot attention: skipping empty tiles for T <= 32") + 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") + if attention_backend == "triton_slot": + module._slot_attention = slot_ring_attention + if self._connector_int("codec_fused_slot_attention", default=0): + module._slot_attention_rows = slot_ring_attention_rows + logger.info("Enabled codec attention backend=%s", attention_backend) 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 c3f226fd1..ae837eddb 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 @@ -21,7 +21,8 @@ checkpoint 1:1 so ``load_weights()`` needs no remapping. from __future__ import annotations -from collections.abc import Sequence +import os +from collections.abc import Callable, Sequence import torch import torch.nn as nn @@ -52,6 +53,7 @@ class _MossTTSLocalAttention(nn.Module): self.n_head = n_head self.head_dim = hidden_size // n_head self.embed_dim = hidden_size + self._short_attention = os.environ.get("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "0") == "1" self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True) self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True) inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim)) @@ -134,7 +136,19 @@ class _MossTTSLocalAttention(nn.Module): 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) + if ( + self._short_attention + and kv_cache is not None + and query.is_cuda + and query.dtype == torch.bfloat16 + and self.head_dim <= 128 + and key.shape[2] <= 16 + ): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + attn_output = local_short_attention(query, key, value) + else: + 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) @@ -193,6 +207,7 @@ class MossTTSLocalDepthTransformer(nn.Module): self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) self._compiled_forward_prefix = None + self._compiled_audio_sampler = None def _forward_prefix( self, @@ -216,6 +231,14 @@ class MossTTSLocalDepthTransformer(nn.Module): dynamic=True, options={"epilogue_fusion": False}, ) + if os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1": + # Keep the existing top-k/top-p algorithm and torch RNG. Explicit + # per-request generators use the original helper below. The + # binary continue/stop head is deliberately unchanged. + self._compiled_audio_sampler = torch.compile( + _sample_token, fullgraph=True, dynamic=True, options={"fallback_random": True} + ) + logger.info("MOSS-TTS Local compiled audio-channel sampler enabled") logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile") def _run_prefix( @@ -239,6 +262,7 @@ class MossTTSLocalDepthTransformer(nn.Module): do_sample: bool, generator: torch.Generator | None, generators: Sequence[torch.Generator | None] | None = None, + sample_token: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: """Compute channel logits, apply repetition penalty, sample, store.""" channel_logits = audio_lm_heads[channel_index](local_hidden).float() @@ -250,7 +274,7 @@ class MossTTSLocalDepthTransformer(nn.Module): pos = sel > 0 sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty) channel_logits.index_copy_(-1, hist_t, sel) - channel_token = _sample_token( + channel_token = (sample_token or _sample_token)( channel_logits, temperature, top_k, @@ -351,6 +375,7 @@ class MossTTSLocalDepthTransformer(nn.Module): do_sample=do_sample, generator=generator, generators=generators, + sample_token=self._compiled_audio_sampler if generator is None and generators is None else None, ) if channel_index + 1 < n_vq: 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 bc4a5b716..303724e2e 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 @@ -37,6 +37,7 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth impo ) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker.sampling_utils import get_tts_local_seed logger = init_logger(__name__) @@ -1675,6 +1676,31 @@ class MossTTSLocalTalkerForGeneration(nn.Module): return input_embeds_out, output_codes # ------------------------------------------------------------------ + # MRV2 reads model-owned capabilities; retain the V1 hook and platform + # graph-safety setting as the canonical implementation. + mtp = talker_mtp + get_mtp_seed = staticmethod(get_tts_local_seed) + + def create_mrv2_model_state(self, vllm_config, encoder_cache, device): + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + + state_cls = OmniModelState + if getattr(self.config, "mrv2_gpu_slot_state", False): + from .local_model_state import MossLocalModelState + + state_cls = MossLocalModelState + return state_cls(vllm_config, self, encoder_cache, device) + + @property + def mtp_output_key(self) -> tuple[str, str]: + """Return the request-buffer destination for one local audio frame.""" + return self.talker_mtp_output_key + + @property + def mtp_graph_safe(self) -> bool: + """Honor the same platform graph-safety override as the V1 runner.""" + return self.talker_mtp_graph_safe + # Package runner-generated audio frames # ------------------------------------------------------------------ diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 0d6151313..766c0d136 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -97,6 +97,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=0, model_stage="moss_tts_local", + supports_native_mrv2_data_plane=True, execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, @@ -114,6 +115,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=1, model_stage="moss_tts_local_codec", + supports_native_mrv2_data_plane=True, execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py index b1516d90a..e492cf4cd 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -25,7 +25,10 @@ so importing it from the API-server process does not pull the talker/codec. from __future__ import annotations import asyncio +import contextlib import hashlib +import os +from collections import OrderedDict from collections.abc import Awaitable, Callable from functools import lru_cache from typing import Any @@ -35,22 +38,93 @@ import torch from numpy.typing import NDArray from vllm.logger import init_logger +from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import ( + FINE_BATCH_SIZES, + FINE_BUCKET_SECONDS, + MossReferenceEncoderGraphs, + install_windowed_attention, +) +from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import ( + SharedReferenceEncoderClient, + SharedReferenceEncoderHost, + SharedReferenceEncoderStartupError, + elect_host, +) + logger = init_logger(__name__) # Coalescing defaults. Kept as module constants (not env/CLI) to match the # ``_REF_AUDIO_RESOLVE_CACHE_MAX_*`` convention in serving_speech.py. _REF_ENCODE_BATCH_WINDOW_MS = 10.0 -_REF_ENCODE_MAX_BATCH = 8 +_REF_ENCODE_MAX_BATCH = 32 +# Experiment switches for cold-reference latency (see the A/B notes): +# batch window in ms; batches in flight per API process with a shared encoder; +# "high" puts reference encodes on the highest-priority CUDA stream. +_REF_ENCODE_WINDOW_ENV = "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS" +_REF_ENCODE_INFLIGHT_ENV = "VLLM_OMNI_MOSS_REF_INFLIGHT" +_REF_ENCODE_PRIORITY_ENV = "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY" +# "1" resamples and loudness-normalizes clips on the encoder's GPU instead of +# the API process's CPU (not bit-identical to the CPU path). +_REF_ENCODE_GPU_PREP_ENV = "VLLM_OMNI_MOSS_REF_GPU_PREP" _INT32_MAX = 2**31 +# Inline (data: URI) references already resolved and encoded, keyed by the +# serving resolve key. Codes are small; the URI itself is not retained. +_INLINE_INDEX_MAX_ENTRIES = 4096 +# Optional directory (e.g. under /dev/shm) that API processes of one server +# share, so an inline reference encoded by any frontend is reused by all. +_SHARED_CODES_DIR_ENV = "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR" +# "0" keeps one reference encoder per API process even with a shared directory. +_SHARED_ENCODER_ENV = "VLLM_OMNI_MOSS_REF_SHARED_ENCODER" +_shared_role: str | None = None +# Concurrent encode workers in the shared host, each with its own CUDA graphs. +_SHARED_WORKERS_ENV = "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS" +_SHARED_WORKER_BATCH_SECONDS = 96.0 + + +def shared_encoder_role() -> str: + """``"host"``, ``"client"`` or ``""`` (no shared encoder) for this process. + + Decided once, before the processor places its tokenizer on the GPU, so + client processes can leave it on the CPU. + """ + global _shared_role + if _shared_role is None: + shared_dir = os.environ.get(_SHARED_CODES_DIR_ENV) + if not shared_dir or os.environ.get(_SHARED_ENCODER_ENV, "1") == "0": + _shared_role = "" + else: + os.makedirs(shared_dir, exist_ok=True) + _shared_role = "host" if elect_host(shared_dir) else "client" + return _shared_role + + +# Opt in to reference graphs; existing deployments keep the eager encoder. +_REF_ENCODE_GRAPHS_ENV = "VLLM_OMNI_MOSS_REF_GRAPHS" + + +def reference_graphs_enabled() -> bool: + return os.environ.get(_REF_ENCODE_GRAPHS_ENV, "0") == "1" + + +def reference_stream_priority() -> int: + """CUDA stream priority for reference encodes (lower is more urgent).""" + if os.environ.get(_REF_ENCODE_PRIORITY_ENV, "") == "high" and torch.cuda.is_available(): + return torch.cuda.Stream.priority_range()[1] + return 0 + + +# "sdpa" keeps the tokenizer's dense masked SDPA attention. By default the +# reference encoder attends through a local-window flash kernel, the backend +# the tokenizer config selects (flash_attention_2) when flash-attn is present. +_REF_ENCODE_ATTN_ENV = "VLLM_OMNI_MOSS_REF_ATTN" +# "0" captures the encoder's eager kernels. By default the captured encoder is +# compiled first, which fuses its many small elementwise kernels. +_REF_ENCODE_COMPILE_ENV = "VLLM_OMNI_MOSS_REF_COMPILE" _Waveform = NDArray[np.float32] | list[float] -def _sha1(s: str) -> str: - return hashlib.sha1((s or "").encode("utf-8")).hexdigest() - - @lru_cache(maxsize=16) def _reference_resampler(sr: int, sr_target: int) -> torch.nn.Module: import torchaudio @@ -72,6 +146,13 @@ def _prep_wav_sync(waveform: _Waveform, sr: int, sr_target: int) -> torch.Tensor return wav +def _inline_reference_key(ref_str: str) -> str | None: + """SHA-1 of a ``data:`` URI, matching the serving resolve key for it.""" + if not isinstance(ref_str, str) or ref_str[:5].lower() != "data:": + return None + return hashlib.sha1(ref_str.encode("utf-8")).hexdigest() + + def _to_compact_codes(codes: torch.Tensor) -> torch.Tensor: """Downcast RVQ codes to int32 for compact caching.""" codes = codes.detach().cpu().contiguous() @@ -119,18 +200,26 @@ class _RefEncodeBatcher: *, window_ms: float, max_batch: int, + max_inflight: int = 1, ): # encode_batch_fn: sync, takes [(waveform, sr), ...], returns a list of # (codes_tensor | Exception) aligned to the input order. self._encode_batch_fn = encode_batch_fn self._window_s = max(0.0, float(window_ms) / 1000.0) self._max_batch = max(1, int(max_batch)) + # Batches encoded at once. Clips that arrive while every slot is busy + # are queued and form the next batch. + self._max_inflight = max(1, int(max_inflight)) self._queue: asyncio.Queue | None = None + self._slots: asyncio.Semaphore | None = None + self._inflight: set[asyncio.Task] = set() self._drainer: asyncio.Task | None = None def _ensure_started(self) -> None: if self._queue is None: self._queue = asyncio.Queue() + if self._slots is None: + self._slots = asyncio.Semaphore(self._max_inflight) if self._drainer is None or self._drainer.done(): self._drainer = asyncio.create_task(self._drain_loop()) @@ -141,16 +230,27 @@ class _RefEncodeBatcher: return await fut async def _drain_loop(self) -> None: - assert self._queue is not None + assert self._queue is not None and self._slots is not None while True: - first = await self._queue.get() - jobs = None + await self._slots.acquire() try: - jobs = await self._coalesce(first) - await self._run_batch(jobs) - finally: - # Do not retain the completed batch while waiting for new work. - del first, jobs + jobs = await self._coalesce(await self._queue.get()) + except BaseException: + self._slots.release() + raise + task = asyncio.create_task(self._run_and_release(jobs)) + self._inflight.add(task) + task.add_done_callback(self._inflight.discard) + # Do not retain the batch here while waiting for new work. + del jobs + + async def _run_and_release(self, jobs: list[tuple[_Waveform, int, asyncio.Future]]) -> None: + try: + await self._run_batch(jobs) + finally: + del jobs + assert self._slots is not None + self._slots.release() async def _coalesce(self, first: tuple) -> list[tuple]: """Group ``first`` with jobs arriving within the batch window.""" @@ -200,6 +300,8 @@ class _RefEncodeBatcher: except (asyncio.CancelledError, Exception): # noqa: BLE001 pass self._drainer = None + for task in list(self._inflight): + task.cancel() class MossReferenceEncoder: @@ -213,7 +315,7 @@ class MossReferenceEncoder: n_vq: int, sr_target: int, speaker_cache: Any, - window_ms: float = _REF_ENCODE_BATCH_WINDOW_MS, + window_ms: float | None = None, max_batch: int = _REF_ENCODE_MAX_BATCH, ): self._processor = processor @@ -224,8 +326,33 @@ class MossReferenceEncoder: # model_type namespaces the whole family so a moss_tts server never # collides with another model's speaker-cache entries. self._model_type = f"moss_tts_{variant}_nq{int(n_vq)}" - self._inflight: dict[str, asyncio.Task] = {} - self._batcher = _RefEncodeBatcher(self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch) + self._inflight: dict[str | tuple[str, str], asyncio.Task] = {} + self._inline_index: OrderedDict[str, torch.Tensor] = OrderedDict() + self._shared_codes_dir = os.environ.get(_SHARED_CODES_DIR_ENV) or None + if self._shared_codes_dir: + os.makedirs(self._shared_codes_dir, exist_ok=True) + # Realtime encodes with its own codec and never shares the encoder. + self._shares_encoder = bool(self._shared_codes_dir) and type(self)._encode_prepared is ( + MossReferenceEncoder._encode_prepared + ) + self._shared_host: SharedReferenceEncoderHost | None = None + self._shared_client: SharedReferenceEncoderClient | None = None + self._graphs: MossReferenceEncoderGraphs | None = None + self._attention_installed = False + self._graphs_enabled = reference_graphs_enabled() + if window_ms is None: + window_ms = float(os.environ.get(_REF_ENCODE_WINDOW_ENV, _REF_ENCODE_BATCH_WINDOW_MS)) + # A shared encoder serves several batches at once (one per host + # worker); a process's own graphs replay one batch at a time. + inflight = 1 + if self._shares_encoder and os.environ.get(_SHARED_ENCODER_ENV, "1") != "0": + inflight = max(1, int(os.environ.get(_REF_ENCODE_INFLIGHT_ENV, "1"))) + self._local_stream: torch.cuda.Stream | None = None + self._gpu_prep = os.environ.get(_REF_ENCODE_GPU_PREP_ENV, "0") == "1" + self._gpu_resamplers: dict[torch.device, torch.nn.Module] = {} + self._batcher = _RefEncodeBatcher( + self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch, max_inflight=inflight + ) def _make_cache_key(self, name: str, created_at: int) -> tuple: return self._speaker_cache.make_cache_key(name, model_type=self._model_type, created_at=int(created_at)) @@ -251,6 +378,7 @@ class MossReferenceEncoder: caller salts those requests with ``voice_created_at`` instead). """ voice_name, created_at = _registered_voice(voice_name, voice_created_at) + inline_key: str | None = None if voice_name: # A named voice has a stable key that does not depend on the @@ -260,6 +388,16 @@ class MossReferenceEncoder: if cached is not None: return _clone_out(cached["codes"]), None else: + inline_key = _inline_reference_key(ref_str) + if inline_key is not None: + codes = self._inline_index.get(inline_key) + if codes is None: + codes = self._load_shared_codes(inline_key) + if codes is not None: + # A data: URI is immutable content, and its serving + # resolve key is this same digest: skip the resolve hop. + self._inline_index.move_to_end(inline_key) + return _clone_out(codes), inline_key # Anonymous refs have no pre-resolve hot path: the content key must # come from the resolve itself (mtime/size from a single stat) so an # on-disk edit invalidates the cached codes; the flight body @@ -267,17 +405,55 @@ class MossReferenceEncoder: # when the resolve cache is warm. The flight key is the request-side # reference (not the content hash), so concurrent requests for the # same ref_str also share the resolve/download, not just the encode. - flight_key = "ref:" + _sha1(ref_str) + # Key by the locator itself: equality is exact, and hashing a + # large inline data URI once per request object is much cheaper + # than a separate SHA-1 of its UTF-8 copy on the event loop. + flight_key = ("ref", ref_str) codes, resolve_key = await self._single_flight( flight_key, lambda: self._resolve_and_encode(ref_str, resolve_ref_audio, get_artifact_key, voice_name, created_at), ) + if not voice_name and inline_key is not None and resolve_key == inline_key: + # Only index when the resolver's key is exactly this digest. + self._remember_inline(inline_key, codes) + self._store_shared_codes(inline_key, codes) return _clone_out(codes), resolve_key + def _remember_inline(self, key: str, codes: torch.Tensor) -> None: + self._inline_index[key] = codes + self._inline_index.move_to_end(key) + while len(self._inline_index) > _INLINE_INDEX_MAX_ENTRIES: + self._inline_index.popitem(last=False) + + def _load_shared_codes(self, key: str) -> torch.Tensor | None: + if not self._shared_codes_dir: + return None + try: + array = np.load(os.path.join(self._shared_codes_dir, key + ".npy"), allow_pickle=False) + except (OSError, ValueError): + return None + if array.ndim != 2 or array.shape[1] != self._n_vq or array.dtype not in (np.int32, np.int64): + return None + codes = torch.from_numpy(array) + self._remember_inline(key, codes) + return codes + + def _store_shared_codes(self, key: str, codes: torch.Tensor) -> None: + if not self._shared_codes_dir: + return + path = os.path.join(self._shared_codes_dir, key + ".npy") + tmp = f"{path}.{os.getpid()}.tmp" + try: + with open(tmp, "wb") as handle: + np.save(handle, codes.detach().cpu().numpy(), allow_pickle=False) + os.replace(tmp, path) + except OSError: + logger.warning("MOSS ref codes: could not write shared entry %s", path, exc_info=True) + async def _single_flight( self, - flight_key: str, + flight_key: str | tuple[str, str], start_flight: Callable[[], Awaitable[tuple[torch.Tensor, str]]], ) -> tuple[torch.Tensor, str]: """Join the in-flight encode for ``flight_key``, starting one if absent. @@ -300,7 +476,7 @@ class MossReferenceEncoder: # a duplicate resolve/encode instead of joining this one. Identity # guard: only drop the slot if it still holds *our* task (a later # request may have replaced it after ours completed). - def _retire(t: asyncio.Task, key: str = flight_key) -> None: + def _retire(t: asyncio.Task, key: str | tuple[str, str] = flight_key) -> None: if self._inflight.get(key) is t: self._inflight.pop(key, None) @@ -350,29 +526,228 @@ class MossReferenceEncoder: prepared_idx: list[int] = [] for i, (waveform, sr) in enumerate(payload): try: - prepared.append(_prep_wav_sync(waveform, sr, self._sr_target)) + prepared.append(self._encoder_input(_prep_wav_sync(waveform, sr, self._sr_target))) prepared_idx.append(i) except Exception as exc: # noqa: BLE001 — isolate this clip's failure results[i] = exc if not prepared: return results - try: - codes_list = self._encode_prepared(prepared) - for local_i, orig_i in enumerate(prepared_idx): - results[orig_i] = codes_list[local_i] - except Exception: # noqa: BLE001 — batch forward failed; retry item-by-item - logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared)) - for local_i, orig_i in enumerate(prepared_idx): - try: - results[orig_i] = self._encode_prepared([prepared[local_i]])[0] - except Exception as exc: # noqa: BLE001 — isolate this clip's failure - results[orig_i] = exc + for orig_i, codes in zip(prepared_idx, self._encode_shared_or_local(prepared)): + results[orig_i] = codes return results - def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: - with torch.no_grad(): - return self._processor.encode_audios_from_wav(prepared, sampling_rate=self._sr_target, n_vq=self._n_vq) + def prepare(self) -> None: + """Capture the CUDA graphs (or start the shared encoder) before the first request.""" + role = shared_encoder_role() if self._shares_encoder else "" + if role == "host": + self._host().wait_until_ready() + elif role == "client": + if self._shared_client is None: + self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir) + self._shared_client.wait_until_ready() + elif role == "": + self._reference_graphs() + + def _host(self) -> SharedReferenceEncoderHost: + if self._shared_host is None: + # Graph capture runs in the background; clips queue until it ends. + self._shared_host = SharedReferenceEncoderHost( + self._shared_codes_dir, + self._host_workers(), + stream_priority=reference_stream_priority(), + batch_window_ms=float(os.environ.get("VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", "0")), + ) + return self._shared_host + + def _encode_shared_or_local(self, prepared: list[torch.Tensor]) -> list: + role = shared_encoder_role() if self._shares_encoder else "" + if role == "host": + return self._host().encode(prepared) + if role == "client": + if self._shared_client is None: + self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir) + try: + return self._shared_client.encode(prepared) + except SharedReferenceEncoderStartupError: + raise + except Exception: # noqa: BLE001 — the host is unreachable: encode here + logger.warning("MOSS shared reference encoder unavailable; encoding locally", exc_info=True) + self._place_tokenizer_for_local_encode() + with self._local_stream_context(): + return self._encode_local(prepared) + + def _local_stream_context(self): + tokenizer = getattr(self._processor, "audio_tokenizer", None) + params = next(tokenizer.parameters(), None) if isinstance(tokenizer, torch.nn.Module) else None + if params is None or params.device.type != "cuda": + return contextlib.nullcontext() + if self._local_stream is None: + self._local_stream = torch.cuda.Stream(device=params.device, priority=reference_stream_priority()) + return torch.cuda.stream(self._local_stream) + + def _host_workers(self) -> list: + """Encode workers of the shared host, each with its own graphs and stream.""" + count = max(1, int(os.environ.get(_SHARED_WORKERS_ENV, "4"))) + if not self._graph_capable(): + return [(self._encode_local, None)] * count + workers = [] + for _ in range(count): + graphs = MossReferenceEncoderGraphs( + self._processor.audio_tokenizer, + n_vq=self._n_vq, + batch_sizes=FINE_BATCH_SIZES, + bucket_seconds=FINE_BUCKET_SECONDS, + max_batch_seconds=_SHARED_WORKER_BATCH_SECONDS, + compile_core=self._compile_graphs(), + batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1", + singleton_bucket_seconds=self._singleton_bucket_seconds(), + ) + workers.append(((lambda wavs, g=graphs: self._encode_local(wavs, g)), graphs.capture)) + return workers + + def _encode_local(self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None) -> list: + """One batched forward, retried per clip so one bad clip fails alone.""" + try: + return list(self._encode_prepared(prepared, graphs)) + except Exception: # noqa: BLE001 — batch forward failed; retry item-by-item + logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared)) + results: list = [] + for wav in prepared: + try: + results.append(self._encode_prepared([wav], graphs)[0]) + except Exception as exc: # noqa: BLE001 — isolate this clip's failure + results.append(exc) + return results + + def _place_tokenizer_for_local_encode(self) -> None: + """A client process keeps the tokenizer off the GPU until it must encode.""" + tokenizer = getattr(self._processor, "audio_tokenizer", None) + device = getattr(self._processor, "_vllm_omni_ref_encoder_device", None) + if tokenizer is not None and device is not None and next(tokenizer.parameters()).device != device: + self._processor.audio_tokenizer = tokenizer.to(device) + + def _encode_prepared( + self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None + ) -> list[torch.Tensor]: + # ``prepared`` clips come from _encoder_input, at ``_clip_rate()``. + graphs = graphs or self._reference_graphs() + if graphs is None: + with torch.no_grad(): + return self._processor.encode_audios_from_wav( + prepared, sampling_rate=self._clip_rate(), n_vq=self._n_vq + ) + codes = graphs.encode([self._finish_prepare(wav) for wav in prepared]) + missing = [i for i, c in enumerate(codes) if c is None] + if missing: + with torch.no_grad(): + eager = self._processor.encode_audios_from_wav( + [prepared[i] for i in missing], sampling_rate=self._clip_rate(), n_vq=self._n_vq + ) + for i, c in zip(missing, eager): + codes[i] = c + return codes + + def _graph_capable(self) -> bool: + """True when the tokenizer can run graphs; installs windowed attention once.""" + if not self._graphs_enabled: + return False + tokenizer = getattr(self._processor, "audio_tokenizer", None) + device = next(tokenizer.parameters()).device if tokenizer is not None else None + if device is None or device.type != "cuda" or not hasattr(tokenizer, "_codec_inference_autocast"): + self._graphs_enabled = False + return False + if not self._attention_installed and os.environ.get(_REF_ENCODE_ATTN_ENV, "flash") == "flash": + logger.info( + "MOSS reference encoder: %d attention layers use local-window flash attention", + install_windowed_attention( + tokenizer, + fa_version=int(os.environ.get("VLLM_OMNI_MOSS_REF_FA_VERSION", "2")), + cache_rope=os.environ.get("VLLM_OMNI_MOSS_REF_CACHE_ROPE", "0") == "1", + skip_padded_query_mask=os.environ.get("VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", "0") == "1", + ), + ) + self._attention_installed = True + return True + + @staticmethod + def _compile_graphs() -> bool: + return os.environ.get(_REF_ENCODE_COMPILE_ENV, "1") != "0" + + @staticmethod + def _singleton_bucket_seconds() -> tuple[float, ...] | None: + step_ms = int(os.environ.get("VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", "0")) + if step_ms == 0: + return None + if not 80 <= step_ms <= 1000: + raise ValueError("MOSS reference singleton bucket step must be 80..1000 ms, or 0 to disable") + return tuple(ms / 1000 for ms in range(3000, 10001, step_ms)) + + def _reference_graphs(self) -> MossReferenceEncoderGraphs | None: + """Capture on first use, on the encode worker thread that replays them.""" + if self._graphs is None: + if not self._graph_capable(): + return None + graphs = MossReferenceEncoderGraphs( + self._processor.audio_tokenizer, + n_vq=self._n_vq, + compile_core=self._compile_graphs(), + batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1", + singleton_bucket_seconds=self._singleton_bucket_seconds(), + ) + graphs.capture() + self._graphs = graphs + return self._graphs + + def _tokenizer_rate(self) -> int: + return int(getattr(getattr(self._processor, "model_config", None), "sampling_rate", self._sr_target)) + + def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor: + """A prepped clip at the tokenizer rate, one or two channels. + + This is the costly half of the processor's per-clip preparation; it runs + in the process that received the request, even with a shared encoder. + """ + import torchaudio + + if wav.ndim == 1: + wav = wav.unsqueeze(0) + if wav.shape[0] > 2: + wav = wav[:2] + if self._clip_rate() != self._sr_target: + # Channels resample independently, so a mono clip is resampled + # once before it is duplicated; the result is the same. + wav = torchaudio.functional.resample(wav, self._sr_target, self._tokenizer_rate()) + return wav + + def _clip_rate(self) -> int: + """Sample rate of the clips _encoder_input hands to the encoder.""" + return self._sr_target if self._gpu_prep else self._tokenizer_rate() + + def _finish_prepare(self, wav: torch.Tensor) -> torch.Tensor: + """The rest of the processor's preparation: two channels, normalized loudness.""" + if self._gpu_prep: + device = next(self._processor.audio_tokenizer.parameters()).device + wav = wav.to(device) + if self._clip_rate() != self._tokenizer_rate(): + wav = self._gpu_resampler(device)(wav) + if wav.shape[0] == 1: + wav = wav.repeat(2, 1) + return self._processor.loudness_normalize(wav) + + def _gpu_resampler(self, device: torch.device) -> torch.nn.Module: + resampler = self._gpu_resamplers.get(device) + if resampler is None: + import torchaudio + + # The same windowed-sinc filter as torchaudio.functional.resample's defaults. + resampler = torchaudio.transforms.Resample(self._sr_target, self._tokenizer_rate()).to(device) + self._gpu_resamplers[device] = resampler + return resampler + + def _processor_prepare(self, wav: torch.Tensor) -> torch.Tensor: + """Per-clip preparation of the processor's ``encode_audios_from_wav``.""" + return self._finish_prepare(self._encoder_input(wav)) async def aclose(self) -> None: """Release the batcher drainer + any in-flight encodes (tests/shutdown).""" @@ -385,7 +760,10 @@ class MossReferenceEncoder: class MossRealtimeReferenceEncoder(MossReferenceEncoder): """Use the realtime codec with the shared reference cache and batcher.""" - def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: + def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor: + return wav # the realtime codec takes the clip at the working rate + + def _encode_prepared(self, prepared: list[torch.Tensor], graphs: Any = None) -> list[torch.Tensor]: with torch.no_grad(): encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq) return [ diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py new file mode 100644 index 000000000..52dead755 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graphs for batched MOSS reference-audio encoding. + +Reference clips are encoded by the upstream MOSS audio tokenizer that the +processor loads. Its eager forward is dominated by kernel launches for the +small batches that serving produces. This runner replays graphs of the same +encoder and quantizer modules, captured per ``(batch, length)`` bucket. For a +given padded input the codes are bit-identical to the tokenizer's own +``_encode_frame``. Only the per-clip crop, which needs host values, runs +outside the graph. + +The encoder is causal, so padding a clip to its length bucket leaves its valid +frames unchanged up to kernel selection for the padded shape, the same effect +that batching already has when the tokenizer pads a batch to its longest clip. +""" + +from __future__ import annotations + +import math +import types +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +DEFAULT_BATCH_SIZES = (1, 2, 4, 8, 16, 32) +DEFAULT_BUCKET_SECONDS = (4.0, 6.0, 8.0, 10.0, 12.0, 16.0, 20.0) +# Fine steps keep padding small: a padded row or padded time is encoded in +# full. Worth the extra graphs where several encoders share one GPU. +FINE_BATCH_SIZES = (1, 2, 3, 4, 5, 6, 8) +FINE_BUCKET_SECONDS = (3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0, 16.0, 20.0) +# Graphs are captured only while batch * bucket stays within this much audio, +# which bounds the shared graph memory pool by its largest batch. +DEFAULT_MAX_BATCH_SECONDS = 256.0 + + +@torch.compile(dynamic=False) +def _reference_rope_factors(seqlen: int, head_dim: int, max_period: float, device: torch.device): + # Use the same compiled FP32 arithmetic as the encoder. Eager libdevice + # sin/cos and compiled sin/cos can round differently near a VQ boundary. + ds = torch.arange(head_dim // 2, device=device, dtype=torch.float32) + freqs = torch.exp(ds * (-math.log(max_period) * 2 / head_dim)) + angles = torch.arange(seqlen, device=device, dtype=torch.float32)[:, None] * freqs[None, :] + return torch.cos(angles), torch.sin(angles) + + +def _cached_reference_rope(self, q: torch.Tensor, k: torch.Tensor, cache: dict): + if self.rope is None: + return q, k + batch, heads, seqlen, dim = q.shape + key = (dim, self.rope.max_period, q.device) + if key not in cache or cache[key][0].shape[0] < seqlen: + cache[key] = _reference_rope_factors(seqlen, dim, self.rope.max_period, q.device) + cos, sin = cache[key] + cos, sin = cos[:seqlen], sin[:seqlen] + qr, qi = q.view(batch, heads, seqlen, dim // 2, 2).unbind(-1) + kr, ki = k.view(batch, heads, seqlen, dim // 2, 2).unbind(-1) + qr, qi, kr, ki = qr.float(), qi.float(), kr.float(), ki.float() + # Preserve the original arithmetic, casts, and B/H/T/D output layout. + qo = torch.stack(((qr * cos - qi * sin).to(q.dtype), (qr * sin + qi * cos).to(q.dtype)), dim=-1) + ko = torch.stack(((kr * cos - ki * sin).to(k.dtype), (kr * sin + ki * cos).to(k.dtype)), dim=-1) + return qo.view(batch, heads, seqlen, dim), ko.view(batch, heads, seqlen, dim) + + +def _windowed_attention( + self, + x: torch.Tensor, + input_lengths: torch.Tensor, + *, + sdpa, + fa_version: int = 2, + rope_cache: dict | None = None, + skip_padded_query_mask: bool = False, +) -> torch.Tensor: + """The tokenizer's non-streaming attention with a local-window flash kernel. + + Same projection, RoPE and padded-row zeroing as the tokenizer's SDPA path. + That path masks each query to ``context`` keys (distance 0 .. context-1) + inside a dense block; the flash kernel visits only that window. + """ + batch, seqlen, _ = x.shape + q, k, v = self._project_qkv(x) + q, k = self._apply_dense_rope(q, k) if rope_cache is None else _cached_reference_rope(self, q, k, rope_cache) + if q.dtype not in (torch.bfloat16, torch.float16): + return sdpa(x, input_lengths) + from vllm.vllm_flash_attn import flash_attn_varlen_func + + heads, dim = q.shape[1], q.shape[-1] + + def packed(t: torch.Tensor) -> torch.Tensor: + return t.transpose(1, 2).reshape(batch * seqlen, heads, dim) + + starts = torch.arange(0, (batch + 1) * seqlen, seqlen, device=x.device, dtype=torch.int32) + window = [self.context - 1, 0] if self.causal and self.context is not None else [-1, -1] + out = flash_attn_varlen_func( + packed(q), + packed(k), + packed(v), + max_seqlen_q=seqlen, + cu_seqlens_q=starts, + max_seqlen_k=seqlen, + cu_seqlens_k=starts, + causal=self.causal, + window_size=window, + fa_version=fa_version, + ).view(batch, seqlen, heads, dim) + if not skip_padded_query_mask: + valid = (torch.arange(seqlen, device=x.device).view(1, seqlen) < input_lengths.view(-1, 1)).view( + batch, seqlen, 1, 1 + ) + out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype)) + return out.reshape(batch, seqlen, self.embed_dim) + + +def _validate_unmasked_reference_encoder(tokenizer: torch.nn.Module) -> None: + """Restrict the optimization to the audited HF encoder's valid-prefix contract. + + Causal attention and pointwise transformer operations cannot read future + padded rows. Its patch-downsample floors lengths, so every valid output + group consists entirely of valid input rows. The quantizer is pointwise + and the caller crops codes by those lengths. Padded *outputs* may differ. + This argument does not apply to the streaming decoder or to upsampling. + """ + transformers = 0 + for module in tokenizer.encoder: + name = type(module).__name__ + if name == "MossAudioTokenizerPatchedPretransform" and module.is_downsample: + continue + if name == "MossAudioTokenizerProjectedTransformer": + layers = module.transformer.layers + if layers and all(layer.self_attn.causal for layer in layers): + transformers += 1 + continue + raise ValueError("Unmasked reference queries require the causal MOSS patch-downsample encoder") + if not transformers: + raise ValueError("Unmasked reference queries require at least one causal MOSS transformer") + if type(tokenizer.quantizer).__name__ not in { + "MossAudioTokenizerResidualVQ", + "MossAudioTokenizerResidualLFQ", + }: + raise ValueError("Unmasked reference queries require the pointwise MOSS quantizer") + for module in tokenizer.quantizer.modules(): + if isinstance(module, torch.nn.Conv1d) and module.kernel_size != (1,): + raise ValueError("Unmasked reference queries do not support temporal quantizer convolutions") + + +def install_windowed_attention( + tokenizer: torch.nn.Module, + *, + fa_version: int = 2, + cache_rope: bool = False, + skip_padded_query_mask: bool = False, +) -> int: + """Route the encoder's non-streaming attention through a local-window flash kernel.""" + if skip_padded_query_mask: + _validate_unmasked_reference_encoder(tokenizer) + logger.info("MOSS reference encoder: skipping padded query masks (valid-prefix outputs only)") + count = 0 + rope_cache = {} if cache_rope else None + if cache_rope: + tokenizer._moss_reference_rope_cache = rope_cache + tokenizer._moss_reference_rope_prepared_samples = 0 + for module in tokenizer.encoder.modules(): + sdpa = getattr(module, "_forward_non_streaming_sdpa", None) + if sdpa is None or not hasattr(module, "_project_qkv") or not getattr(module, "causal", False): + continue + + def attention(self, x, input_lengths, _sdpa=sdpa): + return _windowed_attention( + self, + x, + input_lengths, + sdpa=_sdpa, + fa_version=fa_version, + rope_cache=rope_cache, + skip_padded_query_mask=skip_padded_query_mask, + ) + + module._forward_non_streaming_sdpa = types.MethodType(attention, module) + count += 1 + return count + + +@dataclass +class _Graph: + graph: torch.cuda.CUDAGraph + audio: torch.Tensor # (B, C, T) float32 input + lengths: torch.Tensor # (B,) int64 input, in samples per channel + codes: torch.Tensor # (n_vq, B, frames) output + code_lengths: torch.Tensor # (B,) output + + +class MossReferenceEncoderGraphs: + """Encode prepared reference waveforms with captured tokenizer graphs.""" + + def __init__( + self, + tokenizer: torch.nn.Module, + *, + n_vq: int, + batch_sizes: tuple[int, ...] = DEFAULT_BATCH_SIZES, + bucket_seconds: tuple[float, ...] = DEFAULT_BUCKET_SECONDS, + max_batch_seconds: float = DEFAULT_MAX_BATCH_SECONDS, + compile_core: bool = False, + batched_transfer: bool = False, + singleton_bucket_seconds: tuple[float, ...] | None = None, + ) -> None: + self._tokenizer = tokenizer + self._compile_core = compile_core + self._batched_transfer = batched_transfer + self._n_vq = int(n_vq) + self._device = next(tokenizer.parameters()).device + self._channels = int(tokenizer.number_channels) + frame = int(tokenizer.downsample_rate) + sampling_rate = int(tokenizer.sampling_rate) + # Whole frames, so the tokenizer's own frame padding is a no-op. + self._buckets = sorted({int(math.ceil(s * sampling_rate / frame)) * frame for s in bucket_seconds}) + self._singleton_buckets = sorted( + set(self._buckets) + | {int(math.ceil(s * sampling_rate / frame)) * frame for s in (singleton_bucket_seconds or ())} + ) + self._batch_sizes = sorted({int(b) for b in batch_sizes if int(b) > 0}) + self._max_batch_samples = float(max_batch_seconds) * sampling_rate + self._graphs: dict[tuple[int, int], _Graph] = {} + self._pool = None + self._transfer_buffers = None + + @property + def max_batch(self) -> int: + return self._batch_sizes[-1] + + @property + def captured(self) -> list[tuple[int, int]]: + return sorted(self._graphs) + + def _core(self, audio: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Mirrors the tokenizer's _encode_frame without its host-side crop. + tokenizer = self._tokenizer + values, value_lengths = tokenizer._flatten_channels_for_codec(audio, lengths) + with tokenizer._codec_inference_autocast(): + hidden, hidden_lengths = values, value_lengths + for module in tokenizer.encoder: + hidden, hidden_lengths = module(hidden, hidden_lengths) + _, codes, code_lengths = tokenizer.quantizer(hidden.float(), hidden_lengths, self._n_vq) + return codes, code_lengths + + @torch.no_grad() + def capture(self) -> None: + if self._device.type != "cuda" or self._graphs: + return + self._pool = torch.cuda.graph_pool_handle() + stream = torch.cuda.Stream(device=self._device) + prepared = getattr(self._tokenizer, "_moss_reference_rope_prepared_samples", None) + max_samples = max(self._singleton_buckets) + if prepared is not None and prepared < max_samples: + # Prime immutable FP32 factors at the largest length before any + # capture or compilation. Smaller lengths use a prefix of the + # same table; the dictionary stays fixed throughout capture and + # workers share tables without introducing dynamic-shape guards + # for a separate dictionary key at each sequence length. + audio = torch.zeros(1, self._channels, max_samples, device=self._device) + lengths = torch.full((1,), max_samples, dtype=torch.long, device=self._device) + self._core(audio, lengths) + self._tokenizer._moss_reference_rope_prepared_samples = max_samples + core = self._core + if self._compile_core: + # Fuses the encoder's many small elementwise kernels (casts, RoPE, + # residuals). Only capture runs the compiled code; replays do not. + core = torch.compile(self._core, dynamic=True) + # Each encoder stage is its own module shape: allow a compile per stage. + with torch._dynamo.config.patch(recompile_limit=64): + self._capture_all(core, stream) + if self._batched_transfer and self._graphs: + # One reusable slab per worker, rather than one per graph bucket. + entries = list(self._graphs.values()) + self._transfer_buffers = ( + torch.empty(max(e.audio.numel() for e in entries), dtype=entries[0].audio.dtype, pin_memory=True), + torch.empty(max(e.lengths.numel() for e in entries), dtype=entries[0].lengths.dtype, pin_memory=True), + torch.empty(max(e.codes.numel() for e in entries), dtype=entries[0].codes.dtype, pin_memory=True), + torch.empty( + max(e.code_lengths.numel() for e in entries), dtype=entries[0].code_lengths.dtype, pin_memory=True + ), + torch.cuda.Event(), + ) + logger.info( + "MOSS reference encoder CUDA graphs: batch=%s buckets(s)=%s compiled=%s captured=%d", + self._batch_sizes, + [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._buckets], + self._compile_core, + len(self._graphs), + ) + if self._singleton_buckets != self._buckets: + logger.info( + "MOSS reference encoder singleton buckets(s): %s", + [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._singleton_buckets], + ) + + def _capture_all(self, core, stream: torch.cuda.Stream) -> None: + for batch in reversed(self._batch_sizes): + buckets = self._singleton_buckets if batch == 1 else self._buckets + for samples in reversed(buckets): + if batch > 1 and batch * samples > self._max_batch_samples: + continue + audio = torch.zeros(batch, self._channels, samples, device=self._device) + lengths = torch.full((batch,), samples, dtype=torch.long, device=self._device) + try: + stream.wait_stream(torch.cuda.current_stream(self._device)) + with torch.cuda.stream(stream): + # Lazy library setup must happen before capture; one run does it. + core(audio, lengths) + torch.cuda.current_stream(self._device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool, stream=stream): + codes, code_lengths = core(audio, lengths) + except Exception: + logger.warning( + "MOSS reference encoder graph (B=%d, samples=%d) not captured; using eager", + batch, + samples, + exc_info=True, + ) + continue + self._graphs[(batch, samples)] = _Graph(graph, audio, lengths, codes, code_lengths) + + @torch.no_grad() + def encode(self, wavs: list[torch.Tensor]) -> list[torch.Tensor | None]: + """Codes ``(frames, n_vq)`` int64 on CPU per clip; None where no graph fits. + + ``wavs`` are prepared the way the processor prepares them: ``(C, T)`` + float32 after channel handling, resampling and loudness normalization. + Clips are grouped longest first, so each replay pads to the bucket of + similar-length clips. + """ + out: list[torch.Tensor | None] = [None] * len(wavs) + order = sorted(range(len(wavs)), key=lambda i: -int(wavs[i].shape[-1])) + pos = 0 + while pos < len(order): + longest = int(wavs[order[pos]].shape[-1]) + # Fine singleton buckets reduce padding without fragmenting an + # already assembled batch into separate graph replays. + buckets = self._singleton_buckets if len(order) - pos == 1 else self._buckets + samples = next((s for s in buckets if s >= longest and (1, s) in self._graphs), None) + if samples is None: + pos += 1 # longer than every bucket: left to the caller + continue + widest = max(b for b in self._batch_sizes if (b, samples) in self._graphs) + rows = order[pos : pos + widest] + batch = next(b for b in self._batch_sizes if b >= len(rows) and (b, samples) in self._graphs) + for row, codes in zip(rows, self._replay(self._graphs[(batch, samples)], [wavs[i] for i in rows])): + out[row] = codes + pos += len(rows) + return out + + def _replay(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]: + if self._transfer_buffers is not None and all(w.device.type == "cpu" for w in wavs): + return self._replay_batched_transfer(entry, wavs) + count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1] + # Unused rows encode full-length silence: a fully masked row would be + # all-NaN inside attention, and rows never mix. + entry.audio.zero_() + lengths = [int(w.shape[-1]) for w in wavs] + [samples] * (batch - count) + entry.lengths.copy_(torch.tensor(lengths, dtype=torch.long)) + for row, wav in enumerate(wavs): + entry.audio[row, :, : wav.shape[-1]].copy_(wav) + entry.graph.replay() + codes = entry.codes[:, :count].cpu() + code_lengths = entry.code_lengths[:count].tolist() + return [codes[:, row, : int(code_lengths[row])].transpose(0, 1).contiguous().long() for row in range(count)] + + def _replay_batched_transfer(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]: + """Stage one batch on pinned host memory, then wait once for its outputs.""" + audio_slab, lengths_slab, codes_slab, code_lengths_slab, copied = self._transfer_buffers + count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1] + audio = audio_slab[: entry.audio.numel()].view_as(entry.audio) + lengths = lengths_slab[:batch] + audio.zero_() + lengths.fill_(samples) + for row, wav in enumerate(wavs): + audio[row, :, : wav.shape[-1]].copy_(wav) + lengths[row] = wav.shape[-1] + entry.audio.copy_(audio, non_blocking=True) + entry.lengths.copy_(lengths, non_blocking=True) + entry.graph.replay() + result = entry.codes[:, :count] + codes = codes_slab[: result.numel()].view_as(result) + code_lengths = code_lengths_slab[:count] + codes.copy_(result, non_blocking=True) + code_lengths.copy_(entry.code_lengths[:count], non_blocking=True) + copied.record(torch.cuda.current_stream(self._device)) + copied.synchronize() + # Returned codes must outlive this worker's next replay, including + # one-frame clips whose transpose is already contiguous. + return [ + codes[:, row, : int(code_lengths[row])].transpose(0, 1).to(dtype=torch.long, copy=True).contiguous() + for row in range(count) + ] diff --git a/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py new file mode 100644 index 000000000..fec5c2860 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One reference-audio encoder shared by the API processes of a server. + +Each API process otherwise loads its own copy of the audio tokenizer on the +GPU and encodes only the clips it receives. When the processes share a +directory (``VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR``), the first one to take the lock in it +hosts the encoder and serves the others over a Unix socket in that directory. +Requests from all processes are merged into batches, and only the host keeps +the tokenizer and its CUDA graphs on the GPU. + +Messages carry raw arrays and a JSON header; nothing is unpickled. A client +that cannot reach the host encodes locally instead. +""" + +from __future__ import annotations + +import collections +import fcntl +import json +import os +import threading +import time +from collections.abc import Callable +from contextlib import nullcontext +from multiprocessing.connection import Client, Connection, Listener + +import numpy as np +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_LOCK_NAME = "ref-encoder.lock" +_SOCKET_NAME = "ref-encoder.sock" +# Clips one worker merges into an encode, across all API processes. Several +# workers encode concurrently on their own streams: one serial queue of large +# batches makes a burst of new references wait behind each other. A free +# worker takes what is queued without waiting for more: API processes already +# gather their clips before sending them. +_MAX_BATCH = 8 +# A client gives up on the host after this long and encodes locally. +_CONNECT_TIMEOUT_S = 120.0 +_REQUEST_TIMEOUT_S = 120.0 +# First-time compilation can outlast a normal encode request. Readiness has a +# separate deadline and must finish before the API starts accepting requests. +_STARTUP_TIMEOUT_S = 1200.0 + +EncodeFn = Callable[[list[torch.Tensor]], list] + +# Open for the life of the process: closing it would release the host lock. +_host_lock_handle = None + + +class SharedReferenceEncoderStartupError(RuntimeError): + """The shared encoder cannot safely accept requests yet.""" + + +def elect_host(shared_dir: str) -> bool: + """True in the one process that holds the encoder lock for ``shared_dir``. + + The lock is held for the life of the process, so a restarted server + elects a new host even if the old socket file is still present. + """ + global _host_lock_handle + path = os.path.join(shared_dir, _LOCK_NAME) + handle = open(path, "a+") # noqa: SIM115 - kept open to hold the lock + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + handle.close() + return False + _host_lock_handle = handle + return True + + +def _send(conn: Connection, header: dict, arrays: list[np.ndarray]) -> None: + conn.send_bytes(json.dumps(header).encode()) + for array in arrays: + conn.send_bytes(np.ascontiguousarray(array).tobytes()) + + +def _recv(conn: Connection) -> tuple[dict, list[bytes]]: + header = json.loads(conn.recv_bytes().decode()) + return header, [conn.recv_bytes() for _ in range(int(header["count"]))] + + +class _Job: + __slots__ = ("wav", "done", "result", "queued_at") + + def __init__(self, wav: torch.Tensor) -> None: + self.wav = wav + self.done = threading.Event() + self.result: torch.Tensor | BaseException | None = None + self.queued_at = time.monotonic() + + +class SharedReferenceEncoderHost: + """Serves reference encodes for every API process sharing ``shared_dir``. + + ``workers`` pairs each worker's encode function with an optional warmup + (for example, capturing that worker's CUDA graphs). + """ + + def __init__( + self, + shared_dir: str, + workers: list[tuple[EncodeFn, Callable[[], object] | None]], + *, + stream_priority: int = 0, + batch_window_ms: float = 0.0, + ) -> None: + self._stream_priority = stream_priority + self._batch_window_s = max(0.0, float(batch_window_ms)) / 1000 + self._pending: collections.deque[_Job] = collections.deque() + # Includes jobs already taken by a worker, so shutdown can release + # every caller without waiting for an in-flight GPU operation. + self._outstanding: set[_Job] = set() + self._ready = threading.Condition() + self._path = os.path.join(shared_dir, _SOCKET_NAME) + if os.path.exists(self._path): + os.unlink(self._path) # left by a previous host; we hold the lock + self._listener = Listener(self._path, family="AF_UNIX") + os.chmod(self._path, 0o600) + self._closed = False + self._started = threading.Event() + self._startup_error: Exception | None = None + threading.Thread( + target=self._start_workers, args=(workers,), name="moss-ref-encoder-start", daemon=True + ).start() + threading.Thread(target=self._accept_loop, name="moss-ref-encoder-accept", daemon=True).start() + logger.info("MOSS shared reference encoder serving at %s with %d workers", self._path, len(workers)) + + def encode(self, wavs: list[torch.Tensor]) -> list: + """Encode on behalf of the host process itself.""" + self.wait_until_ready() + jobs = [_Job(wav) for wav in wavs] + with self._ready: + if self._closed: + raise RuntimeError("shared reference encoder is closed") + self._outstanding.update(jobs) + self._pending.extend(jobs) + self._ready.notify(len(jobs)) + for job in jobs: + job.done.wait() + return [job.result for job in jobs] + + def close(self) -> None: + with self._ready: + self._closed = True + self._started.set() + for job in self._outstanding: + job.result = RuntimeError("shared reference encoder closed before encoding completed") + job.done.set() + self._outstanding.clear() + self._pending.clear() + self._ready.notify_all() + try: + self._listener.close() + except OSError: + pass + + def wait_until_ready(self) -> None: + if not self._started.wait(_STARTUP_TIMEOUT_S): + raise SharedReferenceEncoderStartupError("shared reference encoder warmup timed out") + if self._closed or self._startup_error is not None: + raise SharedReferenceEncoderStartupError("shared reference encoder startup failed") from self._startup_error + + def _accept_loop(self) -> None: + while not self._closed: + try: + conn = self._listener.accept() + except OSError: + return + threading.Thread(target=self._serve, args=(conn,), name="moss-ref-encoder-conn", daemon=True).start() + + def _serve(self, conn: Connection) -> None: + with conn: + while True: + try: + header, payloads = _recv(conn) + except (EOFError, OSError): + return + if header.get("op") == "ready": + try: + self.wait_until_ready() + error = None + except SharedReferenceEncoderStartupError as exc: + error = str(exc) + try: + _send(conn, {"count": 0, "ready": error is None, "error": error}, []) + except OSError: + return + continue + wavs = [ + torch.from_numpy(np.frombuffer(data, dtype=np.float32).reshape(shape).copy()) + for data, shape in zip(payloads, header["shapes"]) + ] + results = self.encode(wavs) + errors = [None if isinstance(r, torch.Tensor) else repr(r) for r in results] + codes = [r.to(torch.int32).numpy() for r in results if isinstance(r, torch.Tensor)] + reply = {"count": len(codes), "shapes": [list(c.shape) for c in codes], "errors": errors} + try: + _send(conn, reply, codes) + except OSError: + return + + def _start_workers(self, workers: list[tuple[EncodeFn, Callable[[], object] | None]]) -> None: + # Every warmup (graph capture) finishes before any worker encodes: CUDA + # work in another thread would invalidate a capture in progress. + try: + for _, warmup in workers: + if warmup is not None: + warmup() + if self._closed: + return + for index, (encode_fn, _) in enumerate(workers): + threading.Thread( + target=self._encode_loop, args=(encode_fn,), name=f"moss-ref-encoder-{index}", daemon=True + ).start() + logger.info("MOSS shared reference encoder ready: %d workers", len(workers)) + except Exception as exc: # noqa: BLE001 - wake readiness waiters with the failure + self._startup_error = exc + logger.exception("MOSS shared reference encoder startup failed") + finally: + self._started.set() + + def _encode_loop(self, encode_fn: EncodeFn) -> None: + stream = torch.cuda.Stream(priority=self._stream_priority) if torch.cuda.is_available() else None + with torch.cuda.stream(stream) if stream is not None else nullcontext(): + self._serve_jobs(encode_fn) + + def _serve_jobs(self, encode_fn: EncodeFn) -> None: + while True: + with self._ready: + while True: + if self._closed: + return + if not self._pending: + self._ready.wait() + continue + remaining = self._pending[0].queued_at + self._batch_window_s - time.monotonic() + if remaining > 0 and len(self._pending) < _MAX_BATCH: + # Keep clips in the common queue while collecting; + # workers must not each reserve a singleton first. + self._ready.wait(remaining) + continue + break + jobs = [self._pending.popleft() for _ in range(min(_MAX_BATCH, len(self._pending)))] + try: + results = encode_fn([job.wav for job in jobs]) + if len(results) != len(jobs): + raise RuntimeError( + f"shared reference encoder returned {len(results)} results for {len(jobs)} clips" + ) + except Exception as error: # noqa: BLE001 - every waiter gets the failure + results = [error] * len(jobs) + with self._ready: + for job, result in zip(jobs, results, strict=True): + # close() may already have completed this caller with an + # error. A late GPU result must not replace that outcome. + if not job.done.is_set(): + job.result = result + job.done.set() + self._outstanding.discard(job) + + +class SharedReferenceEncoderClient: + """Sends reference clips to the host process of ``shared_dir``.""" + + def __init__(self, shared_dir: str) -> None: + self._path = os.path.join(shared_dir, _SOCKET_NAME) + # One connection per request in flight; idle ones are reused. + self._idle: list[Connection] = [] + self._lock = threading.Lock() + self._connected = False + + def _connect(self) -> Connection: + # The host may still be starting until the first connection succeeds; + # later reconnects are quick. + deadline = time.monotonic() + (1.0 if self._connected else _CONNECT_TIMEOUT_S) + while True: + try: + conn = Client(self._path, family="AF_UNIX") + self._connected = True + break + except (FileNotFoundError, ConnectionRefusedError): + if time.monotonic() >= deadline: + raise + time.sleep(0.2) + try: + _send(conn, {"op": "ready", "count": 0}, []) + if not conn.poll(_STARTUP_TIMEOUT_S): + raise SharedReferenceEncoderStartupError("shared reference encoder readiness timed out") + header, _ = _recv(conn) + if not header.get("ready"): + raise SharedReferenceEncoderStartupError(header.get("error", "shared reference encoder not ready")) + return conn + except BaseException: + conn.close() + raise + + def wait_until_ready(self) -> None: + """Startup handshake, without encoding or populating the reference cache.""" + try: + conn = self._connect() + except Exception as exc: + raise SharedReferenceEncoderStartupError("shared reference encoder is not ready") from exc + with self._lock: + self._idle.append(conn) + + def encode(self, wavs: list[torch.Tensor]) -> list: + """Codes (or an exception) per clip; raises when the host is unreachable.""" + with self._lock: + conn = self._idle.pop() if self._idle else None + try: + if conn is None: + conn = self._connect() + arrays = [w.detach().to(torch.float32).cpu().numpy() for w in wavs] + _send(conn, {"count": len(arrays), "shapes": [list(a.shape) for a in arrays]}, arrays) + if not conn.poll(_REQUEST_TIMEOUT_S): + raise TimeoutError(f"no reply from the shared reference encoder in {_REQUEST_TIMEOUT_S}s") + header, payloads = _recv(conn) + except BaseException: + if conn is not None: + conn.close() + raise + with self._lock: + self._idle.append(conn) + codes = iter( + torch.from_numpy(np.frombuffer(data, dtype=np.int32).reshape(shape).copy()).long() + for data, shape in zip(payloads, header["shapes"]) + ) + return [next(codes) if error is None else RuntimeError(error) for error in header["errors"]] diff --git a/vllm_omni/model_executor/models/moss_tts/slot_attention.py b/vllm_omni/model_executor/models/moss_tts/slot_attention.py new file mode 100644 index 000000000..f7f6d2783 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/slot_attention.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Direct-slot ring attention for MOSS streaming codec. + +Inputs q/k/v are (execution_batch, heads, padded_T, head_dim). cache is +contiguous (2, state_capacity, heads, ring_capacity, head_dim). end_offset +contains absolute *pre-call* positions per state slot. slot_ids and valid_lengths +contain one entry per execution row. valid_lengths=0 is a scratch/padding row. + +This op writes only valid tokens, attends to the resulting ring with the same +position semantics as RingKVCache.complete, then advances only active offsets. +Slots must be unique and in bounds; 0 <= valid_lengths <= padded_T. Chunks +longer than the ring retain only their final ring_capacity tokens. +These device-value preconditions are caller-owned to avoid host synchronization. +The op deliberately preserves chunk-complete ring overwrite semantics; it does +not recover entries overwritten by later tokens of the same input chunk. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _write( + k_ptr, + v_ptr, + cache_ptr, + end_ptr, + slots_ptr, + lengths_ptr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + ks3: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + vs3: tl.constexpr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + block_size: tl.constexpr, +): + row = tl.program_id(0) + slot = tl.load(slots_ptr + row * slot_stride) + length = tl.load(lengths_ptr + row * length_stride) + start = tl.load(end_ptr + slot * end_stride) + x = tl.program_id(1) * block_size + tl.arange(0, block_size) + dim = x % head_dim + token = (x // head_dim) % num_frames + head = x // (num_frames * head_dim) + # For chunks longer than the ring, write only the final C tokens. This + # avoids duplicate physical addresses (and write races) in one launch. + skip = tl.maximum(length - ring_capacity, 0) + source_token = token + skip + mask = (head < num_heads) & (source_token < length) + pos = (start + source_token) % ring_capacity + dst = ((slot * num_heads + head) * ring_capacity + pos) * head_dim + dim + kval = tl.load(k_ptr + row * ks0 + head * ks1 + source_token * ks2 + dim * ks3, mask, 0) + vval = tl.load(v_ptr + row * vs0 + head * vs1 + source_token * vs2 + dim * vs3, mask, 0) + tl.store(cache_ptr + dst, kval, mask) + tl.store(cache_ptr + state_capacity * num_heads * ring_capacity * head_dim + dst, vval, mask) + + +@triton.jit +def _attend( + q_ptr, + cache_ptr, + end_ptr, + slots_ptr, + lengths_ptr, + out_ptr, + qs0: tl.constexpr, + qs1: tl.constexpr, + qs2: tl.constexpr, + qs3: tl.constexpr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + context_size: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + skip_empty_tiles: tl.constexpr = False, +): + 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 + slot = tl.load(slots_ptr + batch * slot_stride) + length = tl.load(lengths_ptr + batch * length_stride) + start = tl.load(end_ptr + slot * end_stride) + end = start + length + dims = tl.arange(0, head_dim) + q = tl.load( + q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 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) + # Padding rows own scratch slots and must emit zeros without reading KV. + # Keep valid-row accumulation and the ring overwrite semantics unchanged. + if skip_empty_tiles: + block_start = tl.program_id(0) * block_m + live = (length > 0) & (block_start < length) + live = live & (block_start + block_m > tl.maximum(length - ring_capacity, 0)) + else: + live = length > 0 + if live: + # Physical-ring order matches the existing masked_attention accumulation. + last = end - 1 + last_index = last % ring_capacity + if skip_empty_tiles: + # Before the first wrap, only physical keys [0, end) exist. + # Omitted blocks have zero probability for every query; keep the + # physical order and arithmetic of all remaining blocks unchanged. + key_blocks = tl.cdiv(tl.minimum(end, ring_capacity), block_n) + else: + key_blocks = tl.cdiv(ring_capacity, block_n) + for first in range(key_blocks): + cols = first * block_n + tl.arange(0, block_n) + delta = cols - last_index + positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity) + distance = start + rows[:, None] - positions[None, :] + allowed = ( + (rows[:, None] < length) + & (cols[None, :] < ring_capacity) + & (positions[None, :] >= 0) + & (positions[None, :] < end) + & (distance >= 0) + ) + if context_size > 0: + allowed = allowed & (distance < context_size) + base = (slot * num_heads + head) * ring_capacity * head_dim + k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0) + score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5) + score = tl.where(allowed, score, -float("inf")) + next_max = tl.maximum(maximum, tl.max(score, 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, 1) + v = tl.load( + cache_ptr + + state_capacity * num_heads * ring_capacity * head_dim + + base + + cols[:, None] * head_dim + + dims[None, :], + cols[:, None] < ring_capacity, + 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) * num_frames + rows[:, None]) * head_dim + dims[None, :], + out, + rows[:, None] < num_frames, + ) + + +@triton.jit +def _advance( + end_ptr, + slots_ptr, + lengths_ptr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + batch_size: tl.constexpr, + block_size: tl.constexpr, +): + row = tl.arange(0, block_size) + slot = tl.load(slots_ptr + row * slot_stride, row < batch_size, 0) + length = tl.load(lengths_ptr + row * length_stride, row < batch_size, 0) + active = (row < batch_size) & (length > 0) + start = tl.load(end_ptr + slot * end_stride, active, 0) + tl.store(end_ptr + slot * end_stride, start + length, active) + + +@triton.jit +def _write_attend( + q_ptr, + k_ptr, + v_ptr, + cache_ptr, + starts_ptr, + slots_ptr, + lengths_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, + start_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + context_size: tl.constexpr, + write_frames: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_w: 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 + slot = tl.load(slots_ptr + batch * slot_stride) + length = tl.load(lengths_ptr + batch * length_stride) + start = tl.load(starts_ptr + batch * start_stride) + end = start + length + dims = tl.arange(0, head_dim) + base = (slot * num_heads + head) * ring_capacity * head_dim + v_base = state_capacity * num_heads * ring_capacity * head_dim + base + # Every program of this (row, head) writes the whole chunk for its head, + # so it only reads ring entries it wrote itself (duplicate writes across + # query blocks store identical values). Chunks longer than the ring keep + # only their final ring_capacity tokens, as in _write. + skip = tl.maximum(length - ring_capacity, 0) + for first in range(tl.cdiv(write_frames, block_w)): + token = first * block_w + tl.arange(0, block_w) + source = token + skip + mask = (source < length)[:, None] + pos = (start + source) % ring_capacity + kval = tl.load(k_ptr + batch * ks0 + head * ks1 + source[:, None] * ks2 + dims[None, :] * ks3, mask, 0) + vval = tl.load(v_ptr + batch * vs0 + head * vs1 + source[:, None] * vs2 + dims[None, :] * vs3, mask, 0) + dst = pos[:, None] * head_dim + dims[None, :] + tl.store(cache_ptr + base + dst, kval, mask) + tl.store(cache_ptr + v_base + dst, vval, mask) + tl.debug_barrier() + q = tl.load( + q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 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) + if length > 0: + last = end - 1 + last_index = last % ring_capacity + for first in range(tl.cdiv(ring_capacity, block_n)): + cols = first * block_n + tl.arange(0, block_n) + delta = cols - last_index + positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity) + distance = start + rows[:, None] - positions[None, :] + allowed = ( + (rows[:, None] < length) + & (cols[None, :] < ring_capacity) + & (positions[None, :] >= 0) + & (positions[None, :] < end) + & (distance >= 0) + ) + if context_size > 0: + allowed = allowed & (distance < context_size) + k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0) + score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5) + score = tl.where(allowed, score, -float("inf")) + next_max = tl.maximum(maximum, tl.max(score, 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, 1) + v = tl.load( + cache_ptr + v_base + cols[:, None] * head_dim + dims[None, :], + cols[:, None] < ring_capacity, + 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) * num_frames + rows[:, None]) * head_dim + dims[None, :], + out, + rows[:, None] < num_frames, + ) + + +@torch.library.custom_op("vllm_omni::moss_slot_ring_attention", mutates_args=("cache", "end_offset")) +def slot_ring_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + end_offset: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, + skip_empty_tiles: bool = False, +) -> torch.Tensor: + """Fused address/mask path, with three ordered launches and no KV gather. + + q/k are already rotary-position encoded. The caller still advances its + separate MHA/RoPE offset state; this op advances only end_offset in place. + context=-1 means unbounded causal attention over available ring entries. + Outputs beyond each row's valid_lengths are exactly zero. + skip_empty_tiles skips wholly masked tiles for small (T <= 32) chunks. + Larger chunks keep the static loop, which benchmarks faster after wrap. + """ + b, h, t, d = q.shape + if k.shape != q.shape or v.shape != q.shape: + raise ValueError("q/k/v must have identical shapes") + if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share bf16 or fp16 dtype") + if d not in (32, 64, 128) or t < 1: + raise ValueError("requires head_dim 32/64/128 and positive padded_T") + if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d: + raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)") + s, c = cache.shape[1], cache.shape[3] + if not cache.is_contiguous() or cache.dtype != q.dtype: + raise ValueError("cache must be contiguous and share q dtype") + if end_offset.shape != (s,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,): + raise ValueError("invalid metadata shapes") + if any(x.dtype not in (torch.int32, torch.int64) for x in (end_offset, slot_ids, valid_lengths)): + raise ValueError("metadata must be int32/int64") + if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, end_offset, slot_ids, valid_lengths)): + raise ValueError("all tensors must share one CUDA device") + if context == 0 or context < -1: + raise ValueError("context must be -1 or positive") + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + if b == 0: + return out + strides = (end_offset.stride(0), slot_ids.stride(0), valid_lengths.stride(0)) + _write[(b, triton.cdiv(h * t * d, 256))]( + k, v, cache, end_offset, slot_ids, valid_lengths, *k.stride(), *v.stride(), *strides, s, h, t, c, d, 256 + ) + bm = 16 if t <= 32 else 64 + _attend[(triton.cdiv(t, bm), b * h)]( + q, + cache, + end_offset, + slot_ids, + valid_lengths, + out, + *q.stride(), + *strides, + s, + h, + t, + c, + d, + context, + bm, + 64, + skip_empty_tiles=skip_empty_tiles and t <= 32, + num_warps=8 if c <= 256 else 4, + num_stages=2, + ) + _advance[(1,)](end_offset, slot_ids, valid_lengths, *strides, b, triton.next_power_of_2(b)) + return out + + +@slot_ring_attention.register_fake +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + end_offset: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, + skip_empty_tiles: bool = False, +) -> torch.Tensor: + return torch.empty(q.shape, device=q.device, dtype=q.dtype) + + +def _check_inputs(q, k, v, cache, metadata, context): + b, h, t, d = q.shape + if k.shape != q.shape or v.shape != q.shape: + raise ValueError("q/k/v must have identical shapes") + if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share bf16 or fp16 dtype") + if d not in (32, 64, 128) or t < 1: + raise ValueError("requires head_dim 32/64/128 and positive padded_T") + if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d: + raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)") + if not cache.is_contiguous() or cache.dtype != q.dtype: + raise ValueError("cache must be contiguous and share q dtype") + if any(x.dtype not in (torch.int32, torch.int64) for x in metadata): + raise ValueError("metadata must be int32/int64") + if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, *metadata)): + raise ValueError("all tensors must share one CUDA device") + if context == 0 or context < -1: + raise ValueError("context must be -1 or positive") + + +@torch.library.custom_op("vllm_omni::moss_slot_ring_attention_rows", mutates_args=("cache",)) +def slot_ring_attention_rows( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + starts: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, +) -> torch.Tensor: + """``slot_ring_attention`` as one launch, with per-row start positions. + + ``starts`` holds each execution row's pre-call absolute position (the + caller's gathered offset). The ring is written and attended exactly as by + ``slot_ring_attention``; no offset is advanced here, so the caller owns + every offset update and can fuse them with its own. + """ + b, h, t, d = q.shape + _check_inputs(q, k, v, cache, (starts, slot_ids, valid_lengths), context) + if starts.shape != (b,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,): + raise ValueError("invalid metadata shapes") + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + if b == 0: + return out + s, c = cache.shape[1], cache.shape[3] + bm = 16 if t <= 32 else 64 + _write_attend[(triton.cdiv(t, bm), b * h)]( + q, + k, + v, + cache, + starts, + slot_ids, + valid_lengths, + out, + *q.stride(), + *k.stride(), + *v.stride(), + starts.stride(0), + slot_ids.stride(0), + valid_lengths.stride(0), + s, + h, + t, + c, + d, + context, + min(t, c), + bm, + 64, + 16, + num_warps=8 if c <= 256 else 4, + num_stages=2, + ) + return out + + +@slot_ring_attention_rows.register_fake +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + starts: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, +) -> torch.Tensor: + return torch.empty(q.shape, device=q.device, dtype=q.dtype) diff --git a/vllm_omni/model_executor/output_snapshot.py b/vllm_omni/model_executor/output_snapshot.py new file mode 100644 index 000000000..93ab0491e --- /dev/null +++ b/vllm_omni/model_executor/output_snapshot.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Owned, dtype-grouped snapshots for asynchronous runner outputs.""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from torch.utils._pytree import tree_flatten, tree_unflatten + + +@dataclass(frozen=True) +class RequestOutputSnapshot: + """Already partitioned, CPU-owned payloads in the current request order. + + Model finalizers may return this after D2H completes to avoid recursively + cloning and partitioning data whose request ownership is already known. + """ + + inter_stage: list[dict[str, Any] | None] + client: list[dict[str, Any] | None] | None = None + + +class PackedOutputSnapshot(dict): + """A normal payload mapping with an internal batched-copy plan. + + The runner's snapshot-slot event protects the device slabs until D2H is + complete. Host slabs are allocated per output, so downstream consumers can + retain their views without depending on reuse of the device ring slot. + """ + + def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None: + super().__init__(tree_unflatten(leaves, spec)) + self._leaves = leaves + self._spec = spec + self._groups = groups + self.producer_event: torch.cuda.Event | None = None + + def record_producer_event(self, stream: torch.cuda.Stream) -> None: + """Publish readiness when a model produces on a separate CUDA stream.""" + self.producer_event = torch.cuda.Event() + self.producer_event.record(stream) + + def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]: + leaves = list(self._leaves) + for indices, slab in self._groups: + cpu = copy_tensor(slab) + offset = 0 + for index in indices: + value = leaves[index] + size = value.numel() + leaves[index] = cpu.narrow(0, offset, size).view(value.shape) + offset += size + return tree_unflatten(leaves, self._spec) + + +def pack_output_snapshot( + payload: dict[str, Any], + slot: dict[tuple[Any, ...], torch.Tensor], + *, + max_buckets: int, + reuse_existing_storage: bool = False, +) -> PackedOutputSnapshot | None: + """Copy tensor leaves once per dtype/device, retaining their exact values. + + No CUDA synchronization is introduced here. The caller must select the + producer stream and wait for the previous consumer before reusing a slot. + """ + leaves, spec = tree_flatten(payload) + grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list) + for index, value in enumerate(leaves): + if isinstance(value, torch.Tensor): + # Preserve the existing per-tensor path for unusual layouts. + if value.layout != torch.strided or value.is_quantized: + return None + grouped[(value.dtype, value.device)].append(index) + groups = [] + for (dtype, device), indices in grouped.items(): + values = [leaves[index].detach().reshape(-1) for index in indices] + size = sum(value.numel() for value in values) + # A producer may have already written the leaves into one contiguous + # per-forward slab. Reuse that storage directly instead of copying + # it into a second packed slab. The strict offset/stride check keeps + # this optimization limited to views whose order matches the copy + # plan; arbitrary payloads retain the existing copy path. + if reuse_existing_storage and values and all(value.is_contiguous() for value in values): + storage = values[0].untyped_storage() + offset = values[0].storage_offset() + cursor = offset + shares_storage = True + for value in values: + if value.untyped_storage().data_ptr() != storage.data_ptr() or value.storage_offset() != cursor: + shares_storage = False + break + cursor += value.numel() + if shares_storage: + slab = values[0].as_strided((size,), (1,), storage_offset=offset) + offset = 0 + for index, value in zip(indices, values): + shape = leaves[index].shape + leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) + offset += value.numel() + groups.append((indices, slab)) + continue + key = ("packed-output", dtype, device, size) + slab = slot.get(key) + if slab is None: + slab = torch.empty(size, dtype=dtype, device=device) + if len(slot) < max_buckets: + slot[key] = slab + if len(values) == 1: + slab.copy_(values[0]) + else: + torch.cat(values, out=slab) + offset = 0 + for index, value in zip(indices, values): + shape = leaves[index].shape + leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) + offset += value.numel() + groups.append((indices, slab)) + return PackedOutputSnapshot(leaves, spec, groups) diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index e07f552b1..bdb7f8e92 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.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"). @@ -9,15 +12,13 @@ from collections import defaultdict from collections.abc import Mapping from typing import Any +import numpy as np import torch from vllm.inputs import TokensPrompt as OmniTokensPrompt from vllm.logger import init_logger from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct -from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( - parse_chunk_ramp, - ramp_chunk_size, -) +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_chunk_size logger = init_logger(__name__) @@ -237,7 +238,8 @@ def talker2codec_raw_async_chunk( Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally ``[1, n_vq]`` per decode step). This processor buffers those new rows until a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern - de-delay is applied on this path. + de-delay is applied on this path. An optional ``codec_chunk_ramp`` selects + successive chunk sizes and takes precedence over the initial chunk size. """ external_req_id = getattr(request, "external_req_id", None) req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request))) @@ -248,6 +250,8 @@ def talker2codec_raw_async_chunk( transfer_manager.request_payload = {} if not hasattr(transfer_manager, "put_req_chunk"): transfer_manager.put_req_chunk = defaultdict(int) + if not hasattr(transfer_manager, "ramp_chunk_count"): + transfer_manager.ramp_chunk_count = defaultdict(int) pending_frames = transfer_manager.code_prompt_token_ids[req_id] @@ -255,14 +259,16 @@ def talker2codec_raw_async_chunk( codes_dict = multimodal_output.get("codes", {}) or {} new_frames = codes_dict.get("audio") if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0: - frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous() - if frames_cpu.ndim == 1: - frames_cpu = frames_cpu.reshape(1, -1) - if frames_cpu.ndim != 2: - raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}") - valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1) - for frame in frames_cpu[valid_rows]: - pending_frames.append(frame.clone()) + # This runs per request on every decode step. Tiny torch ops cost + # tens of microseconds each here; NumPy on the host copy does not. + frames_np = new_frames.detach().cpu().numpy() + if frames_np.ndim == 1: + frames_np = frames_np.reshape(1, -1) + if frames_np.ndim != 2: + raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_np.shape)}") + valid_rows = (frames_np != _MOSS_AUDIO_PAD_CODE).any(axis=1) + owned = np.array(frames_np if valid_rows.all() else frames_np[valid_rows], dtype=np.int64) + pending_frames.extend(owned) # Raw/local streaming should mirror the non-streaming path: the codec # decodes only generated audio rows. Reference audio conditions the # talker, but feeding its codes into the codec streaming state adds a @@ -288,26 +294,27 @@ def talker2codec_raw_async_chunk( initial_chunk_frames = chunk_frames pending = len(pending_frames) - # codec_chunk_ramp (opt-in): static ladder of chunk sizes replacing the - # initial_codec_chunk_frames -> codec_chunk_frames cliff. The chunk index - # is maintained by the connector (ramp_chunk_count, incremented per - # emitted chunk); past the ladder, ramp_chunk_size falls back to steady. - if not hasattr(transfer_manager, "_moss_ramp_parsed"): - transfer_manager._moss_ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_frames) - _ramp = transfer_manager._moss_ramp_parsed - if _ramp is not None: - if not hasattr(transfer_manager, "ramp_chunk_count"): - transfer_manager.ramp_chunk_count = defaultdict(int) - _chunk_idx = int(transfer_manager.ramp_chunk_count.get(req_id, 0)) - threshold = ramp_chunk_size(_chunk_idx, _ramp, chunk_frames) - else: - emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 - threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + emitted_chunks = int(transfer_manager.put_req_chunk.get(req_id, 0)) + emitted_any = emitted_chunks > 0 + threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + # The connector configuration is fixed for this transfer manager's lifetime. + if not hasattr(transfer_manager, "_moss_chunk_ramp"): + transfer_manager._moss_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_frames) + ramp = transfer_manager._moss_chunk_ramp + if ramp is not None: + # The ladder is indexed by the connector's segment-local counter, which + # restarts at each segment boundary; put_req_chunk is request-global. + ramp_index = int(transfer_manager.ramp_chunk_count.get(req_id, 0)) + threshold = ramp_chunk_size(ramp_index, ramp, chunk_frames) if pending <= 0: if is_finished: transfer_manager.code_prompt_token_ids.pop(req_id, None) transfer_manager.request_payload.pop(req_id, None) return OmniPayloadStruct( + # Explicitly replace the previous chunk (or the prewarmed + # placeholder) with no codes. Metadata alone is not an + # empty audio snapshot in the native connector contract. + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), meta=MetaStruct( req_id=[req_id], left_context_size=0, @@ -325,13 +332,10 @@ def talker2codec_raw_async_chunk( emit_frames = pending if is_finished else threshold chunk_rows = pending_frames[:emit_frames] del pending_frames[:emit_frames] - chunk_codes = torch.stack( - [row.to(torch.long).cpu() for row in chunk_rows], - dim=0, - ).contiguous() + chunk_np = np.stack([np.asarray(row, dtype=np.int64) for row in chunk_rows]) finished = bool(is_finished and len(pending_frames) == 0) - codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long) + codec_flat = torch.from_numpy(np.ascontiguousarray(chunk_np.T).reshape(-1)) if finished: transfer_manager.code_prompt_token_ids.pop(req_id, None) @@ -342,7 +346,7 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_chunk_frames=int(chunk_codes.shape[0]), + codec_chunk_frames=int(chunk_np.shape[0]), codec_left_context_frames=0, code_flat_numel=int(codec_flat.numel()), stream_finished=torch.tensor(finished, dtype=torch.bool), diff --git a/vllm_omni/worker_v2/first_audio_sender.py b/vllm_omni/worker_v2/first_audio_sender.py index 7e9246b5b..e1d09f367 100644 --- a/vllm_omni/worker_v2/first_audio_sender.py +++ b/vllm_omni/worker_v2/first_audio_sender.py @@ -111,9 +111,10 @@ class _PreparedDelivery: class _EngineOutputSink: - def __init__(self, output_queue: Any, scheduler: Any) -> None: + def __init__(self, output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True) -> None: self.output_queue = output_queue self.scheduler = scheduler + self.upstream_first_audio = upstream_first_audio def prepare(self, request_ids: list[str]) -> _PreparedDelivery: routes = {} @@ -130,15 +131,14 @@ class _EngineOutputSink: for request_id, pcm in zip(request_ids, pcm_rows, strict=True): if request_id not in routes: continue + payload = {"model_outputs": pcm, "sr": sample_rate} + if self.upstream_first_audio: + payload[FIRST_AUDIO_KEY] = torch.tensor(True) by_client.setdefault(routes[request_id], []).append( OmniEngineCoreOutput( request_id=request_id, new_token_ids=[], - multimodal_output={ - "model_outputs": pcm, - "sr": sample_rate, - FIRST_AUDIO_KEY: torch.tensor(True), - }, + multimodal_output=payload, ) ) for client_index, outputs in by_client.items(): @@ -163,5 +163,8 @@ class _EngineOutputSink: routes.pop(request_id) -def engine_output_queue_sink(output_queue: Any, scheduler: Any) -> _EngineOutputSink: - return _EngineOutputSink(output_queue, scheduler) +def engine_output_queue_sink( + output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True +) -> _EngineOutputSink: + """Route prepared audio; final-stage PCM must not carry the upstream marker.""" + return _EngineOutputSink(output_queue, scheduler, upstream_first_audio=upstream_first_audio) diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py index be7e19dbb..7475d3d7d 100644 --- a/vllm_omni/worker_v2/model_states/__init__.py +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -38,6 +38,12 @@ def init_omni_model_state( getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs") ) if uses_omni_lifecycle: + factory = getattr(type(model), "create_mrv2_model_state", None) + if callable(factory): + state = factory(model, vllm_config, encoder_cache, device) + if not isinstance(state, ModelState): + raise TypeError("create_mrv2_model_state must return a ModelState") + return state from vllm_omni.worker_v2.model_states.omni_model_state import ( OmniModelState, ) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 73434b8a1..c9ef50236 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -138,6 +138,23 @@ class OmniARModelRunner(OmniGPUModelRunner): # sample_tokens: OmniOutput handling + pooler_output + async D2H # ------------------------------------------------------------------ + def sample(self, hidden_states, input_batch, grammar_output): + # An explicitly declared model-state hook may return already determined + # tokens. Unsupported sampling features retain the upstream path. + sample_determined = getattr(type(self.model_state), "sample_determined_tokens", None) + if ( + sample_determined is not None + and grammar_output is None + and self.batch_sharder is None + and self.pp_handler is None + and input_batch.num_draft_tokens == 0 + and input_batch.num_reqs > 0 + ): + output = sample_determined(self.model_state, input_batch, self.sampler) + if output is not None: + return output, output.num_sampled, output.num_rejected + return super().sample(hidden_states, input_batch, grammar_output) + @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -311,6 +328,11 @@ class OmniARModelRunner(OmniGPUModelRunner): # Freshly allocated each step and filled after sampling; nothing # overwrites them before the host copy. return outputs + # A producer-side PackedOutputSnapshot already owns its device slabs + # for this forward and carries the event that makes them visible to + # the output copy stream. Repacking it would add a second GPU copy. + if isinstance(outputs, PackedOutputSnapshot): + return outputs slot_index = self._async_mm_snapshot_cursor if self._async_mm_snapshot_pending[slot_index]: raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().") @@ -567,6 +589,8 @@ def _async_copy_mm( if not mm_outputs: return {} if isinstance(mm_outputs, PackedOutputSnapshot): + if mm_outputs.producer_event is not None: + mm_outputs.producer_event.wait(copy_stream) return mm_outputs.copy_to_cpu( lambda tensor: _async_copy_tensor(tensor, copy_stream=copy_stream, pin_memory=pin_memory) ) diff --git a/vllm_omni/worker_v2/output_snapshot.py b/vllm_omni/worker_v2/output_snapshot.py index a8ae384db..e496a986a 100644 --- a/vllm_omni/worker_v2/output_snapshot.py +++ b/vllm_omni/worker_v2/output_snapshot.py @@ -1,91 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Owned, dtype-grouped snapshots for asynchronous runner outputs.""" +"""Compatibility imports for the shared model/runner output snapshot contract.""" -from __future__ import annotations +from vllm_omni.model_executor.output_snapshot import ( + PackedOutputSnapshot, + RequestOutputSnapshot, + pack_output_snapshot, +) -from collections import defaultdict -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -import torch -from torch.utils._pytree import tree_flatten, tree_unflatten - - -@dataclass(frozen=True) -class RequestOutputSnapshot: - """Already partitioned, CPU-owned payloads in the current request order. - - Model finalizers may return this after D2H completes to avoid recursively - cloning and partitioning data whose request ownership is already known. - """ - - inter_stage: list[dict[str, Any] | None] - client: list[dict[str, Any] | None] | None = None - - -class PackedOutputSnapshot(dict): - """A normal payload mapping with an internal batched-copy plan. - - The runner's snapshot-slot event protects the device slabs until D2H is - complete. Host slabs are allocated per output, so downstream consumers can - retain their views without depending on reuse of the device ring slot. - """ - - def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None: - super().__init__(tree_unflatten(leaves, spec)) - self._leaves = leaves - self._spec = spec - self._groups = groups - - def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]: - leaves = list(self._leaves) - for indices, slab in self._groups: - cpu = copy_tensor(slab) - offset = 0 - for index in indices: - value = leaves[index] - size = value.numel() - leaves[index] = cpu.narrow(0, offset, size).view(value.shape) - offset += size - return tree_unflatten(leaves, self._spec) - - -def pack_output_snapshot( - payload: dict[str, Any], slot: dict[tuple[Any, ...], torch.Tensor], *, max_buckets: int -) -> PackedOutputSnapshot | None: - """Copy tensor leaves once per dtype/device, retaining their exact values. - - No CUDA synchronization is introduced here. The caller must select the - producer stream and wait for the previous consumer before reusing a slot. - """ - leaves, spec = tree_flatten(payload) - grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list) - for index, value in enumerate(leaves): - if isinstance(value, torch.Tensor): - # Preserve the existing per-tensor path for unusual layouts. - if value.layout != torch.strided or value.is_quantized: - return None - grouped[(value.dtype, value.device)].append(index) - groups = [] - for (dtype, device), indices in grouped.items(): - values = [leaves[index].detach().reshape(-1) for index in indices] - size = sum(value.numel() for value in values) - key = ("packed-output", dtype, device, size) - slab = slot.get(key) - if slab is None: - slab = torch.empty(size, dtype=dtype, device=device) - if len(slot) < max_buckets: - slot[key] = slab - if len(values) == 1: - slab.copy_(values[0]) - else: - torch.cat(values, out=slab) - offset = 0 - for index, value in zip(indices, values): - shape = leaves[index].shape - leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) - offset += value.numel() - groups.append((indices, slab)) - return PackedOutputSnapshot(leaves, spec, groups) +__all__ = ["PackedOutputSnapshot", "RequestOutputSnapshot", "pack_output_snapshot"]