anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Optimize MOSS Local 1.5 MRV2 serving and progressive chunks (#8213)

Browse Source 
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
Sy03andYueqian Lin authored and GitHub committed 2026-10-02 18:03:14 +00:00
1 parent 527982d886
commit 3bc3f1a7d0
70 files changed
+8465 -293
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/configuration/environment_variables.md
+3 -3
View File
@@ -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
docs/configuration/stage_configs.md
+47
View File
@@ -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.
recipes/OpenMOSS/MOSS-TTS.md
+297
View File
@@ -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
```
tests/config/test_config_factory.py
+65 -3
View File
@@ -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
tests/config/test_environment_variables.py
+2 -2
View File
@@ -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,
tests/core/sched/test_generation_batch_coalescing.py
+253
View File
@@ -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
tests/core/sched/test_generation_scheduler_init.py
+102
View File
@@ -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))
tests/core/sched/test_omni_sched_deferred_free_fence.py
+5
View File
@@ -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]
tests/engine/test_stage_engine_core_proc.py
+66
View File
@@ -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)
tests/entrypoints/openai_api/test_moss_local_prompt.py
+85
View File
@@ -0,0 +1,85 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _Processor:
"""Segment API of the model's processor with a deterministic tokenizer."""
def __init__(self):
self.model_config = SimpleNamespace(
n_vq=3,
im_start_token_id=1,
im_end_token_id=2,
audio_start_token_id=3,
audio_pad_token_id=1024,
audio_assistant_slot_token_id=5,
)
self.encoded = []
def _encode_text(self, text):
self.encoded.append(text)
return [100 + ord(c) % 50 for c in text]
def _user_prompt_prefix_ids(self):
return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
def _user_prompt_after_reference_ids(self, language, fields):
return self._encode_text(f"lang={fields.get('language', language)}")
def _assistant_prompt_prefix_ids(self):
return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
def reference(self, text, language, codes):
# Mirrors the processor's _build_continuation_codes + _pad (one item).
fields = {} if language is None else {"language": language}
ids = (
self._user_prompt_prefix_ids()
+ self._encode_text("None")
+ self._user_prompt_after_reference_ids(language, fields)
+ self._encode_text(text)
+ self._assistant_prompt_prefix_ids()
+ [3]
)
text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
text_rows[:, 0] = torch.tensor(ids)
audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
audio_rows[:, 0] = 5
audio_rows[:, 1:] = codes
unified = torch.cat([text_rows, audio_rows])
return unified[:, 0].tolist(), unified[:, 1:].contiguous()
@pytest.mark.parametrize("language", [None, "English"])
def test_matches_processor_and_caches_constant_segments(language):
proc = _Processor()
codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
kwargs = {"text": "ref words target words"}
if language is not None:
kwargs["language"] = language
ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
assert ids == expected_ids
assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
proc.encoded.clear()
_local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
assert proc.encoded == ["other"]
def test_unsupported_layouts_fall_back():
proc = _Processor()
codes = torch.zeros((2, 3), dtype=torch.long)
assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None
tests/entrypoints/openai_api/test_tts_adapter.py
+39
View File
@@ -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)
tests/entrypoints/test_pd_disaggregation.py
+3
View File
@@ -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,
tests/metrics/test_emit_calls.py
+2 -2
View File
@@ -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,
)
tests/metrics/test_stats.py
+39
View File
@@ -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=")
tests/model_executor/models/moss_tts/test_codec_async_output.py
+338
View File
@@ -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()
tests/model_executor/models/moss_tts/test_codec_ring_headroom.py
+146
View File
@@ -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)
tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py
+547
View File
@@ -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)
tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py
+32
View File
@@ -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
tests/model_executor/models/moss_tts/test_local_compiled_sampler.py
+37
View File
@@ -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)
tests/model_executor/models/moss_tts/test_local_determined_tokens.py
+124
View File
@@ -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")
tests/model_executor/models/moss_tts/test_local_eager_mtp.py
+166
View File
@@ -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
tests/model_executor/models/moss_tts/test_local_model_state.py
+286
View File
@@ -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)
tests/model_executor/models/moss_tts/test_local_prefill_batch.py
+86
View File
@@ -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)
tests/model_executor/models/moss_tts/test_local_short_attention.py
+82
View File
@@ -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)
tests/model_executor/models/moss_tts/test_reference_encoder.py
+233
View File
@@ -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"
tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py
+320
View File
@@ -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)
tests/model_executor/models/moss_tts/test_shared_reference_encoder.py
+339
View File
@@ -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()
tests/model_executor/models/moss_tts/test_slot_attention.py
+301
View File
@@ -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"))
tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py
+86
View File
@@ -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
tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py
+25 -7
View File
@@ -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"}
tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py
+80
View File
@@ -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]
tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py
+93
View File
@@ -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]
tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py
+137
View File
@@ -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
tests/worker/test_omni_connector_mixin.py
+97 -9
View File
@@ -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()
tests/worker_v2/test_first_audio_sender.py
+16
View File
@@ -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
tests/worker_v2/test_omni_ar_model_runner.py
+31
View File
@@ -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)
tests/worker_v2/test_omni_model_state.py
+55
View File
@@ -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
vllm_omni/config/environment_variable_inventory.py
+18
View File
@@ -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",
vllm_omni/core/sched/omni_generation_scheduler.py
+144
View File
@@ -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()
vllm_omni/core/sched/omni_scheduler_mixin.py
+7 -7
View File
@@ -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
vllm_omni/core/sched/omni_scheduling_coordinator.py
+5 -2
View File
@@ -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()
vllm_omni/deploy/moss_tts_local.yaml
+3
View File
@@ -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
vllm_omni/deploy/moss_tts_local_mrv2.yaml
+30
View File
@@ -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
vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml
+53
View File
@@ -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
vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml
+66
View File
@@ -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/<dir> 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
vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py
+17
View File
@@ -1169,6 +1169,11 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
self._payload_finished(payload_data),
)
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:
@@ -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
vllm_omni/engine/stage_engine_core_proc.py
+17 -3
View File
@@ -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,13 +93,24 @@ 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
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
def _bind_native_data_plane_ready_sink(model_executor: Any, scheduler: Any) -> bool:
vllm_omni/entrypoints/omni_base.py
+1 -1
View File
@@ -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",
vllm_omni/entrypoints/openai/tts_adapters/base.py
+12 -1
View File
@@ -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,
vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py
+91 -2
View File
@@ -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,6 +344,14 @@ 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._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
@@ -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 "")
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")
if fast is not None:
text_ids, audio_codes = fast
else:
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)
text_ids = unified[:, 0].tolist()
audio_codes = unified[:, 1:].contiguous().to(torch.int64)
params = {
"prompt_token_ids": text_ids,
vllm_omni/metrics/stats.py
+63 -36
View File
@@ -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}]",
vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py
+80 -9
View File
@@ -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,6 +801,40 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
if self.rope:
q, k = self.rope(q, k, offset, time_before_heads=False)
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:
k, v, pos_k = self._complete_kv(k, v, execution_context)
pos_k = pos_k[:, None]
@@ -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
vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py
+28 -7
View File
@@ -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_()
vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py
+333
View File
@@ -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()
vllm_omni/model_executor/models/moss_tts/local_model_state.py
+447
View File
@@ -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)
vllm_omni/model_executor/models/moss_tts/local_short_attention.py
+100
View File
@@ -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)
vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py
+402 -37
View File
@@ -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,6 +112,13 @@ class _MossCodecStreamSession:
if not callable(initialize_state_pool):
raise RuntimeError("The streaming codec does not implement a decoder state pool.")
with torch.no_grad():
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":
@@ -101,14 +144,32 @@ 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:
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()
slot = self._free_stream_slots.pop(0 if oldest else -1)
self._leased_slots.add(slot)
return slot
@@ -119,9 +180,24 @@ class _MossCodecStreamSession:
return
if not state_already_reset:
self.reset_slots([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)
if not callable(reset_streaming_slots):
@@ -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,
)
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,10 +316,68 @@ class _MossCodecStreamSession:
audio_tensor = result.audio
if terminal_slots:
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]
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)
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):
@@ -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.
"""
# 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)
slot = self._stream_req_slots.get(request_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
vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py
+27 -2
View File
@@ -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,6 +136,18 @@ class _MossTTSLocalAttention(nn.Module):
v_cache[:, :, position : position + 1].copy_(value)
key = k_cache[:, :, : position + 1]
value = v_cache[:, :, : position + 1]
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:
vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+26
View File
@@ -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
# ------------------------------------------------------------------
vllm_omni/model_executor/models/moss_tts/pipeline.py
+2
View File
@@ -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,
vllm_omni/model_executor/models/moss_tts/reference_encoder.py
+410 -32
View File
@@ -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(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:
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
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]:
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._sr_target, n_vq=self._n_vq)
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 [
vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py
+399
View File
@@ -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)
]
vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py
+334
View File
@@ -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"]]
vllm_omni/model_executor/models/moss_tts/slot_attention.py
+475
View File
@@ -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)
vllm_omni/model_executor/output_snapshot.py
+125
View File
@@ -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)
vllm_omni/model_executor/stage_input_processors/moss_tts.py
+37 -33
View File
@@ -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
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),
vllm_omni/worker_v2/first_audio_sender.py
+11 -8
View File
@@ -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)
vllm_omni/worker_v2/model_states/__init__.py
+6
View File
@@ -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,
)
vllm_omni/worker_v2/omni_ar_model_runner.py
+24
View File
@@ -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)
)
vllm_omni/worker_v2/output_snapshot.py
+7 -87
View File
@@ -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"]