From e1303e5067bb5617eb12d4cb291aef6eb8b42634 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sat, 26 Sep 2026 07:51:04 -0700 Subject: [PATCH] [Model] Optimize Qwen3-TTS streaming throughput and first audio on CUDA (#8163) Signed-off-by: Sy03 <1370724210@qq.com> --- .../common/ci_source_file_dependencies.yml | 4 + docs/configuration/stage_configs.md | 73 +++- ...qwen3_omni_tts_performance_optimization.md | 75 +++- tests/config/test_environment_variables.py | 2 +- tests/config/test_omni_config.py | 51 +++ .../test_qwen3_tts_customvoice.py | 47 +++ .../test_async_omni_engine_stage_init.py | 101 +++++ tests/engine/test_cuda_mps.py | 200 ++++++++++ tests/engine/test_data_entry_keys.py | 8 + tests/engine/test_orchestrator_first_audio.py | 345 +++++++++++++++++ .../qwen3_tts/test_cached_code_predictor.py | 50 +++ .../qwen3_tts/test_code_predictor_dtype.py | 34 ++ .../qwen3_tts/test_code_predictor_sampling.py | 68 ++++ .../qwen3_tts/test_first_audio_graphs.py | 126 +++++++ .../qwen3_tts/test_qwen3_tts_code2wav.py | 34 ++ .../qwen3_tts/test_short_kv_attention.py | 44 +++ .../qwen3_tts/test_skip_delivered_first.py | 65 ++++ .../models/qwen3_tts/test_time_major_conv.py | 112 ++++++ .../qwen3_tts/test_time_major_decoder.py | 85 +++++ tests/worker_v2/test_first_audio_sender.py | 237 ++++++++++++ tests/worker_v2/test_omni_model_state.py | 201 ++++++++++ .../config/environment_variable_inventory.py | 2 + vllm_omni/config/omni_config.py | 3 + vllm_omni/config/stage_config.py | 14 +- vllm_omni/data_entry_keys.py | 6 + vllm_omni/deploy/qwen3_tts.yaml | 1 + .../deploy/qwen3_tts_high_concurrency.yaml | 1 + .../qwen3_tts_high_concurrency_mrv2.yaml | 1 + ..._tts_high_concurrency_mrv2_single_gpu.yaml | 41 ++- vllm_omni/deploy/qwen3_tts_mrv2.yaml | 5 + vllm_omni/engine/cuda_mps.py | 98 +++++ vllm_omni/engine/orchestrator.py | 128 ++++++- vllm_omni/engine/stage_engine_core_proc.py | 18 + vllm_omni/engine/stage_init_utils.py | 19 +- vllm_omni/engine/stage_runtime.py | 72 +++- .../models/common/qwen3_code_predictor.py | 101 ++--- .../models/qwen3_tts/cached_code_predictor.py | 133 +++++++ .../qwen3_tts/code_predictor_sampling.py | 75 ++++ .../models/qwen3_tts/first_audio.py | 33 ++ .../models/qwen3_tts/first_frame_decoder.py | 118 ++++++ .../models/qwen3_tts/qwen3_tts_code2wav.py | 53 ++- .../qwen3_tts_code_predictor_vllm.py | 90 ++++- .../models/qwen3_tts/qwen3_tts_talker.py | 125 ++++++- .../qwen3_tts/segmented_graph_wrapper.py | 49 ++- .../models/qwen3_tts/short_kv_attention.py | 90 +++++ .../modeling_qwen3_tts_tokenizer_v2.py | 98 +++-- .../tokenizer_12hz/time_major_conv.py | 347 ++++++++++++++++++ .../tokenizer_12hz/time_major_stack.py | 100 +++++ .../stage_input_processors/qwen3_tts.py | 12 +- vllm_omni/utils/device_copy.py | 27 ++ vllm_omni/worker_v2/first_audio_sender.py | 167 +++++++++ vllm_omni/worker_v2/model_states/eager_mtp.py | 226 ++++++++++++ .../model_states/omni_model_state.py | 222 +++++++---- vllm_omni/worker_v2/omni_ar_model_runner.py | 9 + vllm_omni/worker_v2/omni_model_runner.py | 6 + 55 files changed, 4222 insertions(+), 230 deletions(-) create mode 100644 tests/engine/test_cuda_mps.py create mode 100644 tests/engine/test_orchestrator_first_audio.py create mode 100644 tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py create mode 100644 tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py create mode 100644 tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py create mode 100644 tests/model_executor/models/qwen3_tts/test_short_kv_attention.py create mode 100644 tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py create mode 100644 tests/model_executor/models/qwen3_tts/test_time_major_conv.py create mode 100644 tests/model_executor/models/qwen3_tts/test_time_major_decoder.py create mode 100644 tests/worker_v2/test_first_audio_sender.py create mode 100644 vllm_omni/engine/cuda_mps.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/code_predictor_sampling.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/first_audio.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/first_frame_decoder.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_conv.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_stack.py create mode 100644 vllm_omni/utils/device_copy.py create mode 100644 vllm_omni/worker_v2/first_audio_sender.py create mode 100644 vllm_omni/worker_v2/model_states/eager_mtp.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index d1336e66c..02d51197b 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -20,6 +20,7 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml - vllm_omni/entrypoints/openai/serving_speech.py - vllm_omni/entrypoints/openai/tts_adapters/ - vllm_omni/engine/stage_pool.py @@ -30,6 +31,9 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/config/omni_config.py - vllm_omni/engine/orchestrator.py - vllm_omni/engine/stage_engine_core_proc.py + - vllm_omni/engine/cuda_mps.py + - vllm_omni/engine/stage_runtime.py + - vllm_omni/engine/stage_init_utils.py - vllm_omni/deploy/qwen3_tts_mrv2.yaml - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 9836bcff9..305c68360 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -337,6 +337,41 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains its existing `talker_mtp` entry point for V1. +### Qwen3-TTS first-frame delivery and rollback + +The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`, +`qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles +enable the `talker_first_audio` connector option, as does the experimental +single-GPU profile. This changes the default CUDA MRV2 streaming path: +residual prediction runs eagerly after the Talker sample, the Talker loads +an additional first-frame decoder with its weights and CUDA graphs, and the +orchestrator orders its audio before subsequent Code2Wav chunks. + +Code2Wav retains full-audio prefix graphs alongside the optional state-only +graphs. Each request's delivery marker selects the graph and audio trimming; +enabling the option alone never suppresses audio or disables prefix batching +for requests that retain regular codec delivery. + +The path requires asynchronous chunks, TP/PP 1, an in-process executor and +disabled prefix caching. Requests with reference codes and unsupported +runners or platforms retain regular Code2Wav delivery. + +To restore regular codec delivery and avoid loading the Talker's additional +decoder, use a deploy overlay with the relevant base profile: + +```yaml +base_config: qwen3_tts.yaml +connectors: + connector_of_shared_memory: + extra: + talker_first_audio: false +``` + +First-packet latency measures when PCM starts arriving. Time to first audible +audio (TTFA) also includes any leading silence in the generated audio. An +earlier first packet therefore does not necessarily improve TTFA; measure +both for the intended voice and workload. + ### Included performance work - Request snapshots have a fast path for immutable scalar leaves, including @@ -372,14 +407,34 @@ speaker similarity before adopting either batching preset for a production workload. Floating-point decoder outputs can differ across batch sizes; this PR does not claim bitwise or quality equivalence. -### Optional MPS deployment +### Experimental MPS deployment -NVIDIA MPS is an optional operator setting for colocated CUDA processes, not a -YAML option or a library default. This PR does not establish a throughput or -first-packet latency benefit from MPS. Measure the exact deployment with and -without MPS before enabling it. +NVIDIA MPS lets colocated CUDA stage processes share GPU execution resources. +It is disabled by default. The experimental +`qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile sets `cuda_mps: true`, +alongside cached residual prediction, fused sampling, time-major codec +convolutions, first-frame delivery, and larger graph batches. -Use only assigned GPUs and an independent MPS pipe directory. A private MPS -server does not provide exclusive GPU ownership or MIG isolation. For a -single-GPU deployment, explicitly place both stages on that GPU; the supplied -high-concurrency profile places its two stages on different GPUs by default. +The runtime requires `nvidia-cuda-mps-control` on `PATH` and one explicit CUDA +GPU per local EngineCore stage, with `parallel_stage_init: false`. +Use numeric GPU ordinals for stage placement and visibility so initialization +locks identify the physical GPU before MPS remaps it. It starts a private MPS daemon for each selected +GPU and stops its own daemon after the stages exit. If +`CUDA_MPS_PIPE_DIRECTORY` already names an operator-managed daemon, the runtime +reuses it without stopping it. Diffusion and remote stages are unsupported. +Set `cuda_mps: false` in a deploy overlay to disable automatic MPS management. + +Stage `env.CUDA_MPS_PIPE_DIRECTORY` overrides the parent setting when selecting +the daemon; an explicit empty string selects a private daemon. Stages sharing +a GPU reuse the first stage's daemon. Later stages may omit the setting or +explicitly match that policy; a conflicting explicit setting fails before +that stage starts. The parent environment is unchanged. + +MPS can improve throughput under concurrent load while increasing first-packet +or first-audible-audio latency, particularly at low request rates. For a +latency-sensitive workload, compare the same profile with `cuda_mps: false`. +An inherited operator-managed MPS daemon must also be disabled by its owner +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. diff --git a/docs/design/qwen3_omni_tts_performance_optimization.md b/docs/design/qwen3_omni_tts_performance_optimization.md index f90fbaad9..8de13ccce 100644 --- a/docs/design/qwen3_omni_tts_performance_optimization.md +++ b/docs/design/qwen3_omni_tts_performance_optimization.md @@ -93,13 +93,13 @@ Each optimization stacks on the previous one. The summary plots below show the c **Benchmark environment:** -| | Qwen3-Omni | Qwen3-TTS | -| --- |-----------------------------| --- | -| **GPU** | A100 | H200 | +| | Qwen3-Omni | Qwen3-TTS | +| --- | ----------------------------- | --- | +| **GPU** | A100 | H200 | | **Model** | Qwen3-Omni-30B-A3B-Instruct | Qwen3-TTS-12Hz-1.7B-CustomVoice | -| **vLLM** | v0.17.0 | v0.18.0 | -| **vllm-omni** | commit 199f7832 | v0.18.0rc2 | -| **CUDA** | 12.9 | 12.8 | +| **vLLM** | v0.17.0 | v0.18.0 | +| **vllm-omni** | commit 199f7832 | v0.18.0rc2 | +| **CUDA** | 12.9 | 12.8 | This post walks through each optimization in the same order they are typically enabled in practice, then ends with deployment playbooks for both models. @@ -310,7 +310,68 @@ self._compiled_model_fwd = torch.compile( `mode="default"` is used instead of `mode="reduce-overhead"` to avoid conflicts with vLLM's own CUDA graph capture on the main Talker model. `dynamic=True` handles the growing sequence length without recompilation. -These optimizations are always-on in the current codebase - all Qwen3-TTS benchmark results in this post include them. +The historical benchmark results above use re-prefill. It remains the default +for Qwen3-TTS and the shared Qwen3-Omni predictor. + +### Experimental Qwen3-TTS single-GPU throughput profile + +The opt-in `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile uses a +Qwen3-TTS-specific execution path on CUDA with BF16 weights: + +- `code_predictor_kv_cache` reuses each layer's keys and values across the 15 + residual-codebook predictions within one audio frame. The first step writes + two positions; subsequent steps write one. Scratch storage is bounded by the + configured batch capacity, overwritten for every frame, and never shared + between requests across frames. The worker's serialized CUDA stream owns it; + it is separate from the Talker's paged KV cache. +- A short causal/GQA Triton kernel masks unwritten K/V loads. Dynamic-batch + `torch.compile` warms the predictor before the runner captures whole-MTP + CUDA graphs. This path creates no inner CUDA graphs. Unsupported devices and + dtypes retain re-prefill; FP16 retains its existing precision safeguards. +- `code_predictor_fused_sampling` fuses temperature scaling, top-k filtering, + and Gumbel argmax when the runner supplies FP32 uniforms. It preserves top-k + ties and uses the supplied randomness. Explicit generator paths retain the + existing sampler. Qwen3-Omni does not read either TTS option. +- `decode_cudnn_benchmark` selects convolution algorithms during Code2Wav + warmup and graph capture. The process's cuDNN flags are restored afterwards, + including on capture failure. Captured graphs retain the selected algorithms. + +- `decode_time_major_conv` keeps codec convolution activations in time-major + layout to reduce transposes and copies between layers. +- `talker_first_audio` completes residual prediction immediately after the + first codebook sample and decodes the first frame in the Talker process. + The codec advances its state without repeating the delivered frame; the + orchestrator orders that frame before subsequent codec chunks. This path + requires CUDA, MRV2, asynchronous chunks, a single-process executor, TP/PP 1, + and disabled prefix caching. Reference-code requests retain codec delivery. + Supported Qwen3-TTS profiles enable it explicitly; other models do not opt in. + +The predictor, sampling, convolution-layout and autotuning options default to +false outside this throughput profile. It uses fixed 1/25-frame chunks, +72 frames of left context, Talker/codec capacities of 128/64, codec batches +through 8, and synchronous codec scheduling. Client concurrency is independent +of those capacities. This profile also sets `cuda_mps: true`; the runtime owns +a private MPS daemon unless an operator supplies an existing MPS pipe directory. +Ordinary profiles leave MPS disabled. See +[experimental MPS deployment](../configuration/stage_configs.md#experimental-mps-deployment) +for requirements and ownership. Startup includes compilation and graph capture. + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml \ + --api-server-count 4 --stage-init-timeout 1800 --init-timeout 2400 +``` + +KV reuse changes GEMM shapes and attention reduction order; cuDNN autotuning +can also change waveform numerics. This path is not bitwise equivalent to +re-prefill, and aggregate WER alone cannot establish perceptual or speaker +similarity equivalence. Validate quality on the intended workload before +adopting it. Qwen3-Omni keeps its existing predictor and audio-delivery behavior. + +For a controlled ablation, copy the profile and disable individual options while keeping capacities, chunking, hardware, client concurrency, +and warmup identical. Report repeated runs, first-audio latency and failures +alongside throughput. Warm requests measure a ready service; they do not measure +cold-start latency. --- diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 081c4d203..45b15c8e4 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -166,7 +166,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): assert category_counts == { EnvironmentVariableCategory.PUBLIC_OMNI: 30, EnvironmentVariableCategory.INHERITED_VLLM: 20, - EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, + EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, EnvironmentVariableCategory.MODEL_SPECIFIC: 58, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 2, diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 26e6f9e9c..da69ec5f1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -886,6 +886,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): def test_runtime_config_fields_match_structured_runtime_scope(): assert {f.name for f in fields(OmniStageRuntimeConfig)} == { + "cuda_mps", "additional_config", "distributed_executor_backend", "worker_cls", @@ -2051,6 +2052,28 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder): builder(pipeline, user_deploy_config=deploy) +@pytest.mark.parametrize("scope", ["pipeline", "stage"]) +@pytest.mark.parametrize("from_yaml", [False, True]) +@pytest.mark.parametrize("builder", [merge_pipeline_deploy, VllmOmniConfig.from_pipeline_config]) +def test_async_chunk_rejects_quoted_false_before_selecting_processors(tmp_path, scope, from_yaml, builder): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + if from_yaml: + path = tmp_path / "quoted_false.yaml" + path.write_text( + 'async_chunk: "false"\n' if scope == "pipeline" else 'stages:\n - stage_id: 0\n async_chunk: "false"\n' + ) + deploy = load_deploy_config(path) + elif scope == "pipeline": + deploy = DeployConfig(async_chunk="false") + else: + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, async_chunk="false")]) + with pytest.raises(ValueError, match="async_chunk must be a boolean"): + if builder is merge_pipeline_deploy: + builder(pipeline, deploy) + else: + builder(pipeline, user_deploy_config=deploy) + + @pytest.mark.parametrize("explicit", [False, True]) def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit): from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig @@ -2169,3 +2192,31 @@ def test_structured_diffusion_stage_keeps_shared_globals_outside_diffusion(): def test_structured_diffusion_stage_rejects_explicit_shared_engine_field(field_name, config_kwargs): with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): _build_single_diffusion_config(**config_kwargs) + + +def test_mps_is_explicit_only_in_experimental_single_gpu_profile(): + assert not load_deploy_config(_DEPLOY_DIR / "qwen3_tts.yaml").cuda_mps + deploy = load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml") + assert deploy.cuda_mps + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + stages = merge_pipeline_deploy(pipeline, deploy) + assert all(stage.yaml_runtime["cuda_mps"] is True for stage in stages) + + +def test_mps_config_rejects_string_boolean(tmp_path): + path = tmp_path / "deploy.yaml" + path.write_text('cuda_mps: "false"\n') + with pytest.raises(ValueError, match="cuda_mps must be a boolean"): + load_deploy_config(path) + + +def test_mps_stays_in_runtime_instead_of_engine_arguments(): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + config = _from_pipeline_key( + "qwen3_tts", deploy_config_path=str(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml") + ) + for stage in config.stage_configs: + assert stage.runtime_config.cuda_mps + args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + assert "cuda_mps" not in args diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py index 6192f8410..8b78e30c1 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py @@ -66,6 +66,53 @@ default_tts_server_params = [ ) ] +# Exercise the throughput path in the existing L4 lane with smaller capacities. +# The production H200 profile retains its own admission and graph buckets. +_FAST_PATH_STAGE_CONFIG = modify_stage_config( + get_deploy_config_path("qwen3_tts_high_concurrency_mrv2_single_gpu.yaml"), + updates={ + # The temporary overlay lives outside deploy/, so anchor its parent. + "base_config": get_deploy_config_path("qwen3_tts_high_concurrency_mrv2.yaml"), + "cuda_mps": False, + "connectors.connector_of_shared_memory.extra.decode_batch_max_size": 2, + "connectors.connector_of_shared_memory.extra.decode_cudagraph_batch_sizes": [1, 2], + "stages": { + 0: { + "max_num_seqs": 4, + "max_num_batched_tokens": 128, + "compilation_config.cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64, 128], + "compilation_config.max_cudagraph_capture_size": 128, + }, + 1: {"max_num_seqs": 4}, + }, + }, +) + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize( + "omni_server", + [OmniServerParams(model=MODEL, stage_config_path=_FAST_PATH_STAGE_CONFIG, server_args=["--trust-remote-code"])], + indirect=True, +) +def test_cached_predictor_streaming_audio(omni_server, online_client) -> None: + """Cover the optimized predictor, first audio and codec path with real weights at L3.""" + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The quick brown fox jumps over the lazy dog.", + "stream": True, + "stream_format": "audio", + "response_format": "wav", + "task_type": "CustomVoice", + "voice": "vivian", + }, + request_num=4, + ) + @pytest.mark.core_model @pytest.mark.advanced_model diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c836ec152..46ac15f9b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -10,6 +10,7 @@ import threading import time import types from dataclasses import dataclass, field +from subprocess import CompletedProcess import pytest from omegaconf import OmegaConf @@ -70,6 +71,49 @@ def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch): assert env_key not in os.environ +@pytest.mark.parametrize( + "invalid_key,invalid_value", + [("INVALID=ENV", "value"), ("INVALID\0ENV", "value"), ("INVALID_ENV", "value\0")], + ids=["equals-in-key", "nul-in-key", "nul-in-value"], +) +def test_stage_runtime_env_restores_partial_application(monkeypatch, invalid_key, invalid_value): + existing_key = "VLLM_OMNI_TEST_EXISTING_STAGE_ENV" + new_key = "VLLM_OMNI_TEST_NEW_STAGE_ENV" + monkeypatch.setenv(existing_key, "original") + monkeypatch.delenv(new_key, raising=False) + runtime_config = OmniStageRuntimeConfig( + env={existing_key: "overridden", new_key: "temporary", invalid_key: invalid_value} + ) + + with pytest.raises(ValueError): + with stage_runtime_env(0, runtime_config): + pytest.fail("Invalid environment must prevent stage launch") + + assert os.environ[existing_key] == "original" + assert new_key not in os.environ + + +def test_stage_runtime_env_restores_after_launch_error(monkeypatch): + env_key = "VLLM_OMNI_TEST_STAGE_ENV_LAUNCH_ERROR" + monkeypatch.setenv(env_key, "original") + + with pytest.raises(RuntimeError, match="stage launch failed"): + with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "temporary"})): + assert os.environ[env_key] == "temporary" + raise RuntimeError("stage launch failed") + + assert os.environ[env_key] == "original" + + +def test_stage_runtime_env_restores_first_value_for_normalized_keys(monkeypatch): + monkeypatch.setenv("123", "original") + + with stage_runtime_env(0, {"env": {123: "first", "123": "second"}}): + assert os.environ["123"] == "second" + + assert os.environ["123"] == "original" + + def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch): engine = object.__new__(AsyncOmniEngine) engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True) @@ -1019,6 +1063,63 @@ def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeyp assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None +@pytest.mark.parametrize("client_count", [1, 2]) +def test_stage_launch_uses_configured_mps_pipe_and_rejects_conflicts(monkeypatch, client_count): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine import cuda_mps + + runtime = _make_stage_runtime() + plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)] + for stage_id, plan in enumerate(plans): + plan.replicas[0].engine_args_dict = {} + plan.replicas[0].metadata.runtime_cfg = OmniStageRuntimeConfig( + cuda_mps=True, + devices="0", + env={"CUDA_MPS_PIPE_DIRECTORY": f"/operator/stage-{stage_id}"}, + ) + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent") + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0") + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) + monkeypatch.setattr(runtime_mod, "physical_gpu_uuid", lambda _: "GPU-example") + monkeypatch.setattr(runtime_mod.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(cuda_mps.shutil, "which", lambda _: "/bin/mps-control") + controls = [] + + def run_control(args, **kwargs): + controls.append(kwargs) + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent" + assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage-0" + return CompletedProcess(args, 0, stdout="") + + monkeypatch.setattr(cuda_mps.subprocess, "run", run_control) + launched = [] + + @contextlib.contextmanager + def launch(**kwargs): + launched.append(kwargs["stage_id"]) + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage-0" + assert kwargs["stage_visible_devices"] == "GPU-example" + yield StageReplicaResources( + addresses=EngineZmqAddresses( + inputs=[f"ipc://input-{i}" for i in range(client_count)], + outputs=[f"ipc://output-{i}" for i in range(client_count)], + ) + ) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch) + try: + with pytest.raises(ValueError, match="Conflicting CUDA_MPS_PIPE_DIRECTORY"): + with runtime.launch_stage_engines(client_count): + pytest.fail("Conflicting colocated MPS policies must prevent launch") + assert launched == [0] + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent" + finally: + runtime.shutdown() + assert len(controls) == 1 + assert controls[0]["input"] == "get_server_list\n" + + @pytest.mark.parametrize("client_count", [1, 2]) def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count): import fcntl diff --git a/tests/engine/test_cuda_mps.py b/tests/engine/test_cuda_mps.py new file mode 100644 index 000000000..7efc74893 --- /dev/null +++ b/tests/engine/test_cuda_mps.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from vllm_omni.config.omni_config import OmniStageRuntimeConfig +from vllm_omni.engine import cuda_mps +from vllm_omni.engine.stage_runtime import StageRuntime + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def controls(monkeypatch): + monkeypatch.delenv("CUDA_MPS_PIPE_DIRECTORY", raising=False) + monkeypatch.setattr(cuda_mps.shutil, "which", lambda name: "/bin/mps-control") + run = Mock(return_value=SimpleNamespace(stdout="")) + monkeypatch.setattr(cuda_mps.subprocess, "run", run) + return run + + +def test_private_server_has_scoped_environment_and_idempotent_close(controls): + original = dict(os.environ) + server = cuda_mps.CudaMPSServer("GPU-example") + pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"]) + assert pipe.is_dir() + assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"] + assert controls.call_args_list[0].kwargs["env"]["CUDA_VISIBLE_DEVICES"] == "GPU-example" + assert dict(os.environ) == original + server.close() + count = controls.call_count + server.close() + assert controls.call_count == count + assert controls.call_args.kwargs["input"] == "quit\n" + assert not pipe.exists() + + +def test_operator_daemon_is_checked_but_never_stopped(controls, monkeypatch): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/mps") + server = cuda_mps.CudaMPSServer("GPU-example") + server.close() + assert controls.call_count == 1 + assert controls.call_args.kwargs["input"] == "get_server_list\n" + + +def test_explicit_pipe_overrides_parent_without_taking_ownership(controls, monkeypatch): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent") + server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="/operator/stage") + + assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage" + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent" + server.close() + assert controls.call_count == 1 + + +def test_explicit_empty_pipe_selects_private_daemon(controls, monkeypatch): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent") + server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="") + pipe = server.env["CUDA_MPS_PIPE_DIRECTORY"] + try: + assert pipe != "/operator/parent" + assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"] + assert controls.call_args_list[0].kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent" + finally: + server.close() + assert controls.call_args.kwargs["input"] == "quit\n" + assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe + assert not Path(pipe).exists() + + +@pytest.fixture +def mps_runtime(monkeypatch): + from vllm_omni.engine import stage_runtime + + runtime = StageRuntime.__new__(StageRuntime) + runtime._mps_servers = {} + monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example") + monkeypatch.setattr(stage_runtime.current_omni_platform, "is_cuda", lambda: True) + yield runtime + runtime._close_mps_servers() + + +@pytest.mark.parametrize("pipe_directory", ["/operator/stage", ""]) +def test_runtime_shares_first_stage_pipe_policy(controls, monkeypatch, mps_runtime, pipe_directory): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent") + config = OmniStageRuntimeConfig(cuda_mps=True, env={"CUDA_MPS_PIPE_DIRECTORY": pipe_directory}) + first_env = mps_runtime._mps_environment("0", config) + selected_pipe = first_env["CUDA_MPS_PIPE_DIRECTORY"] + if pipe_directory: + assert selected_pipe == pipe_directory + else: + assert selected_pipe != "/operator/parent" + assert Path(selected_pipe).is_dir() + call_count = controls.call_count + + assert mps_runtime._mps_environment("0", {"cuda_mps": True}) is first_env + assert mps_runtime._mps_environment("0", config) is first_env + assert controls.call_count == call_count + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent" + + +@pytest.mark.parametrize( + "first_pipe,conflicting_pipe", + [("/operator/first", "/operator/other"), ("/operator/first", ""), ("", "/operator/other")], +) +def test_runtime_rejects_conflicting_explicit_stage_pipes(controls, mps_runtime, first_pipe, conflicting_pipe): + mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": first_pipe}}) + call_count = controls.call_count + + with pytest.raises(ValueError, match="Conflicting CUDA_MPS_PIPE_DIRECTORY"): + mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": conflicting_pipe}}) + + assert controls.call_count == call_count + + +def test_startup_failure_closes_only_the_private_daemon(controls): + controls.side_effect = [SimpleNamespace(stdout=""), RuntimeError("health failed"), SimpleNamespace(stdout="")] + with pytest.raises(RuntimeError, match="health failed"): + cuda_mps.CudaMPSServer("GPU-example") + assert controls.call_args.kwargs["input"] == "quit\n" + + +def test_runtime_opt_in_and_shared_gpu_ownership(monkeypatch): + from vllm_omni.engine import stage_runtime + + runtime = StageRuntime.__new__(StageRuntime) + runtime._mps_servers = {} + create = Mock(return_value=SimpleNamespace(env={"CUDA_VISIBLE_DEVICES": "GPU-example"}, close=Mock())) + monkeypatch.setattr(stage_runtime, "CudaMPSServer", create) + monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example") + monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True)) + assert runtime._mps_environment("0", {}) == {} + assert create.call_count == 0 + for _ in range(2): + assert runtime._mps_environment("0", {"cuda_mps": True}) == {"CUDA_VISIBLE_DEVICES": "GPU-example"} + assert create.call_count == 1 + runtime._close_mps_servers() + runtime._close_mps_servers() + create.return_value.close.assert_called_once() + + +@pytest.mark.parametrize("devices", [None, "0,1"]) +def test_runtime_rejects_ambiguous_placement(devices, monkeypatch): + from vllm_omni.engine import stage_runtime + + runtime = StageRuntime.__new__(StageRuntime) + monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True)) + with pytest.raises(ValueError, match="exactly one explicit GPU"): + runtime._mps_environment(devices, {"cuda_mps": True}) + + +def test_failed_shutdown_preserves_socket_for_retry(controls): + server = cuda_mps.CudaMPSServer("GPU-example") + pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"]) + controls.side_effect = RuntimeError("quit failed") + with pytest.raises(RuntimeError, match="quit failed"): + server.close() + assert pipe.is_dir() + controls.side_effect = None + server.close() + assert not pipe.exists() + + +@pytest.mark.parametrize("launch_mode,stage_type", [("remote", "llm"), ("local", "diffusion")]) +def test_mps_rejects_unsupported_stages(launch_mode, stage_type): + runtime = StageRuntime.__new__(StageRuntime) + metadata = SimpleNamespace(runtime_cfg={"cuda_mps": True}, stage_type=stage_type) + stage_cfg = SimpleNamespace(runtime=metadata.runtime_cfg) + plans = [ + SimpleNamespace(replicas=[SimpleNamespace(metadata=metadata, stage_cfg=stage_cfg, launch_mode=launch_mode)]) + ] + with pytest.raises(ValueError, match="local EngineCore"): + runtime._validate_mps_topology(plans) + metadata.runtime_cfg["cuda_mps"] = False + runtime._validate_mps_topology(plans) + + +def test_mps_rejects_parallel_initialization(): + runtime = StageRuntime.__new__(StageRuntime) + runtime._parallel_stage_init = True + replica = SimpleNamespace( + metadata=SimpleNamespace(stage_type="llm"), + stage_cfg=SimpleNamespace(runtime={"cuda_mps": True}), + launch_mode="local", + ) + with pytest.raises(ValueError, match="parallel_stage_init=false"): + runtime._validate_mps_topology([SimpleNamespace(replicas=[replica])]) + + +def test_failed_launch_command_still_checks_private_daemon_cleanup(controls): + controls.side_effect = [RuntimeError("launch failed after fork"), SimpleNamespace(stdout="")] + with pytest.raises(RuntimeError, match="launch failed after fork"): + cuda_mps.CudaMPSServer("GPU-example") + assert controls.call_args.kwargs["input"] == "quit\n" diff --git a/tests/engine/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py index df0d7b2ed..dc9a196ca 100644 --- a/tests/engine/test_data_entry_keys.py +++ b/tests/engine/test_data_entry_keys.py @@ -425,3 +425,11 @@ class TestSerializeDeserializePayload: restored = deserialize_payload(wire) assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor + + +def test_first_audio_metadata_round_trips_as_a_tensor(): + from vllm_omni.data_entry_keys import to_dict, to_struct + + payload = {"meta": {"first_audio": torch.tensor(True)}} + restored = to_dict(to_struct(payload)) + assert restored["meta"]["first_audio"].item() is True diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py new file mode 100644 index 000000000..5e467254b --- /dev/null +++ b/tests/engine/test_orchestrator_first_audio.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import FinishReason + +from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _CodecPool: + final_output = True + + def __init__(self, output_kind, registered, replica_id): + self.output_processor = MultimodalOutputProcessor( + tokenizer=None, log_stats=False, engine_core_output_type="audio" + ) + self.output_kind = output_kind + self.replica_id = replica_id + self.bound_replica = None + if registered: + self.register() + + def register(self): + self.bound_replica = self.replica_id + self.output_processor.request_states["r"] = OmniRequestState( + request_id="r", + external_req_id="r", + parent_req=None, + request_index=0, + lora_request=None, + prompt=None, + prompt_token_ids=[0], + prompt_embeds=None, + logprobs_processor=None, + detokenizer=None, + max_tokens_param=None, + arrival_time=0.0, + queue=None, + log_stats=False, + stream_interval=1, + output_kind=self.output_kind, + ) + self.output_processor.external_req_ids["r"].append("r") + + def get_bound_replica_id(self, request_id): + return self.bound_replica + + async def process_llm_raw_outputs(self, replica_id, raw, **kwargs): + assert replica_id == self.replica_id + return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs + + +def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1): + obj = Orchestrator.__new__(Orchestrator) + obj.stage_pools = [SimpleNamespace(final_output=False), _CodecPool(output_kind, registered, replica_id)] + obj.request_states = {"r": OrchestratorRequestState(request_id="r", final_stage_id=final_stage_id)} + obj.output_async_queue = asyncio.Queue() + obj._cfg_tracker = CfgCompanionTracker() + obj._pd_kv_params = {} + obj._running_counter = None + obj._abort_request_ids = AsyncMock(return_value=[]) + obj._release_request_bindings = Mock() + obj._finish_raw_terminal_requests = AsyncMock() + + async def process(stage, replica, raw, terminals): + await obj._route_upstream_first_audio(stage, replica, raw) + return await obj.stage_pools[stage].process_llm_raw_outputs(replica, raw) + + async def route(stage, replica, outputs): + for output in outputs: + await obj.output_async_queue.put( + OutputMessage( + request_id=output.request_id, + stage_id=stage, + replica_id=replica, + engine_outputs=output, + metrics=None, + finished=output.finished, + ) + ) + + obj._process_llm_stage_outputs = process + obj._handle_processed_outputs = route + return obj + + +def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)): + return OmniEngineCoreOutputs( + outputs=[ + OmniEngineCoreOutput( + request_id="r", + new_token_ids=[], + finish_reason=FinishReason.STOP if terminal else None, + multimodal_output={ + "model_outputs": torch.tensor(samples, dtype=torch.float32), + "sr": torch.tensor(24000), + "_omni_first_audio": torch.tensor(first), + "_omni_first_audio_required": torch.tensor(required), + }, + ) + ] + ) + + +async def _codec_output(obj, raw): + replica = obj.stage_pools[1].replica_id + processed = await obj._process_llm_stage_outputs(1, replica, raw, set()) + await obj._handle_processed_outputs(1, replica, processed) + + +def _messages(obj): + messages = [] + while not obj.output_async_queue.empty(): + messages.append(obj.output_async_queue.get_nowait()) + return messages + + +def _audio(output): + audio = output.outputs[0].multimodal_output["audio"] + return torch.cat(audio) if isinstance(audio, list) else audio + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", list(RequestOutputKind)) +@pytest.mark.parametrize("codec_overtook", [False, True]) +async def test_first_audio_obeys_codec_output_kind(kind, codec_overtook): + obj = _orchestrator(kind) + suffix = _raw(required=True, samples=(20, 21)) + terminal = _raw(terminal=True, samples=(30, 31)) + if codec_overtook: + await _codec_output(obj, suffix) + await _codec_output(obj, terminal) + assert obj.output_async_queue.empty() + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + if kind == RequestOutputKind.FINAL_ONLY and not codec_overtook: + assert obj.output_async_queue.empty() + if not codec_overtook: + await _codec_output(obj, suffix) + await _codec_output(obj, terminal) + messages = _messages(obj) + # The source replica is zero; output ownership belongs to codec replica two. + assert all((message.stage_id, message.replica_id) == (1, 2) for message in messages) + actual = [_audio(message.engine_outputs).tolist() for message in messages] + expected = { + RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]], + RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]], + RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]], + } + assert actual == expected[kind] + assert messages[-1].finished + assert obj.request_states["r"].pending_first_audio_outputs == [] + await obj._route_upstream_first_audio(0, 0, _raw(first=True)) + assert obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_first_audio_waits_for_codec_registration_and_ignores_duplicates(): + obj = _orchestrator(registered=False) + req_state = obj.request_states["r"] + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(90, 91))) + assert obj.output_async_queue.empty() + assert not req_state.upstream_first_audio + assert req_state.pending_upstream_first_audio is not None + obj.stage_pools[1].register() + await obj._flush_upstream_first_audio(req_state) + messages = _messages(obj) + assert len(messages) == 1 + assert _audio(messages[0].engine_outputs).tolist() == [10, 11] + assert req_state.upstream_first_audio + assert req_state.pending_upstream_first_audio is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("processing_failed", [False, True]) +async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed): + obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1) + req_state = obj.request_states["r"] + req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1) + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + pool = obj.stage_pools[1] + pool.stage_type = "llm" + pool.stage_vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None)) + obj.stage_pools[0].get_bound_replica_id = lambda request_id: 0 + obj._stage_receives_async_chunks = lambda stage: True + obj._build_payload_sender_info = Mock(return_value={}) + obj._on_stage_submitted = Mock() + obj._emit_tx_edge = Mock() + mocker.patch( + "vllm_omni.engine.orchestrator.build_engine_core_request_from_tokens", + return_value=SimpleNamespace(request_id="r"), + ) + + async def submit(*args, **kwargs): + pool.register() + + pool.submit_initial = submit + if processing_failed: + pool.process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) + later_pool = SimpleNamespace(submit_initial=AsyncMock()) + obj.stage_pools.append(later_pool) + submitted = await obj._prewarm_async_chunk_stages("r", SimpleNamespace(prompt_token_ids=[0]), req_state) + assert submitted is not processing_failed + messages = _messages(obj) + assert len(messages) == 1 + if processing_failed: + assert isinstance(messages[0], ErrorMessage) + later_pool.submit_initial.assert_not_called() + assert "r" not in obj.request_states + else: + assert _audio(messages[0].engine_outputs).tolist() == [10, 11] + + +@pytest.mark.asyncio +async def test_pending_first_audio_is_discarded_on_cancellation(): + obj = _orchestrator(registered=False) + req_state = obj.request_states["r"] + await obj._route_upstream_first_audio(0, 0, _raw(first=True)) + await obj._cleanup_request_ids(["r"], abort=True) + obj.stage_pools[1].register() + await obj._flush_upstream_first_audio(req_state) + assert obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_cancellation_while_routing_first_audio_does_not_release_codec_outputs(): + obj = _orchestrator() + await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True)) + + async def cancel(stage, replica, outputs): + await obj._cleanup_request_ids(["r"], abort=True) + + obj._handle_processed_outputs = cancel + obj._process_llm_stage_outputs = AsyncMock() + await obj._route_upstream_first_audio(0, 0, _raw(first=True)) + obj._process_llm_stage_outputs.assert_not_called() + assert obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_request_ending_at_source_does_not_emit_first_audio(): + obj = _orchestrator(final_stage_id=0) + raw = _raw(first=True) + await obj._route_upstream_first_audio(0, 0, raw) + assert not raw.outputs and obj.output_async_queue.empty() + assert obj.request_states["r"].pending_upstream_first_audio is None + + +@pytest.mark.asyncio +async def test_later_aligner_receives_full_codec_waveform(): + from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner + + obj = _orchestrator(RequestOutputKind.CUMULATIVE, final_stage_id=2) + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + await _codec_output(obj, _raw(required=True, terminal=True, samples=(20, 21))) + messages = _messages(obj) + assert all(message.stage_id == 1 for message in messages) + inputs = code2wav2aligner( + [messages[-1].engine_outputs], prompt={"additional_information": {"text": ["Hello world"]}} + ) + assert inputs[0]["multi_modal_data"]["audio"][0].tolist() == [10, 11, 20, 21] + + +@pytest.mark.asyncio +async def test_first_audio_processing_failure_fails_only_its_request(): + obj = _orchestrator() + obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1) + obj.stage_pools[1].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) + await obj._route_upstream_first_audio(0, 0, _raw(first=True)) + error = obj.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == "r" and error.stage_id == 1 + assert "processing failed" in error.error + assert set(obj.request_states) == {"healthy"} + + +@pytest.mark.asyncio +async def test_unmarked_upstream_audio_and_regular_codec_outputs_are_untouched(): + obj = _orchestrator() + for stage in (0, 1): + raw = _raw() + await obj._route_upstream_first_audio(stage, 0, raw) + assert len(raw.outputs) == 1 + assert torch.equal(raw.outputs[0].multimodal_output["model_outputs"], torch.ones(4)) + assert obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_cancelled_request_does_not_receive_late_first_audio(): + obj = _orchestrator() + obj.request_states.clear() + raw = _raw(first=True) + await obj._route_upstream_first_audio(0, 0, raw) + assert not raw.outputs and obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_first_audio_supports_both_speech_and_chat_consumers(): + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + + obj = _orchestrator() + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + result = obj.output_async_queue.get_nowait().engine_outputs + assert not result.finished + assert _audio(result).tolist() == [10, 11] + payload, key = OmniOpenAIServingSpeech._extract_audio_output(result) + assert payload is not None and key is not None + audio = payload[key] + audio = torch.cat(audio) if isinstance(audio, list) else audio + assert audio.tolist() == [10, 11] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("audio", [None, [], torch.empty(0)]) +@pytest.mark.parametrize("codec_overtook", [False, True]) +async def test_invalid_first_audio_reports_error_and_releases_request(audio, codec_overtook): + obj = _orchestrator() + obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1) + if codec_overtook: + await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True)) + assert obj.request_states["r"].pending_first_audio_outputs + first = _raw(first=True) + first.outputs[0].multimodal_output["model_outputs"] = audio + await obj._route_upstream_first_audio(0, 0, first) + + error = obj.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == "r" and "first-audio" in error.error + assert not first.outputs and set(obj.request_states) == {"healthy"} + obj._abort_request_ids.assert_awaited_once_with(["r"]) + obj._release_request_bindings.assert_called_once_with(["r"]) + await obj._route_upstream_first_audio(0, 0, _raw(first=True)) + assert obj.output_async_queue.empty() diff --git a/tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py b/tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py new file mode 100644 index 000000000..476ade06e --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Cache isolation and graph replay for the opt-in TTS predictor.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorBaseModel +from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_cached_frames_match_causal_prefill_after_graph_replay(dtype): + from vllm_omni.model_executor.models.qwen3_tts.cached_code_predictor import FrameLocalKVCache + + config = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_code_groups=4, + ) + torch.manual_seed(171) + model = CodePredictorBaseModel(config).to(device="cuda", dtype=dtype).eval() + inputs = torch.randn(3, 5, 64, device="cuda", dtype=dtype) + positions = torch.arange(5, device="cuda").expand(3, -1) + cache = FrameLocalKVCache(model, max_batch=3) + with torch.inference_mode(): + for batch in (1, 3): + for _ in range(2): + for step in (1, 2, 3): + cache(inputs, batch, step) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = torch.cat([cache(inputs, 3, step) for step in (1, 2, 3)], dim=1) + for _ in range(3): + inputs.normal_() + # No previous frame (including a smaller batch) may leak through. + cache(inputs, 1, 1) + cache.buffer.fill_(torch.nan) + graph.replay() + expected = model(inputs, positions)[:, :4] + tolerance = 1e-5 if dtype == torch.float32 else 0.04 + torch.testing.assert_close(output, expected, atol=tolerance, rtol=tolerance) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 587993fca..fe19b44a9 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -681,6 +681,7 @@ class TestCodePredictorGraphReplay: sampling_mode="per_call", return_proj_buf=False, ) + predictor._frame_cache = None predictor._prefix_graphs_enabled = False predictor._prefix_reprefill_enabled = False predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)} @@ -1398,3 +1399,36 @@ class TestMTPExecutionBuckets: ) wrapper2._bucket_sizes = [1, 2, 4] # stand-in for warmed-up state wrapper2.configure_mtp_execution_buckets([3]) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_tts_cache_falls_back_on_cpu(mocker, loaded_target_classes, enabled): + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + cfg = _make_vllm_config(mocker) + cfg.model_config.stage_connector_config = {"extra": {"code_predictor_kv_cache": enabled}} + predictor = loaded_target_classes[2](vllm_config=cfg, config=cp_config, talker_config=talker_config) + common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + setup = mocker.patch.object(common.CodePredictorWrapper, "_setup_compile") + predictor._setup_compile() + setup.assert_called_once_with() + assert predictor._frame_cache is None + + +def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_classes): + cp_config, _ = _make_tiny_config(loaded_target_classes) + cfg = _make_vllm_config(mocker) + cfg.model_config.stage_connector_config = { + "extra": {"code_predictor_kv_cache": True, "code_predictor_fused_sampling": True} + } + common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + predictor = common.CodePredictorWrapper( + vllm_config=cfg, + cp_config=cp_config, + wrapper_config=common.CodePredictorWrapperConfig(), + ) + assert not hasattr(predictor, "_frame_cache") + logits = torch.randn(3, cp_config.vocab_size) + uniforms = torch.rand_like(logits) + expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) + actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms) + torch.testing.assert_close(actual, expected) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py new file mode 100644 index 000000000..30994e36b --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Supplied-uniform samples must match the existing ATen path exactly.""" + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +@pytest.mark.parametrize("batch", [1, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("top_k", [-1, 1, 50, 2048]) +def test_topk_gumbel_exact_samples_and_graph(batch, dtype, top_k): + from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel + + with torch.inference_mode(): + torch.manual_seed(42) + logits = torch.randn(batch, 2048, device="cuda", dtype=dtype) * 5 + uniforms = torch.rand(batch, 15, 2048, device="cuda")[:, 7, :] + uniforms.clamp_(1e-6, 1 - 1e-6) + for case in ("random", "ties", "all_masked", "part_masked", "uniform_zero", "uniform_one", "nan"): + if case == "ties": + logits.zero_() + uniforms.fill_(0.5) + elif case == "all_masked": + logits.fill_(-torch.inf) + elif case == "part_masked": + logits[:, 4:9] = 1 + elif case == "uniform_zero": + logits.normal_() + uniforms.zero_() + elif case == "uniform_one": + uniforms.fill_(1) + elif case == "nan": + uniforms.fill_(0.5) + logits[:, 7] = torch.nan + scaled = logits * (1 / 0.9) + if top_k > 0: + threshold = scaled.topk(top_k, dim=-1).values[:, -1:] + scaled = scaled.masked_fill(scaled < threshold, -torch.inf) + expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) + actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9) + graph.replay() + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("top_k", [0, 1, 50, 97]) +def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k): + from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel + + torch.manual_seed(81) + logits = torch.randn(3, 194, device="cuda", dtype=torch.bfloat16)[:, ::2] + uniforms = torch.rand(3, 15, 194, device="cuda")[:, 3, ::2].clamp_(1e-6, 1 - 1e-6) + scaled = logits * (1 / 0.9) + if top_k > 0: + scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf) + expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) + actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9) + torch.testing.assert_close(actual, expected) diff --git a/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py b/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py new file mode 100644 index 000000000..ce92f06a6 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Prefix graphs must follow per-request first-audio delivery, not a capture hint.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _decoder_stub +from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Decoder, +) + +pytestmark = [pytest.mark.core_model] +DEVICE = torch.device("cuda:0") + + +@pytest.mark.cpu +@pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)]) +@pytest.mark.parametrize("initial_frames", [1, 3]) +@pytest.mark.parametrize("state_only_available", [False, True]) +def test_xvec_first_chunk_replays_prefix_graph(monkeypatch, skip_flags, initial_frames, state_only_available): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = initial_frames + wrapper.codec_chunk_frames = 25 + wrapper.capture_batch_sizes = [1, 2] + wrapper._icl_previous_frames_by_target = {26: 1} + wrapper._xvec_previous_frames_by_target = {initial_frames + 25: initial_frames} + replayed = [] + wrapper.xvec_prefix_states = { + 2: { + "graph": SimpleNamespace(replay=lambda: replayed.append("full")), + "input": {"codes": torch.zeros(2, 2, initial_frames)}, + "output": torch.arange(4 * initial_frames, dtype=torch.float32).view(2, 1, 2 * initial_frames), + "cache": { + "ref_hidden": torch.zeros(2, 2, 0), + "ref_conv": torch.zeros(2, 0, 3), + "prefix_hidden": torch.zeros(2, 0, 3), + "ref_upsample": torch.zeros(2, 3, 0), + "ref_wav": torch.zeros(2, 1, 0), + "suffix_quantized": torch.ones(2, 2, initial_frames), + "suffix_conv": torch.ones(2, initial_frames, 3), + "past_key_values": SimpleNamespace(layers=[]), + }, + } + } + wrapper.xvec_prefix_state_only_states = {} + if state_only_available and initial_frames == 1: + wrapper.xvec_prefix_state_only_states[2] = { + **wrapper.xvec_prefix_states[2], + "graph": SimpleNamespace(replay=lambda: replayed.append("state")), + "state_only": True, + } + wrapper._record_graph_hit = lambda *_args: None + wrapper._record_graph_fallback = lambda *_args: None + wrapper._ensure_suffix_buffers = lambda cache: None + wrapper.decoder = _decoder_stub( + capture_first_audio_state_only=True, + total_upsample=2, + _is_suffix_cache_rolling=lambda previous, cached: False, + _decode_xvec_first_chunk=lambda *_args: pytest.fail("unexpected eager fallback"), + _slice_dynamic_cache=Qwen3TTSTokenizerV2Decoder._slice_dynamic_cache, + ) + + caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags] + outputs = wrapper._batched_request_decode( + [torch.full((1, 2, initial_frames), 3), torch.full((1, 2, initial_frames), 4)], + caches, + ) + + assert replayed == ["state" if all(skip_flags) and state_only_available and initial_frames == 1 else "full"] + torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][0], torch.full((2, initial_frames), 3.0)) + torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][1], torch.full((2, initial_frames), 4.0)) + assert [cache["decoder_prefix_frames"] for cache in caches] == [0, 0] + assert [cache["suffix_frames"] for cache in caches] == [initial_frames, initial_frames] + assert len(outputs) == 2 + for row, skip in enumerate(skip_flags): + expected = wrapper.xvec_prefix_states[2]["output"][row : row + 1, :, 2 if skip else 0 :] + torch.testing.assert_close(outputs[row], expected) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("initial_frames", [1, 3]) +@pytest.mark.parametrize("time_major", [False, True]) +@torch.inference_mode() +def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import ( + CUDAGraphDecoderWrapper as SegmentedWrapper, + ) + + decoder = _make_decoder().to(DEVICE) + decoder.precompute_snake_caches() + if time_major: + decoder.enable_time_major_conv() + decoder.capture_first_audio_state_only = True + wrapper = SegmentedWrapper( + decoder, + capture_modes=("xvec",), + capture_batch_sizes=[1, 2], + num_quantizers=2, + initial_chunk_frames=initial_frames, + ) + # This regression covers prefix graphs. Avoid compiling unrelated suffix + # shapes, especially the time-major kernels for long continuation chunks. + wrapper._xvec_previous_frames_by_target = {} + wrapper.warmup(DEVICE) + assert set(wrapper.xvec_prefix_states) == {1, 2} + assert set(wrapper.xvec_prefix_state_only_states) == ({1, 2} if initial_frames == 1 else set()) + + for skip_flags in ((False, False), (True, True), (False, True), (True, False)): + codes = [torch.randint(0, 32, (1, 2, initial_frames), device=DEVICE) for _ in skip_flags] + caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags] + expected_caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags] + expected = [decoder._decode_stream_first_chunk(code, cache) for code, cache in zip(codes, expected_caches)] + actual = wrapper._decode_xvec_prefix_batch(codes, caches) + assert actual is not None # Eligible batches must not silently become per-request eager calls. + for got, want, cache, expected_cache in zip(actual, expected, caches, expected_caches): + torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4) + for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"): + torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index d4ac126cb..472077c0e 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -752,3 +752,37 @@ def test_load_weights_uses_model_dtype_before_precomputing_caches(dtype): assert _load_weights_noop(model) == {"decoder.fake_weight"} assert model.decoder.weight.dtype is dtype assert cache_dtypes == [dtype] + + +@pytest.mark.parametrize("capture_fails", [False, True]) +def test_decode_autotune_restores_process_flags(mocker, monkeypatch, capture_fails): + monkeypatch.setattr(torch.backends.cudnn, "benchmark", False) + monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 5) + original = (torch.backends.cudnn.enabled, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32) + model = _make_model( + async_chunk=True, + device=torch.device("cuda"), + stage_connector_config={"extra": {"decode_cudnn_benchmark": True}}, + ) + + observed = [] + + def capture(**kwargs): + observed.append( + ( + torch.backends.cudnn.benchmark, + torch.backends.cudnn.benchmark_limit, + torch.backends.cudnn.enabled, + torch.backends.cudnn.deterministic, + torch.backends.cudnn.allow_tf32, + ) + ) + if capture_fails: + raise RuntimeError("capture failed") + + capture_mock = mocker.patch.object(model, "_maybe_enable_decoder_cudagraph", side_effect=capture) + _load_weights_noop(model) + capture_mock.assert_called_once() + assert observed == [(True, 10, *original)] + assert torch.backends.cudnn.benchmark is False + assert torch.backends.cudnn.benchmark_limit == 5 diff --git a/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py b/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py new file mode 100644 index 000000000..15dcecb14 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Causal/GQA correctness and graph replay for the frame-local KV kernel.""" + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +@pytest.mark.parametrize("batch", [1, 3, 64]) +@pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)]) +@pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]]) +def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions): + from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention + + with torch.inference_mode(): + torch.manual_seed(42) + pos = torch.tensor(positions, device="cuda") + q = torch.randn(batch, len(positions), heads, dim, device="cuda", dtype=torch.bfloat16).transpose(1, 2) + cache = torch.randn(2, batch + 3, kv_heads, 17, dim, device="cuda", dtype=torch.bfloat16) + k, v = cache[0, :batch], cache[1, :batch] + mask = torch.arange(17, device="cuda")[None, :] <= pos[:, None] + kr = k.float().repeat_interleave(heads // kv_heads, dim=1) + vr = v.float().repeat_interleave(heads // kv_heads, dim=1) + scores = (q.float() @ kr.transpose(-1, -2)) * dim**-0.5 + expected = scores.masked_fill(~mask, -torch.inf).softmax(-1) @ vr + actual = short_kv_attention(q, k, v, pos, dim**-0.5) + torch.testing.assert_close(actual.float(), expected, atol=0.016, rtol=0.008) + + # Unwritten cache slots must not affect a new frame, even if poisoned. + k[:, :, max(positions) + 1 :] = torch.nan + v[:, :, max(positions) + 1 :] = torch.nan + poisoned = short_kv_attention(q, k, v, pos, dim**-0.5) + torch.testing.assert_close(poisoned, actual, atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = short_kv_attention(q, k, v, pos, dim**-0.5) + graph.replay() + torch.testing.assert_close(captured, actual, atol=0, rtol=0) diff --git a/tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py b/tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py new file mode 100644 index 000000000..20d017793 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Code2Wav drops only frames explicitly marked as delivered by the Talker.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Decoder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +UPSAMPLE = 4 + + +def _decoder(skip: bool): + decoder = Qwen3TTSTokenizerV2Decoder.__new__(Qwen3TTSTokenizerV2Decoder) + decoder.__dict__["capture_first_audio_state_only"] = skip + decoder.__dict__["total_upsample"] = UPSAMPLE + calls = [] + + def full(codes, caches): + calls.append("full") + frames = int(codes.shape[-1]) + return torch.arange(frames * UPSAMPLE, dtype=torch.float32).view(1, 1, -1) + + def state_only(codes, caches): + calls.append("state") + return torch.zeros(1) + + decoder.__dict__["_decode_xvec_first_chunk"] = full + decoder.__dict__["_decode_xvec_first_chunk_state_only"] = state_only + return decoder, calls + + +def test_one_frame_first_chunk_only_builds_state(): + decoder, calls = _decoder(skip=True) + out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {"skip_first_audio": True}) + assert out.shape[-1] == 0 and calls == ["state"] + + +def test_longer_first_chunk_keeps_all_but_the_delivered_frame(): + decoder, calls = _decoder(skip=True) + out = decoder._decode_stream_first_chunk( + torch.zeros(1, 16, 3, dtype=torch.long), + {"skip_first_audio": True}, + ) + assert calls == ["full"] + assert torch.equal(out, torch.arange(UPSAMPLE, 3 * UPSAMPLE, dtype=torch.float32).view(1, 1, -1)) + + +def test_without_skip_the_whole_first_chunk_is_decoded(): + decoder, calls = _decoder(skip=False) + out = decoder._decode_stream_first_chunk( + torch.zeros(1, 16, 3, dtype=torch.long), + {"skip_first_audio": False}, + ) + assert calls == ["full"] and out.shape[-1] == 3 * UPSAMPLE + + +def test_request_without_first_audio_uses_regular_decoder_even_when_graph_skip_is_enabled(): + decoder, calls = _decoder(skip=True) + out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {}) + assert calls == ["full"] and out.shape[-1] == UPSAMPLE diff --git a/tests/model_executor/models/qwen3_tts/test_time_major_conv.py b/tests/model_executor/models/qwen3_tts/test_time_major_conv.py new file mode 100644 index 000000000..9422f01ea --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_time_major_conv.py @@ -0,0 +1,112 @@ +# 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 + +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import time_major_conv as tmc + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _snake_nct(x, exp_alpha, inv_beta): + return x + inv_beta[None, :, None] * torch.sin(x * exp_alpha[None, :, None]).square() + + +@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 1), (7, 3), (3, 9)]) +@pytest.mark.parametrize("with_bias_and_snake", [False, True]) +def test_im2col_gemm_matches_causal_conv1d(kernel_size, dilation, with_bias_and_snake): + torch.manual_seed(0) + bsz, t, c_in, c_out = 2, 11, 6, 5 + conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation) + x = torch.randn(bsz, c_in, t) + in_bias = torch.randn(c_in) if with_bias_and_snake else None + snake = (torch.rand(c_in) + 0.5, torch.rand(c_in) + 0.5) if with_bias_and_snake else None + + ref_in = x if in_bias is None else x + in_bias[None, :, None] + if snake is not None: + ref_in = _snake_nct(ref_in, *snake) + expected = conv(F.pad(ref_in, ((kernel_size - 1) * dilation, 0))) + + cols = tmc.snake_im2col(x.transpose(1, 2).contiguous(), kernel_size, dilation, in_bias, snake) + actual = torch.addmm(conv.bias, cols, tmc.conv_gemm_weight(conv.weight)).view(bsz, t, c_out) + torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize(("stride", "kernel_size"), [(2, 2), (3, 6), (8, 16)]) +def test_transposed_overlap_add_matches_trimmed_conv_transpose(stride, kernel_size): + torch.manual_seed(0) + bsz, t, c_in, c_out = 2, 5, 4, 3 + conv = torch.nn.ConvTranspose1d(c_in, c_out, kernel_size, stride=stride) + x = torch.randn(bsz, c_in, t) + expected = conv(x)[..., : t * stride] # causal: drop the kernel - stride right tail + + z = (x.transpose(1, 2).reshape(bsz * t, c_in) @ tmc.transposed_gemm_weight(conv.weight)).view( + bsz, t, kernel_size, c_out + ) + actual = tmc.transposed_overlap_add(z, stride, conv.bias) + torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5) + + +def test_transposed_overlap_add_rejects_other_kernels(): + with pytest.raises(ValueError, match="kernel stride or 2"): + tmc.transposed_overlap_add(torch.zeros(1, 2, 5, 3), 2, torch.zeros(3)) + + +@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 3)]) +@pytest.mark.parametrize("with_residual", [False, True]) +def test_causal_conv_matches_conv1d(kernel_size, dilation, with_residual): + torch.manual_seed(0) + bsz, t, c_in, c_out = 2, 13, 6, 5 + conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation) + x = torch.randn(bsz, c_in, t) + residual = torch.randn(bsz, t, c_out) if with_residual else None + expected = conv(F.pad(x, ((kernel_size - 1) * dilation, 0))).transpose(1, 2) + if residual is not None: + expected = expected + residual + actual = tmc.causal_conv( + x.transpose(1, 2).contiguous(), tmc.conv_gemm_weight(conv.weight), conv.bias, kernel_size, dilation, residual + ) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def _gpu_tolerances(dtype): + # bf16 results may differ by one ulp (sin rounding, fp32 vs bf16 accumulation order). + return {"rtol": 1.6e-2, "atol": 1e-2} if dtype == torch.bfloat16 else {"rtol": 1e-4, "atol": 1e-4} + + +@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_triton_kernels_match_reference(dtype): + torch.manual_seed(0) + tol = _gpu_tolerances(dtype) + x = torch.randn(3, 37, 96, device="cuda", dtype=dtype) + bias = torch.randn(96, device="cuda", dtype=dtype) + snake = (torch.rand(96, device="cuda", dtype=dtype) + 0.5, torch.rand(96, device="cuda", dtype=dtype) + 0.5) + for kernel_size, dilation in ((1, 1), (7, 9)): + torch.testing.assert_close( + tmc.snake_im2col(x, kernel_size, dilation, bias, snake), + tmc.snake_im2col_reference(x, kernel_size, dilation, bias, snake), + **tol, + ) + z = torch.randn(3, 37, 10, 48, device="cuda", dtype=dtype) + torch.testing.assert_close( + tmc.transposed_overlap_add(z, 5, bias[:48]), tmc.transposed_overlap_add_reference(z, 5, bias[:48]), **tol + ) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize(("c_in", "c_out", "kernel_size", "dilation"), [(96, 96, 7, 9), (192, 96, 1, 1), (96, 1, 7, 1)]) +def test_triton_causal_conv_matches_reference(dtype, c_in, c_out, kernel_size, dilation): + torch.manual_seed(0) + x = torch.randn(3, 101, c_in, device="cuda", dtype=dtype) + weight = torch.randn(kernel_size * c_in, c_out, device="cuda", dtype=dtype) / (kernel_size * c_in) ** 0.5 + bias = torch.randn(c_out, device="cuda", dtype=dtype) + residual = torch.randn(3, 101, c_out, device="cuda", dtype=dtype) + expected = tmc.causal_conv_reference( + x.float(), weight.float(), bias.float(), kernel_size, dilation, residual.float() + ) + actual = tmc.causal_conv(x, weight, bias, kernel_size, dilation, residual) + torch.testing.assert_close(actual.float(), expected, **_gpu_tolerances(dtype)) diff --git a/tests/model_executor/models/qwen3_tts/test_time_major_decoder.py b/tests/model_executor/models/qwen3_tts/test_time_major_decoder.py new file mode 100644 index 000000000..8097c957c --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_time_major_decoder.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2DecoderConfig, +) +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Decoder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_decoder() -> Qwen3TTSTokenizerV2Decoder: + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + torch.manual_seed(0) + decoder = Qwen3TTSTokenizerV2Decoder(config).eval() + with torch.no_grad(): + for name, param in decoder.named_parameters(): + # Non-trivial SnakeBeta and ConvNeXt scales so every term matters. + if name.endswith(("alpha", "beta")): + param.uniform_(-0.5, 0.5) + elif name.endswith("gamma"): + param.uniform_(0.5, 1.0) + decoder.precompute_snake_caches() + return decoder + + +@torch.inference_mode() +def test_time_major_conv_decode_matches_nct_stack() -> None: + decoder = _make_decoder() + hidden = torch.randn(2, 5, decoder.config.latent_dim) + expected = decoder._conv_decode(hidden) + + decoder.enable_time_major_conv() + actual = decoder._conv_decode(hidden) + + assert actual.shape == expected.shape == (2, 1, 5 * int(decoder.total_upsample)) + scale = expected.abs().max() + torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5) + # The stack holds weight copies outside the module tree. + assert not any("time_major" in key for key in decoder.state_dict()) + + +@torch.inference_mode() +def test_time_major_conv_decode_takes_sliced_hidden() -> None: + decoder = _make_decoder() + hidden = torch.randn(1, 9, decoder.config.latent_dim) + expected = decoder._conv_decode(hidden[:, -4:, :]) + decoder.enable_time_major_conv() + actual = decoder._conv_decode(hidden[:, -4:, :]) # non-contiguous, like the suffix decode's slice + scale = expected.abs().max() + torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5) + + +@torch.inference_mode() +def test_first_frame_single_key_attention_matches_codec_without_graphs(): + from vllm_omni.model_executor.models.qwen3_tts.first_frame_decoder import Qwen3TTSFirstFrameDecoder + + decoder = _make_decoder() + codes = torch.randint(0, 32, (3, 2, 1)) + expected = decoder._decode_xvec_first_chunk(codes, {})[:, 0, :] + first = Qwen3TTSFirstFrameDecoder.__new__(Qwen3TTSFirstFrameDecoder) + torch.nn.Module.__init__(first) + first.decoder = decoder + first._graphs = {} + actual = first.decode(codes[:, :, 0]) + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5) diff --git a/tests/worker_v2/test_first_audio_sender.py b/tests/worker_v2/test_first_audio_sender.py new file mode 100644 index 000000000..181f01d8b --- /dev/null +++ b/tests/worker_v2/test_first_audio_sender.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import queue +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.engine import OmniEngineCoreOutputs +from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _scheduler(**client_by_request): + return SimpleNamespace( + requests={request_id: SimpleNamespace(client_index=client) for request_id, client in client_by_request.items()} + ) + + +def test_sink_groups_outputs_by_client_and_carries_audio(): + output_queue: queue.Queue = queue.Queue() + sink = engine_output_queue_sink(output_queue, _scheduler(a=0, b=1, c=0)) + sample_rate = torch.tensor(24000, dtype=torch.int32) + rows = [torch.full((4,), float(i)) for i in range(3)] + + sink.prepare(["a", "b", "c"])(["a", "b", "c"], rows, sample_rate) + + delivered = {} + while not output_queue.empty(): + client_index, outputs = output_queue.get_nowait() + assert isinstance(outputs, OmniEngineCoreOutputs) + delivered[client_index] = outputs.outputs + assert sorted(delivered) == [0, 1] + assert [o.request_id for o in delivered[0]] == ["a", "c"] + assert [o.request_id for o in delivered[1]] == ["b"] + for output in delivered[0] + delivered[1]: + # First audio carries no tokens and does not finish the request. + assert output.new_token_ids == [] + assert output.finish_reason is None + assert output.multimodal_output["sr"] is sample_rate + assert torch.equal(delivered[0][1].multimodal_output["model_outputs"], rows[2]) + + +def test_sink_skips_requests_that_already_left_the_scheduler(): + output_queue: queue.Queue = queue.Queue() + sink = engine_output_queue_sink(output_queue, _scheduler(kept=0)) + + sink.prepare(["gone", "kept"])(["gone", "kept"], [torch.zeros(2), torch.ones(2)], torch.tensor(24000)) + sink.prepare(["gone"])(["gone"], [torch.zeros(2)], torch.tensor(24000)) + + client_index, outputs = output_queue.get_nowait() + assert client_index == 0 + assert [o.request_id for o in outputs.outputs] == ["kept"] + assert output_queue.empty() + + +class _DoneEvent: + def record(self): + pass + + def synchronize(self): + pass + + +def _run_sender(items): + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + output_queue: queue.Queue = queue.Queue() + sink = engine_output_queue_sink(output_queue, _scheduler(**{rid: 0 for item in items for rid in item[2]})) + sender = FirstAudioSender(sink) + for item in items: + sender._queue.put((*item, sink.prepare(item[2]))) + sender.close() + delivered = [] + while not output_queue.empty(): + _, outputs = output_queue.get_nowait() + delivered.append( + ([o.request_id for o in outputs.outputs], [o.multimodal_output["model_outputs"] for o in outputs.outputs]) + ) + return delivered + + +def test_sender_drops_rows_whose_frame_is_not_audio(): + pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2) + sample_rate = torch.tensor(24000, dtype=torch.int32) + valid = torch.tensor([True, False, True]) + + delivered = _run_sender( + [ + (_DoneEvent(), pcm, ["a", "b", "c"], sample_rate, valid), + # Every row invalid: nothing reaches the sink. + (_DoneEvent(), pcm[:1], ["d"], sample_rate, torch.tensor([False])), + # No validity: every row is delivered (previous behaviour). + (_DoneEvent(), pcm[:1], ["e"], sample_rate, None), + ] + ) + + assert [ids for ids, _rows in delivered] == [["a", "c"], ["e"]] + assert torch.equal(delivered[0][1][1], pcm[2]) + + +def test_prepared_delivery_survives_scheduler_completion(): + output_queue: queue.Queue = queue.Queue() + scheduler = _scheduler(r=3) + sink = engine_output_queue_sink(output_queue, scheduler) + deliver = sink.prepare(["r"]) + scheduler.requests.clear() + deliver(["r"], [torch.ones(2)], torch.tensor(24000)) + client, outputs = output_queue.get_nowait() + assert client == 3 + assert outputs.outputs[0].multimodal_output["_omni_first_audio"] + + +def test_delivery_failure_emits_request_error(): + from vllm.v1.engine import FinishReason + + output_queue: queue.Queue = queue.Queue() + delivery = engine_output_queue_sink(output_queue, _scheduler(r=0)).prepare(["r"]) + delivery.fail(["r"]) + _, outputs = output_queue.get_nowait() + assert outputs.outputs[0].finish_reason == FinishReason.ERROR + + +@pytest.mark.parametrize("client_index", [0, 3]) +def test_submit_reports_only_requests_with_a_prepared_route(monkeypatch, client_index): + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + # Exercise the routing contract without allocating CUDA or pinned memory. + original_empty = torch.empty + monkeypatch.setattr(torch, "empty", lambda *args, pin_memory=False, **kwargs: original_empty(*args, **kwargs)) + monkeypatch.setattr(torch.cuda, "Event", _DoneEvent) + output_queue: queue.Queue = queue.Queue() + sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler(kept=client_index))) + pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2) + try: + accepted = sender.submit(["gone", "kept", "also-gone"], pcm, torch.tensor(24000)) + finally: + sender.close() + + assert accepted == ["kept"] + client, outputs = output_queue.get_nowait() + assert client == client_index + assert [output.request_id for output in outputs.outputs] == ["kept"] + assert torch.equal(outputs.outputs[0].multimodal_output["model_outputs"], pcm[1]) + assert output_queue.empty() + + +def test_submit_without_any_route_leaves_audio_to_the_codec(monkeypatch): + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + def unexpected_copy(*args, **kwargs): + pytest.fail("No delivery is possible; do not start a device copy") + + monkeypatch.setattr(torch, "empty", unexpected_copy) + monkeypatch.setattr(torch.cuda, "Event", unexpected_copy) + output_queue: queue.Queue = queue.Queue() + sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler())) + try: + accepted = sender.submit(["gone"], torch.ones(1, 2), torch.tensor(24000)) + finally: + sender.close() + + assert accepted == [] + assert output_queue.empty() + + +@pytest.mark.parametrize("sink", [lambda ids, rows, sr: None, SimpleNamespace(prepare=None)]) +def test_sender_rejects_sink_without_route_preparation(sink): + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + with pytest.raises(TypeError, match="prepare"): + FirstAudioSender(sink) + + +def test_sender_failure_only_aborts_valid_undelivered_rows(): + from vllm.v1.engine import FinishReason + + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + class FailSecondClientOnce(queue.Queue): + def put_nowait(self, item): + client, outputs = item + if client == 1 and outputs.outputs[0].finish_reason is None: + raise RuntimeError("client enqueue failed") + super().put_nowait(item) + + output_queue = FailSecondClientOnce() + sink = engine_output_queue_sink(output_queue, _scheduler(eos=1, sent=0, failed=1, pending=2, next=3)) + sender = FirstAudioSender(sink) + ids = ["eos", "sent", "failed", "pending", "gone"] + sender._queue.put( + ( + _DoneEvent(), + torch.ones(5, 2), + ids, + torch.tensor(24000), + torch.tensor([False, True, True, True, True]), + sink.prepare(ids), + ) + ) + sender._queue.put((_DoneEvent(), torch.ones(1, 2), ["next"], torch.tensor(24000), None, sink.prepare(["next"]))) + sender.close() + + actual: list[tuple[int, str, FinishReason | None]] = [] + while not output_queue.empty(): + client, outputs = output_queue.get_nowait() + actual.extend((client, o.request_id, o.finish_reason) for o in outputs.outputs) + assert actual == [ + (0, "sent", None), + (1, "failed", FinishReason.ERROR), + (2, "pending", FinishReason.ERROR), + (3, "next", None), + ] + + +def test_sender_copy_failure_reports_all_prepared_routes(): + from vllm.v1.engine import FinishReason + + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + class FailedEvent: + def synchronize(self): + raise RuntimeError("copy failed before validity became readable") + + output_queue: queue.Queue = queue.Queue() + sink = engine_output_queue_sink(output_queue, _scheduler(r=3)) + sender = FirstAudioSender(sink) + sender._queue.put( + (FailedEvent(), torch.ones(2, 2), ["r", "gone"], torch.tensor(24000), None, sink.prepare(["r", "gone"])) + ) + sender.close() + client, outputs = output_queue.get_nowait() + assert client == 3 + assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)] + assert output_queue.empty() diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 1124f08ca..3b585527a 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -391,3 +391,204 @@ def test_mtp_requires_model_declared_output_key(output_key): model = SimpleNamespace(mtp=lambda: None, mtp_output_key=output_key) with pytest.raises(TypeError, match="must declare mtp_output_key"): state._init_mtp_runner(model) + + +_EAGER_DIM = 3 +_CODEBOOK = 2048 +_EOS = 2150 + + +class _EagerBatch(_DummyInputBatch): + def __init__(self, spans, indices=None): + indices = list(range(len(spans))) if indices is None else indices + super().__init__(indices) + starts = [0] + for n_tok in spans: + starts.append(starts[-1] + n_tok) + self.num_scheduled_tokens = list(spans) + self.query_start_loc_np = starts[:-1] + self.query_start_loc = torch.tensor(starts, dtype=torch.int32) + + +def _make_eager_state(max_num_reqs=4): + state = _make_state(max_num_reqs=max_num_reqs, has_preprocess=True) + _init_static(state, max_num_reqs, dim=_EAGER_DIM) + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + + state._eager_state = EagerMTPState(state) + state._first_audio_requests = set() + state._eager_mtp = True + state._eager_fastpath = False + state._eager_rows = None + state._eager_ready = {} + state._eager_settled = {} + state._eager_embeds = torch.zeros((max_num_reqs, _EAGER_DIM)) + state._decode_preprocess = None + state._mtp_sample_uniforms = None + state.vllm_config.cache_config = SimpleNamespace(enable_prefix_caching=False) + model = state.model + model.first_frame_decoder = None + model.mtp_frame_valid = lambda layer0: (layer0 >= 0) & (layer0 < _CODEBOOK) + model.embed_input_ids = lambda ids: ids.float().reshape(-1, 1, 1).expand(-1, 1, _EAGER_DIM) + model.mtp_calls = [] + + def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs): + model.mtp_calls.append((input_ids.clone(), last_hidden.clone(), text_step.clone())) + codes = torch.stack([input_ids, input_ids + 1, input_ids + 2], dim=1) + # Frame embedding sum plus the text step, like the Qwen3 talkers. + return input_embeds + 100 + text_step, codes + + model.mtp = mtp + return state + + +def _eager_outputs(num_tokens): + return { + "codes": {"audio": torch.zeros((num_tokens, 3), dtype=torch.long)}, + "meta": {"codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8)}, + } + + +def test_eager_mtp_publishes_frames_in_the_step_that_sampled_cb0(): + state = _make_eager_state() + _fill_buffers(state, "prefill", "decode", "ended") + # prefill: final 3-token chunk; decode: input CB0 42; ended: input CB0 was + # EOS (async scheduling ran the finished request once more). + batch = _EagerBatch([3, 1, 1]) + input_ids = torch.tensor([0, 0, 0, 42, _EOS]) + state._eager_rows = (batch, [(0, 0, "prefill", True), (1, 1, "decode", False), (2, 2, "ended", False)], input_ids) + text_hidden = torch.arange(5 * _EAGER_DIM, dtype=torch.float32).reshape(5, _EAGER_DIM) + sampled = torch.tensor([[7], [_EOS], [9]]) + outputs = _eager_outputs(5) + + state.run_eager_mtp(batch, text_hidden, sampled, outputs) + + ids, hidden, text_step = state.model.mtp_calls[0] + assert ids.tolist() == [7, _EOS, 9] + # The frame's hidden is this step's last-token hidden of each span. + assert torch.equal(hidden, text_hidden[[2, 3, 4]]) + # The next step adds its own text step. + assert not text_step.any() + codes = outputs["codes"]["audio"] + assert codes[2].tolist() == [7, 8, 9] + assert codes[3].tolist() == [_EOS, _EOS + 1, _EOS + 2] + assert codes[4].tolist() == [9, 10, 11] + assert not codes[:2].any() + # Valid only when both the sampled and (for decode) the input CB0 are codec ids. + assert outputs["meta"]["codec_frame_valid"].tolist() == [0, 0, 1, 0, 0] + assert state._eager_ready == {0: "prefill", 1: "decode", 2: "ended"} + assert torch.equal(state._eager_embeds[0], torch.full((_EAGER_DIM,), 107.0)) + assert state._eager_rows is None + + +def test_eager_mtp_ignores_rows_recorded_for_another_batch(): + state = _make_eager_state() + _fill_buffers(state, "r0") + state._eager_rows = (_EagerBatch([1]), [(0, 0, "r0", False)], torch.tensor([5])) + outputs = _eager_outputs(1) + state.run_eager_mtp(_EagerBatch([1]), torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), outputs) + assert state.model.mtp_calls == [] + assert not outputs["codes"]["audio"].any() + + +def test_eager_mtp_requires_retained_frame_outputs(): + state = _make_eager_state() + _fill_buffers(state, "r0") + batch = _EagerBatch([1]) + state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5])) + with pytest.raises(RuntimeError, match="requires retained codes.audio"): + state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), {"codes": {"audio": None}}) + + +def test_eager_decode_input_is_previous_frame_plus_text_step(): + state = _make_eager_state() + _fill_buffers(state, "r0", "r1") + state._eager_ready = {0: "r0", 1: "r1"} + state._eager_embeds[0] = 5.0 + state._eager_embeds[1] = 7.0 + embeds = torch.zeros((2, _EAGER_DIM)) + batches = [ + (0, 0, (torch.zeros(_EAGER_DIM), torch.full((_EAGER_DIM,), 2.0))), + (1, 1, (None, torch.ones(_EAGER_DIM))), + ] + state._eager_state._apply_eager_frames(batches, embeds, _EagerBatch([1, 1]), None) + assert embeds.tolist() == [[7.0] * _EAGER_DIM, [8.0] * _EAGER_DIM] + # The deferred MTP must not run for an eager decode row. + assert state.model.mtp_calls == [] + + +def test_eager_decode_without_a_frame_fails_loudly(): + state = _make_eager_state() + _fill_buffers(state, "r0") + batches = [(0, 0, (torch.zeros(_EAGER_DIM), torch.zeros(_EAGER_DIM)))] + with pytest.raises(RuntimeError, match="frame missing"): + state._eager_state._apply_eager_frames(batches, torch.zeros((1, _EAGER_DIM)), _EagerBatch([1]), None) + + +def test_run_preprocess_records_rows_that_keep_a_sample(): + state = _make_eager_state() + _fill_buffers(state, "chunk", "final", "decode") + state._eager_ready = {2: "decode"} + state._eager_embeds[2] = 4.0 + text_step = torch.ones(_EAGER_DIM) + + def preprocess(input_ids, input_embeds, **info): + updates = {"mtp_inputs": (torch.zeros(_EAGER_DIM), text_step)} if input_ids.shape[0] == 1 else {} + return input_ids, input_embeds, updates + + state.model.preprocess = preprocess + batch = _EagerBatch([3, 3, 1]) + model_inputs = {"input_ids": torch.zeros(7, dtype=torch.long), "inputs_embeds": torch.zeros((7, _EAGER_DIM))} + req_states = SimpleNamespace( + prompt_len=np.array([10, 10, 10], dtype=np.int32), + num_computed_tokens=np.array([0, 7, 10], dtype=np.int32), + ) + state.run_preprocess(batch, model_inputs, req_states) + + recorded_batch, entries, _ids = state._eager_rows + assert recorded_batch is batch + # A non-final prefill chunk samples nothing that is kept. + assert entries == [(1, 1, "final", True), (2, 2, "decode", False)] + assert model_inputs["inputs_embeds"][6].tolist() == [5.0] * _EAGER_DIM + assert state.model.mtp_calls == [] + + +@pytest.mark.parametrize("first_was_valid", [False, True]) +def test_first_audio_requirement_survives_eos_only_if_audio_was_queued(first_was_valid): + state = _make_eager_state() + _fill_buffers(state, "r0") + state._first_audio_requests.add("r0") + state._eager_state._first_audio_valid = torch.tensor([first_was_valid, False, False, False]) + batch = _EagerBatch([1]) + state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5])) + outputs = _eager_outputs(1) + outputs["meta"]["first_audio"] = torch.zeros(1, dtype=torch.bool) + state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[_EOS]]), outputs) + assert outputs["meta"]["first_audio"].tolist() == [first_was_valid] + + +@pytest.mark.parametrize("accepted", [[], ["r1"], ["r0", "r1"]]) +def test_first_audio_marker_requires_accepted_delivery(monkeypatch, accepted): + from contextlib import nullcontext + + state = _make_eager_state() + _fill_buffers(state, "r0", "r1") + state.model.first_frame_decoder = SimpleNamespace( + sample_rate=24000, decode=lambda codes: torch.ones(codes.shape[0], 2) + ) + state._first_audio_stream = SimpleNamespace(wait_stream=lambda stream: None) + state._first_audio_sender = SimpleNamespace(submit=lambda ids, pcm, sr, valid: accepted) + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: None) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch.Tensor, "record_stream", lambda tensor, stream: None) + batch = _EagerBatch([1, 1]) + state._eager_rows = (batch, [(0, 0, "r0", True), (1, 1, "r1", True)], torch.zeros(2, dtype=torch.long)) + outputs = _eager_outputs(2) + outputs["meta"]["first_audio"] = torch.zeros(2, dtype=torch.bool) + + state.run_eager_mtp(batch, torch.zeros(2, _EAGER_DIM), torch.tensor([[7], [8]]), outputs) + + assert state._first_audio_requests == set(accepted) + # A missing route must leave the normal codec path responsible for frame 0; + # otherwise it skips the frame and the orchestrator waits forever for it. + assert outputs["meta"]["first_audio"].tolist() == ["r0" in accepted, "r1" in accepted] diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 048e340f3..db62d4534 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -112,6 +112,8 @@ _INHERITED_VLLM = ( _PLATFORM_EXTERNAL = ( "ASCEND_LAUNCH_BLOCKING", "ASCEND_RT_VISIBLE_DEVICES", + "CUDA_MPS_LOG_DIRECTORY", + "CUDA_MPS_PIPE_DIRECTORY", "FLASHINFER_DISABLE_VERSION_CHECK", "HF_HOME", "HF_HUB_DISABLE_XET", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 578aa8db1..3a650dd03 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -629,6 +629,8 @@ class OmniStageConnectorConfig: class OmniStageRuntimeConfig: """Per-stage process placement and backend runtime behavior.""" + cuda_mps: bool = False + # LLM backend extensions; diffusion owns these in its config projection. additional_config: dict[str, Any] | None = None distributed_executor_backend: Any = None @@ -2127,6 +2129,7 @@ def _build_runtime_config( kwargs["num_replicas"] = stage_deploy.num_replicas if "env" not in kwargs and stage_deploy is not None and stage_deploy.env is not None: kwargs["env"] = _copy_value(stage_deploy.env) + kwargs["cuda_mps"] = deploy.cuda_mps kwargs["num_gpus"] = parallel_config.world_size return OmniStageRuntimeConfig(**kwargs) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 42da6cfa1..00ebda453 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -582,6 +582,8 @@ class DeployConfig: model_runner: Literal["v1", "v2"] = "v1" # Stage-1 active stream slots; 0 preserves legacy all-stream cycling. active_stream_window: int = 0 + # Experimental local NVIDIA MPS; disabled unless a deploy explicitly opts in. + cuda_mps: bool = False duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig) connectors: dict[str, Any] | None = None edges: list[dict[str, Any]] | None = None @@ -804,7 +806,11 @@ def load_deploy_config(path: str | Path) -> DeployConfig: if model_runner not in ("v1", "v2"): raise ValueError(f"model_runner must be one of ('v1', 'v2'), got {model_runner!r}") + if not isinstance(raw_dict.get("cuda_mps", False), bool): + raise ValueError("cuda_mps must be a boolean") + kwargs: dict[str, Any] = { + "cuda_mps": raw_dict.get("cuda_mps", False), "async_chunk": raw_dict.get("async_chunk", True), "session_mode": raw_dict.get("session_mode", "turn"), "model_runner": model_runner, @@ -936,7 +942,11 @@ def _resolve_execution_mode( def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool: - return bool(deploy.async_chunk and (stage is None or stage.async_chunk is not False)) + if not isinstance(deploy.async_chunk, bool): + raise ValueError("async_chunk must be a boolean") + if stage is not None and stage.async_chunk is not None and not isinstance(stage.async_chunk, bool): + raise ValueError(f"Stage {stage.stage_id} async_chunk must be a boolean or null") + return deploy.async_chunk and (stage is None or stage.async_chunk is not False) def validate_stage_async_chunk_edges(pipeline: PipelineConfig, deploy: DeployConfig) -> None: @@ -1178,6 +1188,8 @@ def merge_pipeline_deploy( runtime["num_replicas"] = ds.num_replicas if ds.env is not None: runtime["env"] = ds.env + if deploy.cuda_mps: + runtime["cuda_mps"] = True runtime["requires_multimodal_data"] = ps.requires_multimodal_data result.append( diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 52c27f20e..e7c27fa70 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -19,6 +19,10 @@ import msgspec import numpy as np import torch +# Internal output routing markers shared by first-frame producers and orchestration. +FIRST_AUDIO_KEY = "_omni_first_audio" +FIRST_AUDIO_REQUIRED_KEY = "_omni_first_audio_required" + REQUEST_ARTIFACT_DIRS_KEY = "_omni_request_artifact_dirs" TRANSFORM_OWNED_META_KEYS = frozenset({"minimax_h3_prepared_reference_videos"}) @@ -101,6 +105,7 @@ class OmniPayloadMeta(TypedDict, total=False): width: int decode_flag: bool codec_streaming: bool + first_audio: bool | torch.Tensor codec_frame_valid: bool | torch.Tensor ref_code_len: int ref_context_size: int @@ -214,6 +219,7 @@ class MetaStruct(_StructBase): width: int | None = None decode_flag: bool | None = None codec_streaming: bool | None = None + first_audio: torch.Tensor | None = None codec_frame_valid: torch.Tensor | None = None ref_code_len: int | None = None # Expected FINAL length of a growing async-chunk sequence, when the diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index f3ebe80f6..9883c2662 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -24,6 +24,7 @@ connectors: connector_of_shared_memory: name: SharedMemoryConnector extra: + talker_first_audio: true codec_streaming: true connector_get_sleep_s: 0.01 connector_get_max_wait_first_chunk: 3000 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 36c85e41d..4612accf2 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -24,6 +24,7 @@ connectors: connector_of_shared_memory: name: SharedMemoryConnector extra: + talker_first_audio: true codec_streaming: true connector_get_sleep_s: 0.01 connector_get_max_wait_first_chunk: 3000 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml index b6ebcf67f..cb5dc91f6 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml @@ -15,6 +15,7 @@ platforms: connectors: connector_of_shared_memory: extra: + talker_first_audio: true # Keep unqualified prefix-graph optimizations disabled in this profile. code_predictor_prefix_graphs: false ref_audio_artifact_cache_max_entries: 1024 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml index da6a90cf4..05be5ff0b 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml @@ -1,15 +1,36 @@ -# Opt-in single-GPU throughput profile, measured on H200 at concurrency 128 -# with Qwen3-TTS-12Hz-1.7B-Base. Warm the full pipeline before benchmarking. -# B8 uses more graph memory than B4; this is not a low-latency profile. -# MPS is optional and is not enabled by this file. +# Experimental single-GPU CUDA/BF16 throughput profile. +# Optimizations: frame-local predictor KV, fused sampling, time-major codec +# convolutions, eager residual prediction and first-frame delivery; B1-B8 codec +# graphs, FULL Talker graphs and 128 Talker slots. Floating-point results differ: +# validate audio quality for your workload. Ordinary profiles do not enable MPS. +# cuda_mps starts one private NVIDIA MPS daemon shared by this pipeline's stages +# and stops it after they exit. An existing CUDA_MPS_PIPE_DIRECTORY is reused +# without taking ownership. Requires nvidia-cuda-mps-control on PATH. +# Set cuda_mps: false in an overlay to disable it. +# MPS can increase throughput under load but also increase first-packet and +# first-audible-audio latency. Compare cuda_mps: false for latency-sensitive +# workloads, especially at low request rates. +# Warm the full pipeline before measuring latency or throughput. # -# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \ -# --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml +# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ +# --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml \ +# --api-server-count 4 --stage-init-timeout 1800 --init-timeout 2400 base_config: qwen3_tts_high_concurrency_mrv2.yaml +cuda_mps: true connectors: connector_of_shared_memory: extra: + talker_first_audio: true + codec_chunk_ramp: null + codec_chunk_adaptive: false + initial_codec_chunk_frames: 1 + codec_chunk_frames: 25 + codec_left_context_frames: 72 + code_predictor_kv_cache: true + code_predictor_fused_sampling: true + decode_time_major_conv: true + decode_cudnn_benchmark: true decode_batch_max_size: 8 decode_cudagraph_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] @@ -17,7 +38,11 @@ stages: - stage_id: 0 devices: "0" max_num_seqs: 128 + compilation_config: + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512] + max_cudagraph_capture_size: 512 - stage_id: 1 devices: "0" - # Admission candidates are grouped by codec phase before B8 execution. - max_num_seqs: 32 + max_num_seqs: 64 + async_scheduling: false diff --git a/vllm_omni/deploy/qwen3_tts_mrv2.yaml b/vllm_omni/deploy/qwen3_tts_mrv2.yaml index 180d55b5c..b753c83b9 100644 --- a/vllm_omni/deploy/qwen3_tts_mrv2.yaml +++ b/vllm_omni/deploy/qwen3_tts_mrv2.yaml @@ -18,3 +18,8 @@ stages: - stage_id: 0 # Bound per-step prefill work, matching the high-concurrency TTS profile. max_num_batched_tokens: 512 + +connectors: + connector_of_shared_memory: + extra: + talker_first_audio: true diff --git a/vllm_omni/engine/cuda_mps.py b/vllm_omni/engine/cuda_mps.py new file mode 100644 index 000000000..3c9d575e1 --- /dev/null +++ b/vllm_omni/engine/cuda_mps.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Private NVIDIA MPS lifetime for explicitly opted-in local stage processes.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import tempfile +from pathlib import Path + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def physical_gpu_uuid(device: str) -> str: + """Resolve a physical ordinal/UUID through NVML without initializing CUDA.""" + from vllm.utils.import_utils import import_pynvml + + nvml = import_pynvml() + nvml.nvmlInit() + try: + if device.isdigit(): + handle = nvml.nvmlDeviceGetHandleByIndex(int(device)) + elif device.startswith("GPU-"): + handle = nvml.nvmlDeviceGetHandleByUUID(device) + else: + raise ValueError("cuda_mps requires one physical NVIDIA GPU ordinal or GPU UUID") + uuid = nvml.nvmlDeviceGetUUID(handle) + return uuid.decode() if isinstance(uuid, bytes) else str(uuid) + finally: + nvml.nvmlShutdown() + + +class CudaMPSServer: + """Own a private control socket; never stop an operator's existing daemon.""" + + def __init__(self, gpu_uuid: str, *, pipe_directory: str | None = None) -> None: + control = shutil.which("nvidia-cuda-mps-control") + if control is None: + raise RuntimeError("cuda_mps requires nvidia-cuda-mps-control on PATH") + self._control = control + self._directory: Path | None = None + self._closed = False + # None inherits the parent; an explicit empty stage value selects a + # private daemon even when the parent names an operator-managed one. + self.operator_pipe_directory = ( + os.environ.get("CUDA_MPS_PIPE_DIRECTORY") if pipe_directory is None else pipe_directory + ) or None + if self.operator_pipe_directory: + self.env = {"CUDA_MPS_PIPE_DIRECTORY": self.operator_pipe_directory, "CUDA_VISIBLE_DEVICES": gpu_uuid} + self._run(input="get_server_list\n") + logger.info("Using operator-managed MPS at %s for %s", self.operator_pipe_directory, gpu_uuid) + return + + self._directory = root = Path(tempfile.mkdtemp(prefix="vllm-omni-mps-")) + (root / "pipe").mkdir() + (root / "log").mkdir() + self.env = { + "CUDA_VISIBLE_DEVICES": gpu_uuid, + "CUDA_MPS_PIPE_DIRECTORY": str(root / "pipe"), + "CUDA_MPS_LOG_DIRECTORY": str(root / "log"), + } + try: + self._run("-d") + self._run(input="get_server_list\n") + except BaseException: + # A failed/timed-out launcher may already have forked the daemon. + # Always address our private socket; keep it if cleanup also fails. + try: + self.close() + except Exception: + logger.exception("MPS startup cleanup failed; control files remain at %s", root) + raise + logger.info("Started private MPS for %s at %s", gpu_uuid, self.env["CUDA_MPS_PIPE_DIRECTORY"]) + + def _run(self, *args: str, input: str | None = None) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [self._control, *args], + input=input, + env=os.environ | self.env, + check=True, + text=True, + capture_output=True, + timeout=15, + ) + + def close(self) -> None: + if self._closed: + return + if self._directory is not None: + # Keep the control socket if quit fails, so cleanup can be retried. + self._run(input="quit\n") + shutil.rmtree(self._directory) + self._directory = None + self._closed = True diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index c5a242732..64eee1cff 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -21,7 +21,7 @@ import time as _time from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass, field from http import HTTPStatus -from typing import Any +from typing import TYPE_CHECKING, Any import janus import torch @@ -34,6 +34,7 @@ from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats +from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest @@ -66,6 +67,9 @@ from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput +if TYPE_CHECKING: + from vllm_omni.engine import OmniEngineCoreOutput + logger = init_logger(__name__) @@ -226,6 +230,10 @@ class OrchestratorRequestState: running_counter_registered: bool = False request_artifact_dirs: set[str] = field(default_factory=set) native_kv_transfer_id: str | None = None + # First-frame delivery and final-stage outputs use different channels. + upstream_first_audio: bool = False + pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None + pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list) @dataclass @@ -272,7 +280,7 @@ class StreamingInputState: Fresh rather than shared: ``output_metadata`` is handed out by reference. """ - segment = self.segments.get(stage_id) + segment = self.segments.get(stage_id) if stage_id is not None else None return segment if segment is not None else StreamingSegmentState() @@ -440,9 +448,13 @@ class OrchestratorBase: # Start membership watcher if distributed mode is active. membership_watcher: asyncio.Task[None] | None = None if self._membership is not None: + + async def cleanup_member_requests(ids: list[str]) -> None: + await self._cleanup_request_ids(ids, abort=True) + self._membership.install_unregister_handlers( output_queue=self.output_async_queue, - cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), + cleanup_callback=cleanup_member_requests, replica_removed_callback=self._remove_stage_replica_waiting, ) membership_watcher = self._membership.start() @@ -528,7 +540,7 @@ class OrchestratorBase: # process exit as EngineDeadError during teardown. for pool in self.stage_pools: for client in pool.clients: - if hasattr(client, "_shutting_down"): + if client is not None and hasattr(client, "_shutting_down"): client._shutting_down = True # Stage teardown runs once in run()'s finally after the # orchestration loop observes _shutdown_event and exits. @@ -786,6 +798,7 @@ class OrchestratorBase: ``_finish_raw_terminal_requests`` once routing is done. """ pool = self.stage_pools[stage_id] + await self._route_upstream_first_audio(stage_id, replica_id, raw_outputs) await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) for eco in raw_outputs.outputs: # Emit kv_wait_s before _handle_kv_ready_raw_outputs' @@ -800,7 +813,8 @@ class OrchestratorBase: kv_params.get("connector_type") or "unknown", float(kv_wait_s), ) - req_state = self.request_states.get(getattr(eco, "request_id", None)) + request_id = getattr(eco, "request_id", None) + req_state = self.request_states.get(request_id) if request_id is not None else None if req_state is None: continue if ( @@ -861,6 +875,97 @@ class OrchestratorBase: ) return processed + async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_outputs: Any) -> None: + """Order explicitly marked first audio ahead of a codec that skipped it. + + The sender and codec run independently. Keep overtaking codec outputs, + including terminal outputs, on the request until its first PCM arrives. + Cancellation and stage errors use ordinary request cleanup. + """ + kept = [] + final_output = self.stage_pools[stage_id].final_output + for eco in raw_outputs.outputs: + mm = getattr(eco, "multimodal_output", None) + is_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_KEY, False)) + requires_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_REQUIRED_KEY, False)) + request_id = getattr(eco, "request_id", None) + req_state = self.request_states.get(request_id) if request_id is not None else None + if is_first: + assert isinstance(mm, dict) + # Only the dedicated first-frame channel is intercepted. + if ( + req_state is None + or req_state.upstream_first_audio + or req_state.pending_upstream_first_audio is not None + or stage_id + 1 > req_state.final_stage_id + ): + continue + audio = mm.get("model_outputs") + if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + await self._handle_stage_error(stage_id, eco, error="Invalid or empty first-audio payload") + continue + req_state.pending_upstream_first_audio = (stage_id, eco) + await self._flush_upstream_first_audio(req_state) + continue + if ( + req_state is not None + and final_output + and not req_state.upstream_first_audio + and (requires_first or req_state.pending_first_audio_outputs) + and getattr(eco, "finish_reason", None) != FinishReason.ERROR + ): + req_state.pending_first_audio_outputs.append((stage_id, replica_id, eco)) + continue + kept.append(eco) + raw_outputs.outputs = kept + + async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) -> None: + """Feed upstream PCM through the codec's normal output-kind contract. + + The source can finish its first frame while downstream prewarm is + still selecting/registering a replica. Retain that frame until the + codec owns an output-processor state, then accumulate it before any + codec output that omitted it. This preserves FINAL_ONLY/CUMULATIVE + waveforms and downstream consumers such as forced alignment. + """ + from vllm_omni.engine import OmniEngineCoreOutputs + + pending = req_state.pending_upstream_first_audio + if pending is None or self.request_states.get(req_state.request_id) is not req_state: + return + source_stage, first_output = pending + codec_stage = source_stage + 1 + pool = self.stage_pools[codec_stage] + replica_id = pool.get_bound_replica_id(req_state.request_id) + if replica_id is None or req_state.request_id not in pool.output_processor.request_states: + return + req_state.pending_upstream_first_audio = None + req_state.upstream_first_audio = True + try: + processed = await pool.process_llm_raw_outputs(replica_id, OmniEngineCoreOutputs(outputs=[first_output])) + await self._handle_processed_outputs(codec_stage, replica_id, processed) + except EngineDeadError as error: + await self._handle_dead_replica(codec_stage, replica_id, error) + return + except Exception as error: + # A promised first frame must either enter the codec accumulator + # or fail its request; leaving it pending would stall completion. + logger.exception("First-audio output processing failed for %s", req_state.request_id) + await self._handle_stage_error(codec_stage, first_output, error=str(error)) + return + if self.request_states.get(req_state.request_id) is not req_state: + return + ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, [] + for pending_stage, pending_replica, output in ready: + if self.request_states.get(req_state.request_id) is not req_state: + break + terminal: set[str] = set() + processed = await self._process_llm_stage_outputs( + pending_stage, pending_replica, OmniEngineCoreOutputs(outputs=[output]), terminal + ) + await self._handle_processed_outputs(pending_stage, pending_replica, processed) + await self._finish_raw_terminal_requests(pending_stage, pending_replica, terminal) + def _absorb_diffusion_metrics(self, stage_id: int, replica_id: int, diffusion_output: Any) -> bool: """Drain a diffusion output's piggybacked metrics. @@ -1543,7 +1648,7 @@ class OrchestratorBase: self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) if req_state is not None: - cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) + cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", set())) if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: self._running_counter.decrement() req_state.running_counter_registered = False @@ -2201,7 +2306,7 @@ class OrchestratorBase: ) return diffusion_source_outputs = [output, *companion_outputs] - if next_client.custom_process_input_func is not None: + if next_client is not None and next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() _fn = next_client.custom_process_input_func _extra_kwargs: dict[str, Any] = {} @@ -2386,7 +2491,9 @@ class OrchestratorBase: req_state.streaming.source_token_decoder = decode try: - next_inputs = next_client.process_engine_inputs( + assert next_client is not None + process_inputs = getattr(next_client, "process_engine_inputs") + next_inputs = process_inputs( source_outputs, req_state.prompt, streaming_context=req_state.streaming, @@ -2637,6 +2744,11 @@ class OrchestratorBase: req_state, ) + if req_state.pending_upstream_first_audio is not None or req_state.upstream_first_audio: + await self._flush_upstream_first_audio(req_state) + if self.request_states.get(request_id) is not req_state: + return False + # async_chunk pre-submit fires per stage edge (N-1 -> N). Source # replica is stage 0's bound replica (single-replica thinker in # all current configs); fall back to 0 if unknown. diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index cc02ecd65..ecc9576bc 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -79,6 +79,23 @@ def _signal_exit_code(signum: int) -> int: return _SIGNAL_EXIT_BASE + signum +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.""" + if not isinstance(model_executor, UniProcExecutor): + 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) + model = getattr(model_runner, "model", None) + if getattr(model, "first_frame_decoder", None) is None or not hasattr(model_state, "set_first_audio_sink"): + return False + from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink + + assert model_state is not None + model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler)) + return True + + def _bind_native_data_plane_ready_sink(model_executor: Any, scheduler: Any) -> bool: """Bind the TP1 in-process runner control plane directly to its scheduler.""" if not isinstance(model_executor, UniProcExecutor): @@ -107,6 +124,7 @@ class StageEngineCoreProc(EngineCoreProc): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + _bind_first_audio_sink(self.model_executor, self.output_queue, self.scheduler) if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler): logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.") diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index cd38d0e2e..27d37bd59 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1041,14 +1041,17 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None, return previous_env: dict[str, str | None] = {} - for key, value in runtime_env.items(): - env_key = str(key) - previous_env[env_key] = os.environ.get(env_key) - os.environ[env_key] = str(value) - - if previous_env: - logger.info("[stage_init] Stage-%s applied runtime env keys: %s", stage_id, sorted(previous_env)) try: + for key, value in runtime_env.items(): + env_key = str(key) + old_value = os.environ.get(env_key) + os.environ[env_key] = str(value) + # Track only successful writes, preserving the original value + # if distinct keys normalize to the same environment name. + previous_env.setdefault(env_key, old_value) + + if previous_env: + logger.info("[stage_init] Stage-%s applied runtime env keys: %s", stage_id, sorted(previous_env)) yield finally: for key, old_value in previous_env.items(): @@ -1107,7 +1110,7 @@ def _project_omni_stage_engine_args( "default_sampling_params", "has_sampling_extra_args", } - runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"} + runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus", "cuda_mps"} if not is_diffusion: # These values configure OmniDiffusionConfig or its worker process; # OmniEngineArgs has no matching fields for LLM stages. diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 2f3e2e742..dc43575a5 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -30,6 +30,7 @@ from vllm_omni.distributed.omni_coordinator import ( RandomBalancer, RoundRobinBalancer, ) +from vllm_omni.engine.cuda_mps import CudaMPSServer, physical_gpu_uuid from vllm_omni.engine.messages import ( EngineQueueMessage, RegisterRemoteReplicaMessage, @@ -205,6 +206,7 @@ class StageRuntime: # device groups can initialize concurrently. self._replica_launch_lock = threading.Lock() self._init_visible_devices_baseline: str | None = None + self._mps_servers: dict[str, CudaMPSServer] = {} @staticmethod def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]: @@ -423,6 +425,8 @@ class StageRuntime: ) ) + with self._replica_launch_lock: + mps_env = self._mps_environment(physical_devices, replica.metadata.runtime_cfg) launch_context = launch_stage_replica( vllm_config=replica.stage_vllm_config, executor_class=replica.executor_class, @@ -432,7 +436,7 @@ class StageRuntime: stage_config=replica.stage_cfg, omni_master_server=self._get_omni_master_server(), omni_coordinator_address=self._get_coordinator_address(), - stage_visible_devices=physical_devices, + stage_visible_devices=mps_env.get("CUDA_VISIBLE_DEVICES", physical_devices), spawn_device_lock=self._spawn_device_lock, omni_parallel_stage_init=self._parallel_stage_init, num_api_servers=num_api_servers, @@ -441,7 +445,10 @@ class StageRuntime: # Environment overlays are process-global. Serialize spawn; # parallel initialization waits for READY outside this lock. with self._replica_launch_lock: - with stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg): + with ( + stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg), + stage_runtime_env(replica.metadata.stage_id, {"env": mps_env}), + ): stage_resources = launch_context.__enter__() entered_contexts.append(launch_context) if stage_resources is None: @@ -514,6 +521,55 @@ class StageRuntime: if lock_fds: release_device_locks(lock_fds) + @staticmethod + def _mps_enabled(runtime_cfg: Any) -> bool: + return bool( + runtime_cfg.get("cuda_mps", False) + if isinstance(runtime_cfg, Mapping) + else getattr(runtime_cfg, "cuda_mps", False) + ) + + def _validate_mps_topology(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + for plan in stage_plans: + for replica in plan.replicas: + runtime_cfg = getattr(replica.stage_cfg, "runtime_config", getattr(replica.stage_cfg, "runtime", None)) + if not self._mps_enabled(runtime_cfg): + continue + if replica.launch_mode != "local" or replica.metadata.stage_type == "diffusion": + raise ValueError("cuda_mps currently supports local EngineCore stages only") + if self._parallel_stage_init: + raise ValueError( + "cuda_mps requires parallel_stage_init=false for physical GPU initialization locks" + ) + + def _mps_environment(self, devices: str | None, runtime_cfg: Any) -> dict[str, str]: + if not self._mps_enabled(runtime_cfg): + return {} + if not current_omni_platform.is_cuda() or devices is None or len(devices.split(",")) != 1: + raise ValueError("cuda_mps currently requires a local CUDA stage on exactly one explicit GPU") + runtime_env = runtime_cfg.get("env") if isinstance(runtime_cfg, Mapping) else getattr(runtime_cfg, "env", None) + pipe_directory = ( + str(runtime_env["CUDA_MPS_PIPE_DIRECTORY"]) + if isinstance(runtime_env, Mapping) and "CUDA_MPS_PIPE_DIRECTORY" in runtime_env + else None + ) + uuid = physical_gpu_uuid(devices.strip()) + if uuid not in self._mps_servers: + self._mps_servers[uuid] = CudaMPSServer(uuid, pipe_directory=pipe_directory) + server = self._mps_servers[uuid] + if pipe_directory is not None and (pipe_directory or None) != server.operator_pipe_directory: + raise ValueError(f"Conflicting CUDA_MPS_PIPE_DIRECTORY settings for stages sharing GPU {uuid}") + return server.env + + def _close_mps_servers(self) -> None: + for uuid, server in list(self._mps_servers.items()): + try: + server.close() + except Exception: + logger.exception("Failed to close private MPS server for %s; retaining its control directory", uuid) + else: + del self._mps_servers[uuid] + def shutdown(self) -> None: for pool in self.stage_pools: for client in pool.clients: @@ -525,6 +581,7 @@ class StageRuntime: if self._stage_init_executor is not None: self._stage_init_executor.shutdown(wait=True, cancel_futures=True) self._stage_init_executor = None + self._close_mps_servers() def create_membership_controller(self) -> Any | None: """Return a distributed membership controller, if this runtime needs one.""" @@ -540,6 +597,7 @@ class StageRuntime: replicas_per_stage, replica_devices_map, ) + self._validate_mps_topology(stage_plans) self._validate_native_kv_topology(stage_plans) return stage_plans @@ -557,17 +615,20 @@ class StageRuntime: def _cleanup_after_initialize_failure(self) -> None: """Hook for runtimes that own extra infrastructure during init.""" - return None + self._close_mps_servers() def _resolve_replica_physical_devices(self, stage_id: int, runtime_cfg: Any) -> str | None: if runtime_cfg is None: runtime_cfg = {} devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) - return resolve_stage_physical_devices( + physical = resolve_stage_physical_devices( stage_id, devices, visible_baseline=self._init_visible_devices_baseline, ) + if self._mps_enabled(runtime_cfg) and (physical is None or not physical.isdigit()): + raise ValueError("cuda_mps requires one numeric physical GPU for initialization locking") + return physical # ---- Internal methods ---- @@ -664,7 +725,7 @@ class StageRuntime: port = int(extra.get("bootstrap_port", 8998)) + replica_metadata.replica_id extra["bootstrap_addr"] = f"http://{kv_config.kv_ip}:{port}" if isinstance(replica_cfg, BaseVllmOmniStageConfig): - runtime_cfg = replica_cfg.runtime_config + runtime_cfg: Any = replica_cfg.runtime_config runtime_cfg.env = { **(runtime_cfg.env or {}), "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(port), @@ -1374,6 +1435,7 @@ class DistStageRuntime(StageRuntime): self._stage_remote_factory_contexts = self._capture_stage_factory_contexts(stage_plans) def _cleanup_after_initialize_failure(self) -> None: + super()._cleanup_after_initialize_failure() self._cleanup_distributed_infra() def shutdown(self) -> None: diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 306301c94..ce2cd1e5e 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -1077,6 +1077,64 @@ class CodePredictorWrapper(nn.Module): ] ) + def _predict_step_logits( + self, + proj_buf: torch.Tensor, + bsz: int, + padded_bsz: int, + step: int, + is_npu_capturing: bool, + ) -> torch.Tensor: + """Run one residual-codebook step; model wrappers may specialize it.""" + device = proj_buf.device + max_seq = self._num_groups + 1 + model_fwd = self._compiled_model_fwd + lm_heads = self._lm_heads_list + graph_key: int | tuple[int, int] = padded_bsz + seq_len = max_seq + if self._prefix_reprefill_enabled: + actual_seq_len = step + 1 + for prefix_seq_len in self._prefix_reprefill_seq_lens: + if prefix_seq_len < actual_seq_len: + continue + prefix_key = (padded_bsz, prefix_seq_len) + if prefix_key in self._bucket_pos_ids: + graph_key = prefix_key + seq_len = prefix_seq_len + break + pos_ids = self._bucket_pos_ids.get(graph_key) + if pos_ids is None: + pos_ids = ( + torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous() + ) + + # Use captured device graph if available, otherwise call compiled fn. + device_graph_entry = self._device_graphs.get(graph_key) + + # Let the outer graph record the regular forward during capture; + # normal inference still uses the inner graph replay fast path. + if device_graph_entry is not None and not is_npu_capturing: + device_graph_entry[0].replay() + hidden_out = device_graph_entry[1] + else: + hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids) + + return lm_heads[step - 1](hidden_out[:bsz, step, :]) + + def _sample_per_call( + self, + logits: torch.Tensor, + inv_temperature: float, + top_k: int, + generator: _GeneratorLike, + uniforms: torch.Tensor | None, + ) -> torch.Tensor: + scaled = logits * inv_temperature + if top_k > 0: + topk_vals, _ = scaled.topk(top_k, dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms) + @torch.inference_mode() def forward( self, @@ -1114,10 +1172,7 @@ class CodePredictorWrapper(nn.Module): self._ensure_buffers(device, dtype, padded_bsz) proj_buf = self._proj_buf - max_seq = num_groups + 1 projection = self.small_to_mtp_projection - model_fwd = self._compiled_model_fwd - lm_heads = self._lm_heads_list codec_embeds = self._codec_embeds_list # torch-npu cannot replay an inner NPUGraph while the outer talker_mtp # graph is being captured. Capture state is constant for this forward. @@ -1153,39 +1208,7 @@ class CodePredictorWrapper(nn.Module): # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - graph_key: int | tuple[int, int] = padded_bsz - seq_len = max_seq - if self._prefix_reprefill_enabled: - actual_seq_len = step + 1 - for prefix_seq_len in self._prefix_reprefill_seq_lens: - if prefix_seq_len < actual_seq_len: - continue - prefix_key = (padded_bsz, prefix_seq_len) - if prefix_key in self._bucket_pos_ids: - graph_key = prefix_key - seq_len = prefix_seq_len - break - pos_ids = self._bucket_pos_ids.get(graph_key) - if pos_ids is None: - pos_ids = ( - torch.arange(seq_len, device=device, dtype=torch.long) - .unsqueeze(0) - .expand(padded_bsz, -1) - .contiguous() - ) - - # Use captured device graph if available, otherwise call compiled fn. - device_graph_entry = self._device_graphs.get(graph_key) - - # Let the outer graph record the regular forward during capture; - # normal inference still uses the inner graph replay fast path. - if device_graph_entry is not None and not is_npu_capturing: - device_graph_entry[0].replay() - hidden_out = device_graph_entry[1] - else: - hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids) - - logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) + logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) # Sample next code via Gumbel-max. # @@ -1216,12 +1239,8 @@ class CodePredictorWrapper(nn.Module): else: # "per_call" mode: temperature-scaled + top-k -> Gumbel-max if use_sampling: - scaled = logits * inv_temperature - if top_k > 0: - topk_vals, _ = scaled.topk(top_k, dim=-1) - scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None - code = self._sample_codes_gumbel(scaled, generator=sample_generator, uniforms=step_uniforms) + code = self._sample_per_call(logits, inv_temperature, top_k, sample_generator, step_uniforms) else: code = logits.argmax(dim=-1, keepdim=True) diff --git a/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py new file mode 100644 index 000000000..81eb8fdda --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Frame-local KV reuse for the Qwen3-TTS residual codebooks. + +The scratch buffer belongs to one predictor worker, just like its projection +buffer. Calls are serialized on the worker's CUDA stream. Every frame starts +by overwriting positions 0 and 1; later queries can only read that frame's +written positions. No request state survives between frames. +""" + +from collections.abc import Sequence + +import torch + +from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( + CodePredictorAttention, + CodePredictorBaseModel, + CodePredictorDecoderLayer, + _rotate_half, +) + +from .short_kv_attention import short_kv_attention + + +def _attention_cached( + attn: CodePredictorAttention, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, +) -> torch.Tensor: + """Attend only the new positions, reusing earlier positions' keys and values. + + ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``; + the new positions' keys and values are written at ``write_positions`` + (``[n]``). The attention kernel masks keys after each query position, + so unwritten slots are never read. + """ + bsz, seq_len, _ = hidden_states.shape + qkv = attn.qkv_proj(hidden_states) + q_raw, k_raw, v_raw = attn._split_qkv(qkv) + q = attn.q_norm(q_raw.view(bsz, seq_len, attn.num_heads, attn.head_dim)).transpose(1, 2) + k = attn.k_norm(k_raw.view(bsz, seq_len, attn.num_kv_heads, attn.head_dim)).transpose(1, 2) + v = v_raw.view(bsz, seq_len, attn.num_kv_heads, attn.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) + k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype)) + v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype)) + attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, attn.scaling) + attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) + return attn.o_proj(attn_out) + + +def _layer_cached( + layer: CodePredictorDecoderLayer, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, +) -> torch.Tensor: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + hidden_states = _attention_cached( + layer.self_attn, hidden_states, position_embeddings, k_cache, v_cache, write_positions + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = layer.mlp(hidden_states) + return residual + hidden_states + + +def forward_cached( + model: CodePredictorBaseModel, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]], + write_positions: torch.Tensor, +) -> torch.Tensor: + """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V.""" + hidden_states = inputs_embeds + with torch.amp.autocast(inputs_embeds.device.type, enabled=False): + position_embeddings = model.rotary_emb(hidden_states, position_ids) + for layer, (k_cache, v_cache) in zip(model.layers, kv_caches): + hidden_states = _layer_cached(layer, hidden_states, position_embeddings, k_cache, v_cache, write_positions) + hidden_states = model.norm(hidden_states) + return hidden_states + + +class FrameLocalKVCache: + """Non-module execution state; checkpoint parameters remain on the predictor.""" + + def __init__(self, model: CodePredictorBaseModel, max_batch: int) -> None: + self.model = model + weight = next(model.parameters()) + attn = model.layers[0].self_attn + self.buffer = torch.zeros( + len(model.layers), + 2, + max_batch, + attn.num_kv_heads, + attn.max_seq, + attn.head_dim, + device=weight.device, + dtype=weight.dtype, + ) + self.positions = { + step: torch.arange(0 if step == 1 else step, step + 1, device=weight.device) + for step in range(1, attn.max_seq - 1) + } + # Batch buckets share compiled kernels; outer MTP graphs still capture + # fixed addresses. This helper never creates nested CUDA graphs. + self.forward = torch.compile(forward_cached, dynamic=True, options={"epilogue_fusion": False}) + + def __call__(self, proj_buf: torch.Tensor, batch: int, step: int) -> torch.Tensor: + if batch > self.buffer.shape[2]: + raise ValueError("Qwen3-TTS cached predictor batch exceeds its configured capacity") + first = 0 if step == 1 else step + caches = [(layer[0, :batch], layer[1, :batch]) for layer in self.buffer] + positions = self.positions[step] + return self.forward( + self.model, + proj_buf[:batch, first : step + 1], + positions.unsqueeze(0).expand(batch, -1), + caches, + positions, + ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/code_predictor_sampling.py b/vllm_omni/model_executor/models/qwen3_tts/code_predictor_sampling.py new file mode 100644 index 000000000..8788988bd --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/code_predictor_sampling.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Top-k Gumbel sampling with the outer MTP graph's supplied uniforms.""" + +import torch +from vllm.triton_utils import tl, tldevice, triton + + +@triton.jit +def _sample( + logits, + uniforms, + output, + ls: tl.constexpr, + us: tl.constexpr, + lc: tl.constexpr, + uc: tl.constexpr, + vocab: tl.constexpr, + top_k: tl.constexpr, + topk_block: tl.constexpr, + inv_temperature: tl.constexpr, + block: tl.constexpr, +): + row = tl.program_id(0) + idx = tl.arange(0, block) + value = tl.load(logits + row * ls + idx * lc, idx < vocab, other=-float("inf")) + scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype).to(tl.float32) + if top_k > 0: + if top_k == 1: + threshold = tl.max(scaled, 0) + else: + ordered = tl.topk(scaled, topk_block) + threshold = tl.sum(tl.where(tl.arange(0, topk_block) == top_k - 1, ordered, 0)) + scaled = tl.where(scaled < threshold, -float("inf"), scaled) + u = tl.load(uniforms + row * us + idx * uc, idx < vocab, other=0.5) + score = scaled - tldevice.log(-tldevice.log(u)) + score = tl.where(idx < vocab, score, -float("inf")) + max_score = tl.max(score, 0) + chosen = tl.min(tl.where(score == max_score, idx, block), 0) + first_nan = tl.min(tl.where(score != score, idx, block), 0) + tl.store(output + row, tl.where(first_nan < block, first_nan, chosen)) + + +def sample_code_topk_gumbel( + logits: torch.Tensor, uniforms: torch.Tensor, top_k: int, inv_temperature: float +) -> torch.Tensor: + """Match BF16 scale rounding, top-k ties, and first-index Gumbel argmax. + + Random values are supplied by the outer graph, so RNG state and seeded + per-request streams are unchanged. CUDA libdevice logs match ATen's logs. + """ + if logits.shape != uniforms.shape: + raise ValueError("sampling uniforms must match logits shape") + if logits.device != uniforms.device or not logits.is_cuda: + raise ValueError("logits and uniforms must be on the same CUDA device") + if top_k > logits.shape[-1]: + raise ValueError("top_k must not exceed the vocabulary size") + b, v = logits.shape + out = torch.empty((b, 1), device=logits.device, dtype=torch.int64) + _sample[(b,)]( + logits, + uniforms, + out, + logits.stride(0), + uniforms.stride(0), + logits.stride(1), + uniforms.stride(1), + v, + top_k, + triton.next_power_of_2(max(1, top_k)), + inv_temperature, + triton.next_power_of_2(v), + num_warps=4, + ) + return out diff --git a/vllm_omni/model_executor/models/qwen3_tts/first_audio.py b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py new file mode 100644 index 000000000..1396c0dc4 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Qwen3-TTS first-frame decode eligibility.""" + +from typing import Any + +import torch + +from vllm_omni.platforms import current_omni_platform + + +def talker_first_audio_enabled(vllm_config: Any) -> bool: + """First-frame delivery requires CUDA MRv2 streaming with an in-process Talker. + + Unsupported runners retain the regular codec path. Prefix-cache replay + and distributed Talkers likewise retain the regular path. + """ + from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor + + extra = Predictor._stage_connector_extra_config(vllm_config) + model = vllm_config.model_config + parallel = vllm_config.parallel_config + return ( + Predictor._parse_bool_config(extra.get("talker_first_audio")) + and current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + and bool(getattr(model, "use_v2_model_runner", False)) + and bool(getattr(model, "async_chunk", False)) + and parallel.tensor_parallel_size == 1 + and parallel.pipeline_parallel_size == 1 + and parallel.distributed_executor_backend in (None, "uni") + and not vllm_config.cache_config.enable_prefix_caching + ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/first_frame_decoder.py b/vllm_omni/model_executor/models/qwen3_tts/first_frame_decoder.py new file mode 100644 index 000000000..40a66bf02 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/first_frame_decoder.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""First codec frame -> PCM inside the Talker process. + +A stream's first frame has no decoder context, so its waveform does not depend +on the Code2Wav stage's streaming state. Decoding it where the frame is +produced removes the stage hop from time to first audio. Code2Wav advances +its causal state without emitting that delivered frame again. Only streams +without reference-code context use this path. + +The decode runs beside the Talker's next step, so its latency is set by how +many dependent kernels it issues. A single frame with no context attends to +itself only: each softmax is over one key and equals 1 exactly, so attention +returns the value projection. The single-frame path therefore +skips q/k, RoPE, masks and softmax (about two thirds of the kernels). +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn as nn +from torch.cuda import CUDAGraph +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper + +from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config +from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Decoder + +logger = init_logger(__name__) + +_SUBFOLDER = "speech_tokenizer" + + +class Qwen3TTSFirstFrameDecoder(nn.Module): + def __init__(self, model_path: str) -> None: + super().__init__() + self.model_path = model_path + config = Qwen3TTSTokenizerV2Config.from_pretrained(model_path, subfolder=_SUBFOLDER) + self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(config.decoder_config) + self.decoder.eval() + self.num_quantizers = int(config.decoder_config.num_quantizers) + self.sample_rate = int(config.output_sample_rate) + self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {} + + def load(self, vllm_config: VllmConfig) -> set[str]: + """Load decoder weights exactly as Code2Wav does; returns the loaded parameter names.""" + loader = DefaultModelLoader(vllm_config.load_config) + source = DefaultModelLoader.Source( + model_or_path=self.model_path, + revision=vllm_config.model_config.revision, + subfolder=_SUBFOLDER, + ) + loaded = AutoWeightsLoader(self).load_weights( + loader._get_weights_iterator(source), + mapper=WeightsMapper(orig_to_new_prefix={"encoder.": None}), + ) + self.decoder.to(device=vllm_config.device_config.device, dtype=vllm_config.model_config.dtype) + if hasattr(self.decoder, "precompute_snake_caches"): + self.decoder.precompute_snake_caches() + from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor + + extra = Predictor._stage_connector_extra_config(vllm_config) + if Predictor._parse_bool_config(extra.get("decode_time_major_conv")): + self.decoder.enable_time_major_conv() + return loaded + + def _decode_single_frame(self, codes: torch.Tensor) -> torch.Tensor: + """``codes`` [n, num_quantizers, 1] -> waveform; equals ``_decode_xvec_first_chunk``.""" + d = self.decoder + hidden = d.pre_conv(d.quantizer.decode(codes)).transpose(1, 2) + t = d.pre_transformer + h = t.input_proj(hidden) + for layer in t.layers[: t.config.num_hidden_layers]: + attention = layer.self_attn + values = attention.v_proj(layer.input_layernorm(h)) + if attention.num_key_value_groups != 1: + values = values.reshape(*values.shape[:-1], -1, attention.head_dim) + values = values.repeat_interleave(attention.num_key_value_groups, dim=-2).flatten(-2) + h = h + layer.self_attn_layer_scale(attention.o_proj(values)) + h = h + layer.mlp_layer_scale(layer.mlp(layer.post_attention_layernorm(h))) + return d._conv_decode(t.output_proj(t.norm(h))).clamp(min=-1, max=1) + + @torch.inference_mode() + def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None: + device = next(self.decoder.parameters()).device + for batch_size in sorted(set(batch_sizes)): + static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device) + self._decode_single_frame(static_input) + torch.accelerator.synchronize(device) + graph = CUDAGraph() + # A separate pool prevents overlap with Talker/MTP graph allocations. + with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()): + static_output = self._decode_single_frame(static_input) + self._graphs[batch_size] = (graph, static_input, static_output) + logger.info("Captured Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs)) + + @torch.inference_mode() + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor).""" + if not self._graphs: + return self._decode_single_frame(codes.unsqueeze(-1))[:, 0, :].float() + outputs = [] + largest = max(self._graphs) + for start in range(0, int(codes.shape[0]), largest): + chunk = codes[start : start + largest] + rows = int(chunk.shape[0]) + batch_size = min(size for size in self._graphs if size >= rows) + graph, static_input, static_output = self._graphs[batch_size] + static_input.zero_() + static_input[:rows, :, 0].copy_(chunk) + graph.replay() + outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True)) + return torch.cat(outputs, dim=0) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 2701520b2..a9c16bbe5 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -6,6 +6,7 @@ from __future__ import annotations import os from collections import Counter from collections.abc import Iterable +from contextlib import nullcontext from typing import Any import torch @@ -16,6 +17,7 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp @@ -112,6 +114,7 @@ class Qwen3TTSCode2Wav(nn.Module): dec_config = tok_config.decoder_config self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(dec_config) self.decoder.eval() + self._num_quantizers = int(dec_config.num_quantizers) self._output_sample_rate = int(tok_config.output_sample_rate) self._total_upsample = int(self.decoder.total_upsample) @@ -302,6 +305,7 @@ class Qwen3TTSCode2Wav(nn.Module): valid_indices: list[int] = [] left_context_size = [0] * len(request_ids_list) ref_context_size = [0] * len(request_ids_list) + first_audio_flags = [False] * len(request_ids_list) segment_finished_flags = [False] * len(request_ids_list) request_state_ids: list[str | None] = [None] * len(request_ids_list) ref_context_request_ids: list[str | None] = [None] * len(request_ids_list) @@ -336,6 +340,7 @@ class Qwen3TTSCode2Wav(nn.Module): if not isinstance(info, dict): continue meta = info.get("meta", {}) + first_audio_flags[i] = _meta_bool(meta.get("first_audio", False)) if "is_segment_finished" in meta: segment_finished_flags[i] = _meta_bool(meta["is_segment_finished"]) if "left_context_size" in meta: @@ -404,6 +409,7 @@ class Qwen3TTSCode2Wav(nn.Module): ) state = {} self._decoder_state_cache[state_req_id] = state + state.setdefault("skip_first_audio", first_audio_flags[i]) state.setdefault("prefix_frames", 0) if state_req_id.startswith(_DUMMY_REQUEST_ID): state["_is_dummy_run"] = True @@ -515,6 +521,10 @@ class Qwen3TTSCode2Wav(nn.Module): if self._batch_stats_log_every > 0 and self._batch_stats_forwards % self._batch_stats_log_every == 0: self.log_decode_batch_stats() + first_audio_required = [torch.tensor(False)] * num_req + if request_states is not None: + for row, idx in enumerate(valid_indices): + first_audio_required[idx] = torch.tensor(bool(request_states[row].get("skip_first_audio", False))) audios: list[torch.Tensor] = [empty] * num_req srs = [sr_tensor] * num_req @@ -538,7 +548,7 @@ class Qwen3TTSCode2Wav(nn.Module): return OmniOutput( text_hidden_states=None, - multimodal_outputs={"model_outputs": audios, "sr": srs}, + multimodal_outputs={"model_outputs": audios, "sr": srs, FIRST_AUDIO_REQUIRED_KEY: first_audio_required}, ) def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: @@ -714,6 +724,11 @@ class Qwen3TTSCode2Wav(nn.Module): raise ValueError(f"Invalid Qwen3-TTS Code2Wav config decode_batch_max_size={decode_batch_max_size}") self._decode_batch_max_size = decode_batch_max_size decode_enable_tf32 = _get_bool_config("decode_enable_tf32", False) + decode_time_major_conv = _get_bool_config("decode_time_major_conv", False) + self.decoder.capture_first_audio_state_only = self._async_chunk and _get_bool_config( + "talker_first_audio", False + ) + decode_cudnn_benchmark = _get_bool_config("decode_cudnn_benchmark", False) else: codec_chunk_frames = 0 codec_left_context_frames = 0 @@ -722,6 +737,8 @@ class Qwen3TTSCode2Wav(nn.Module): decode_cudagraph_batch_sizes = None decode_cudagraph_capture_sizes = None decode_enable_tf32 = False + decode_cudnn_benchmark = False + decode_time_major_conv = False if decode_enable_tf32 and device.type == "cuda": # PyTorch exposes TF32 controls as process-wide CUDA backend @@ -738,21 +755,39 @@ class Qwen3TTSCode2Wav(nn.Module): torch.get_float32_matmul_precision(), ) + if decode_time_major_conv and device.type == "cuda": + self.decoder.enable_time_major_conv() + self.decoder._initial_codec_chunk_frames = initial_codec_chunk_frames self.decoder._incremental_chunk_frames = codec_chunk_frames or 25 self.decoder._incremental_chunk_ramp = list(codec_chunk_ramp or ()) if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda": try: - self._maybe_enable_decoder_cudagraph( - device=device, - codec_chunk_frames=codec_chunk_frames, - codec_left_context_frames=codec_left_context_frames, - initial_codec_chunk_frames=initial_codec_chunk_frames, - codec_chunk_ramp=codec_chunk_ramp, - decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes, - decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes, + # Autotune only during warmup/capture, then restore the process + # flags. The captured convolution algorithms remain in the + # graphs without changing later models' cuDNN policy. + autotune = ( + torch.backends.cudnn.flags( + enabled=torch.backends.cudnn.enabled, + benchmark=True, + benchmark_limit=10, + deterministic=torch.backends.cudnn.deterministic, + allow_tf32=torch.backends.cudnn.allow_tf32, + ) + if decode_cudnn_benchmark + else nullcontext() ) + with autotune: + self._maybe_enable_decoder_cudagraph( + device=device, + codec_chunk_frames=codec_chunk_frames, + codec_left_context_frames=codec_left_context_frames, + initial_codec_chunk_frames=initial_codec_chunk_frames, + codec_chunk_ramp=codec_chunk_ramp, + decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes, + decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes, + ) except Exception: logger.warning( "Failed to enable CUDA Graph for Code2Wav decoder", diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 9d00cc2ba..675c5eecd 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -1,12 +1,17 @@ -"""Qwen3-TTS Code Predictor -- thin wrapper over CodePredictorWrapper.""" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Qwen3-TTS predictor: per-call sampling and opt-in CUDA frame-local KV reuse.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING import torch from vllm.config import VllmConfig from vllm.config.vllm import set_current_vllm_config +from vllm.logger import init_logger from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorBaseModel, @@ -17,6 +22,11 @@ from vllm_omni.platforms import current_omni_platform from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig +if TYPE_CHECKING: + from .cached_code_predictor import FrameLocalKVCache + +logger = init_logger(__name__) + # Backward-compat alias used by tests Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel @@ -48,6 +58,82 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp # Store talker_config for backward compat (accessed by some callers) self.talker_config = talker_config self._vllm_config = vllm_config + extra = self._stage_connector_extra_config(vllm_config) + self._kv_requested = self._parse_bool_config(extra.get("code_predictor_kv_cache")) + self._fused_sampling = self._parse_bool_config(extra.get("code_predictor_fused_sampling")) + self._frame_cache: FrameLocalKVCache | None = None + + def _setup_compile(self) -> None: + if self._compiled_model_fwd is not None: + return + weight = next(self.model.parameters()) + supported = ( + current_omni_platform.is_cuda() + and weight.device.type == "cuda" + and weight.dtype == torch.bfloat16 + and current_omni_platform.supports_torch_inductor() + ) + if not self._kv_requested or not supported: + if self._kv_requested: + logger.warning("Qwen3-TTS frame-local KV requires CUDA BF16; using re-prefill") + return super()._setup_compile() + + # Import CUDA-only kernels only when this model explicitly opts in. + from .cached_code_predictor import FrameLocalKVCache + + self._model_dtype = weight.dtype + self._lm_heads_list = list(self.lm_head) + self._codec_embeds_list = list(self.model.codec_embedding) + self._bucket_sizes = self._batch_bucket_sizes() + max_batch = max(self._bucket_sizes) + self._ensure_buffers(weight.device, weight.dtype, max_batch) + cache = FrameLocalKVCache(self.model, max_batch) + with torch._dynamo.config.patch( + cache_size_limit=max(torch._dynamo.config.cache_size_limit, 2 * len(self._bucket_sizes)) + ): + for batch in self._bucket_sizes: + for step in range(1, min(3, self._num_groups)): + for _ in range(2): + cache(self._proj_buf, batch, step) + self._synchronize_warmup(weight.device) + self._frame_cache = cache + # Preserve the shared wrapper's initialization sentinel. Its full + # re-prefill callable is unused while _frame_cache is active. + self._compiled_model_fwd = self.model.forward + logger.info("Qwen3-TTS frame-local KV warmed for batches %s", self._bucket_sizes) + + def _predict_step_logits( + self, + proj_buf: torch.Tensor, + bsz: int, + padded_bsz: int, + step: int, + is_npu_capturing: bool, + ) -> torch.Tensor: + if self._frame_cache is None: + return super()._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) + hidden = self._frame_cache(proj_buf, padded_bsz, step) + return self._lm_heads_list[step - 1](hidden[:bsz, -1]) + + def _sample_per_call( + self, + logits: torch.Tensor, + inv_temperature: float, + top_k: int, + generator: torch.Generator | Sequence[torch.Generator | None] | None, + uniforms: torch.Tensor | None, + ) -> torch.Tensor: + if ( + self._fused_sampling + and current_omni_platform.is_cuda() + and logits.is_cuda + and uniforms is not None + and uniforms.dtype == torch.float32 + ): + from .code_predictor_sampling import sample_code_topk_gumbel + + return sample_code_topk_gumbel(logits, uniforms, top_k, inv_temperature) + return super()._sample_per_call(logits, inv_temperature, top_k, generator, uniforms) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with vllm config context (required for VocabParallelEmbedding).""" diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index b795052b7..e05ac3871 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -36,6 +36,7 @@ from vllm_omni.utils.speaker_cache import ( from vllm_omni.worker.sampling_utils import get_tts_local_seed from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig +from .first_audio import talker_first_audio_enabled from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config @@ -361,6 +362,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): } ) + # Set by the MRV2 model state once eager frames are enabled. + eager_frames_active = False + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.vllm_config = vllm_config @@ -404,6 +408,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # per-request additional_information. self.talker_mtp_output_key = ("codes", "audio") self.talker_mtp_graph_safe = True + # A frame's residual codebooks depend only on the Talker hidden and + # CB0 of the step that sampled it, so MRV2 may complete the frame at + # the end of that step instead of in the next step's preprocess. + self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) # The runners bypass only the outer whole-MTP graph when explicit # generators are present, so seeded requests can still share one raw # batched MTP call with independent per-row streams. @@ -722,6 +730,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) async_chunk = bool(getattr(model_config, "async_chunk", False)) + spans = kwargs.get("request_token_spans") + if async_chunk and self.eager_frames_active and spans is not None: + return self._make_eager_omni_output(hidden, info_dicts, spans) audio_codes_list: list[torch.Tensor] = [] ref_code_len_segments: list[tuple[int | torch.Tensor, int, torch.device]] = [] ref_code_list: list[torch.Tensor] = [] @@ -863,6 +874,44 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): mm.setdefault("meta", {})["codec_frame_valid"] = codec_frame_valid[:span_len] return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + def _make_eager_omni_output( + self, + hidden: torch.Tensor, + info_dicts: list[Any], + spans: list[tuple[int, int]], + ) -> OmniOutput: + """Async-chunk output when the runner fills frames after sampling. + + Every row that sampled this step gets its codes and validity from the + runner's eager MTP; the rest (non-final prefill chunks) carry no frame, + so the per-request buffer rows need not be gathered here. + """ + num_tokens = spans[-1][1] if spans else 0 + q = int(self.talker_config.num_code_groups) + mm: OmniPayload = { + "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)}, + "meta": { + "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + }, + } + ref_rows = [ + index + for index, info in enumerate(info_dicts) + if isinstance(info, dict) + and isinstance(ref := info.get("codes", {}).get("ref"), torch.Tensor) + and ref.numel() > 0 + and _should_publish_async_ref_codes(info) + ] + if ref_rows: + empty = torch.empty(0, dtype=torch.long) + ref_code_list = [empty] * len(info_dicts) + for index in ref_rows: + ref_code_list[index] = info_dicts[index]["codes"]["ref"] + _mark_async_ref_codes_published(info_dicts[index]) + mm["codes"]["ref"] = ref_code_list + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + # -------------------- preprocess / postprocess -------------------- def preprocess( @@ -1087,7 +1136,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): f"preprocess_decode_batch expected {len(req_infos)} input ids, got {int(input_ids_flat.numel())}" ) - codec_frame_valid_batch = (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size) + eager = self.eager_frames_active + codec_frame_valid_batch = ( + None if eager else (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size) + ) device = input_ids_flat.device dtype = self._embedding_dtype # Request-independent constant (see :meth:`_init_runtime_buffers`) — @@ -1150,28 +1202,29 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): text_step = tts_pad_embed next_text_offset = text_offset - last_hidden = hs.get("last") - if isinstance(last_hidden, torch.Tensor): - past_hidden = last_hidden.to(device=device, dtype=dtype).reshape(1, -1) - else: - # Match scalar preprocess(): EOS / async-scheduling races can - # arrive before postprocess has written last; zeros are filtered - # downstream the same way as the single-request path. - logger.warning_once( - "Qwen3-TTS batched decode missing last hidden; zero-filling past_hidden " - "to match scalar preprocess()" - ) - past_hidden = torch.zeros_like(text_step) - past_hidden_list.append(past_hidden) text_step_list.append(text_step) - info_update: dict[str, Any] = { "meta": { "talker_text_offset": int(next_text_offset), "codec_streaming": codec_streaming, - "codec_frame_valid": codec_frame_valid_batch[len(updates)].reshape(()), }, } + if not eager: + # Eager frames finished MTP last step: no hidden input or row validity needed. + last_hidden = hs.get("last") + if isinstance(last_hidden, torch.Tensor): + past_hidden = last_hidden.to(device=device, dtype=dtype).reshape(1, -1) + else: + # Match scalar preprocess(): EOS / async-scheduling races can + # arrive before postprocess has written last; zeros are filtered + # downstream the same way as the single-request path. + logger.warning_once( + "Qwen3-TTS batched decode missing last hidden; zero-filling past_hidden " + "to match scalar preprocess()" + ) + past_hidden = torch.zeros_like(text_step) + past_hidden_list.append(past_hidden) + info_update["meta"]["codec_frame_valid"] = codec_frame_valid_batch[len(updates)].reshape(()) if trailing_text_update is not None: info_update["hidden_states"] = {"trailing_text": trailing_text_update.detach()} updates.append(info_update) @@ -1179,11 +1232,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if input_embeds is None: input_embeds = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long)) inputs_embeds_out = input_embeds.to(device=device, dtype=dtype).reshape(len(req_infos), -1) + text_steps = torch.cat(text_step_list, dim=0) return ( input_ids_flat, inputs_embeds_out, - torch.cat(past_hidden_list, dim=0), - torch.cat(text_step_list, dim=0), + text_steps if eager else torch.cat(past_hidden_list, dim=0), + text_steps, updates, ) @@ -1382,6 +1436,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded)) self._build_stacked_codec_embed() + if talker_first_audio_enabled(self.vllm_config): + from .first_frame_decoder import Qwen3TTSFirstFrameDecoder + + self.first_frame_decoder = Qwen3TTSFirstFrameDecoder(self.model_path) + decoder_loaded = self.first_frame_decoder.load(self.vllm_config) + loaded = set(loaded) | {f"first_frame_decoder.{name}" for name in decoder_loaded} return loaded def _build_stacked_codec_embed(self) -> None: @@ -1569,6 +1629,35 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): mtp = talker_mtp get_mtp_seed = staticmethod(get_tts_local_seed) + # Postprocess only stores the last hidden for the deferred decode MTP. + eager_frames_need_postprocess = False + + def eager_decode_settled(self, info: dict[str, Any]) -> bool: + """Decode preprocess is a no-op for this row: its trailing text is spent. + + From then on every decode step feeds the frame embedding plus the pad + text step and changes no per-request state (the non-streaming + CustomVoice/VoiceDesign case after the first decode). + """ + hs = info.get("hidden_states") + tail = hs.get("trailing_text") if isinstance(hs, dict) else None + meta = info.get("meta") + offset = meta.get("talker_text_offset", 0) if isinstance(meta, dict) else 0 + return isinstance(tail, torch.Tensor) and tail.numel() == 0 and not offset + + def eager_settled_text_step(self) -> torch.Tensor: + return self._tts_pad_embed.reshape(1, -1) + + def capture_first_frame_graphs(self) -> None: + # Set by load_weights when talker_first_audio_enabled accepts the connector option. + decoder = getattr(self, "first_frame_decoder", None) + if decoder is not None: + decoder.capture() + + def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor: + """Rows whose CB0 is a codec id rather than EOS/special; matches decode preprocess.""" + return (layer0 >= 0) & (layer0 < self._codebook_vocab_size) + @property def mtp_output_key(self) -> tuple[str, str]: return self.talker_mtp_output_key diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py index a87af5f85..ce9d86f0b 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py @@ -71,6 +71,7 @@ class CUDAGraphDecoderWrapper: self.combined_states: dict[tuple[str, int, int], dict] = {} self.icl_prefix_states: dict[int, dict] = {} self.xvec_prefix_states: dict[int, dict] = {} + self.xvec_prefix_state_only_states: dict[int, dict] = {} self.stateless_states: dict[tuple[int, int], dict] = {} self._device = None @@ -230,6 +231,7 @@ class CUDAGraphDecoderWrapper: *self._retained_kv_caches(self.combined_states), *self._retained_kv_caches(self.icl_prefix_states), *self._retained_kv_caches(self.xvec_prefix_states), + *self._retained_kv_caches(self.xvec_prefix_state_only_states), ] batched_captures = sum(rows for rows, _ in retained) @@ -248,7 +250,7 @@ class CUDAGraphDecoderWrapper: rows_reason=( f"rows across {len(retained)} retained caches in " f"{len(self.combined_states)} suffix / {len(self.icl_prefix_states)} icl-prefix / " - f"{len(self.xvec_prefix_states)} xvec-prefix captures" + f"{len(self.xvec_prefix_states) + len(self.xvec_prefix_state_only_states)} xvec-prefix captures" ), allocation_note=( "counted from caches transformers has materialized; captures whose tensors are still " @@ -424,6 +426,14 @@ class CUDAGraphDecoderWrapper: logger.info("Captured xvec prefix CUDA Graph for batch=%d", batch_size) except Exception: logger.warning("Failed to capture xvec prefix CUDA Graph for batch=%d", batch_size, exc_info=True) + if getattr(self.decoder, "capture_first_audio_state_only", False) and self.initial_chunk_frames == 1: + try: + self._capture_xvec_prefix(batch_size, device, dtype, state_only=True) + logger.info("Captured xvec state-only CUDA Graph for batch=%d", batch_size) + except Exception: + logger.warning( + "Failed to capture xvec state-only graph for batch=%d", batch_size, exc_info=True + ) logger.info( "Starting CUDA Graph warmup for %d shapes: batch_sizes=%s seq_lens=%s", @@ -631,7 +641,9 @@ class CUDAGraphDecoderWrapper: "cache": caches, } - def _capture_xvec_prefix(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> None: + def _capture_xvec_prefix( + self, batch_size: int, device: torch.device, dtype: torch.dtype, *, state_only: bool = False + ) -> None: static_input = torch.zeros( batch_size, self.num_quantizers, @@ -639,8 +651,13 @@ class CUDAGraphDecoderWrapper: dtype=dtype, device=device, ) + # The Talker delivers exactly one frame upstream: a one-frame first + # chunk then only needs its state; a longer one keeps the other frames. + first_chunk = ( + self.decoder._decode_xvec_first_chunk_state_only if state_only else self.decoder._decode_xvec_first_chunk + ) with torch.no_grad(): - _ = self.decoder._decode_xvec_first_chunk(static_input, {}) + _ = first_chunk(static_input, {}) torch.accelerator.synchronize(device) config = self.decoder.config @@ -662,13 +679,15 @@ class CUDAGraphDecoderWrapper: graph = CUDAGraph() with torch.no_grad(): with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): - static_output = self.decoder._decode_xvec_first_chunk(static_input, caches) + static_output = first_chunk(static_input, caches) - self.xvec_prefix_states[batch_size] = { + states = self.xvec_prefix_state_only_states if state_only else self.xvec_prefix_states + states[batch_size] = { "graph": graph, "input": {"codes": static_input}, "output": static_output, "cache": caches, + "state_only": state_only, } def _ensure_suffix_buffers(self, caches: dict) -> None: @@ -808,7 +827,14 @@ class CUDAGraphDecoderWrapper: codes_list: list[torch.Tensor], request_caches: list[dict], ) -> list[torch.Tensor] | None: - available = set(getattr(self, "xvec_prefix_states", {})) + skip_flags = [bool(cache.get("skip_first_audio", False)) for cache in request_caches] + # Full PCM graphs remain available even when the Talker cannot use + # direct delivery (or only some requests qualify). The option is only + # a capture hint; request metadata decides which audio is still owed. + states = getattr(self, "xvec_prefix_states", {}) + if skip_flags and all(skip_flags): + states = getattr(self, "xvec_prefix_state_only_states", {}) or states + available = set(states) if not available: self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list)) return None @@ -835,7 +861,7 @@ class CUDAGraphDecoderWrapper: self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list)) return None - state = self.xvec_prefix_states[batch_size] + state = states[batch_size] static_input = state["input"]["codes"] static_input.zero_() for row, codes in enumerate(codes_list): @@ -855,7 +881,12 @@ class CUDAGraphDecoderWrapper: cache["suffix_conv"] = static_caches["suffix_conv"][row : row + 1].clone() cache["suffix_frames"] = self.initial_chunk_frames self._ensure_suffix_buffers(cache) - outputs.append(state["output"][row : row + 1].clone()) + if state.get("state_only"): + outputs.append(static_input.new_zeros((1, 1, 0), dtype=torch.float32)) + elif skip_flags[row]: + outputs.append(state["output"][row : row + 1, :, self.decoder.total_upsample :].clone()) + else: + outputs.append(state["output"][row : row + 1].clone()) return outputs def _decode_suffix_batch( @@ -983,6 +1014,8 @@ class CUDAGraphDecoderWrapper: def _decode_request_fallback(self, codes: torch.Tensor, cache: dict) -> torch.Tensor: if "suffix_quantized" not in cache: if int(cache["prefix_frames"]) == 0: + if cache.get("skip_first_audio", False): + return self.decoder._decode_stream_first_chunk(codes, cache) return self.decoder._decode_xvec_first_chunk(codes, cache) prefix_frames = int(cache["prefix_frames"]) output = self.decoder._decode_icl_first_chunk(codes, cache, prefix_frames) diff --git a/vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py b/vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py new file mode 100644 index 000000000..35019000b --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Short causal attention for the scratch KV cache inside a code predictor.""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _short_kv_attention( + q_ptr, + k_ptr, + v_ptr, + positions_ptr, + out_ptr, + q_stride_b: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_n: tl.constexpr, + k_stride_b: tl.constexpr, + k_stride_h: tl.constexpr, + k_stride_n: tl.constexpr, + v_stride_b: tl.constexpr, + v_stride_h: tl.constexpr, + v_stride_n: tl.constexpr, + heads: tl.constexpr, + num_queries: tl.constexpr, + num_keys: tl.constexpr, + head_dim: tl.constexpr, + queries_per_kv: tl.constexpr, + scale: tl.constexpr, + block_n: tl.constexpr, + block_d: tl.constexpr, +): + b = tl.program_id(0) + h = tl.program_id(1) + qi = tl.program_id(2) + kvh = h // queries_per_kv + d = tl.arange(0, block_d) + n = tl.arange(0, block_n) + last = tl.load(positions_ptr + qi) + q = tl.load(q_ptr + b * q_stride_b + h * q_stride_h + qi * q_stride_n + d, d < head_dim, other=0).to(tl.float32) + k = tl.load( + k_ptr + b * k_stride_b + kvh * k_stride_h + n[:, None] * k_stride_n + d[None, :], + (n[:, None] < num_keys) & (n[:, None] <= last) & (d[None, :] < head_dim), + other=0, + ).to(tl.float32) + score = tl.sum(k * q[None, :], 1) * scale + score = tl.where((n < num_keys) & (n <= last), score, -float("inf")) + p = tl.exp(score - tl.max(score, 0)) + p = p / tl.sum(p, 0) + v = tl.load( + v_ptr + b * v_stride_b + kvh * v_stride_h + n[:, None] * v_stride_n + d[None, :], + (n[:, None] < num_keys) & (n[:, None] <= last) & (d[None, :] < head_dim), + other=0, + ).to(tl.float32) + out = tl.sum(p[:, None] * v, 0) + tl.store(out_ptr + ((b * num_queries + qi) * heads + h) * head_dim + d, out, d < head_dim) + + +def short_kv_attention( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, positions: torch.Tensor, scale: float +) -> torch.Tensor: + """One/two-query GQA over a frame-local scratch cache with contiguous head dimensions. + + The returned [B, H, Q, D] tensor has a contiguous [B, Q, H, D] view, + avoiding a separate transpose-copy before the output projection. + """ + b, h, nq, d = q.shape + nk = k.shape[2] + out = torch.empty((b, nq, h, d), device=q.device, dtype=q.dtype) + _short_kv_attention[(b, h, nq)]( + q, + k, + v, + positions, + out, + *q.stride()[:3], + *k.stride()[:3], + *v.stride()[:3], + h, + nq, + nk, + d, + h // k.shape[1], + scale, + triton.next_power_of_2(nk), + triton.next_power_of_2(d), + num_warps=4, + ) + return out.transpose(1, 2) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 7d6d721bd..432f2e7e7 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -997,6 +997,33 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): if count > 0: logger.info("Precomputed exp caches for %d SnakeBeta activations", count) + # Set by enable_time_major_conv(); a plain object so it stays out of state_dict(). + _time_major_stack = None + + def enable_time_major_conv(self) -> None: + """Run the upsample + decoder convs time-major: one tensor-core GEMM per conv. + + Call after the decoder reaches its device and dtype (the stack keeps + GEMM-layout weight copies) and before CUDA graph capture. + """ + from .time_major_stack import TimeMajorConvStack + + self._time_major_stack = TimeMajorConvStack(self) + logger.info("Qwen3-TTS decoder: time-major (GEMM) conv stack enabled") + + def _conv_decode(self, hidden: torch.Tensor) -> torch.Tensor: + """Pre-transformer output ``[B, T, latent]`` -> unclamped waveform ``[B, 1, T * total_upsample]``.""" + if self._time_major_stack is not None: + return self._time_major_stack(hidden) + hidden = hidden.permute(0, 2, 1) + for blocks in self.upsample: + for block in blocks: + hidden = block(hidden) + wav = hidden + for block in self.decoder: + wav = block(wav) + return wav + def enable_cudagraph( self, capture_modes: tuple[str, ...] = ("icl", "xvec"), @@ -1056,15 +1083,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): hidden = self.pre_conv(hidden).transpose(1, 2) hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state - hidden = hidden.permute(0, 2, 1) - for blocks in self.upsample: - for block in blocks: - hidden = block(hidden) - - wav = hidden - for block in self.decoder: - wav = block(wav) - return wav.clamp(min=-1, max=1) + return self._conv_decode(hidden).clamp(min=-1, max=1) def _decode_icl_first_chunk( self, @@ -1126,20 +1145,32 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): caches.update({"past_key_values": prefix_cache}) caches.update({"prefix_hidden": prefix_hidden}) - hidden = hidden.permute(0, 2, 1) - - for blocks in self.upsample: - for block in blocks: - hidden = block(hidden) - wav = hidden - - for block in self.decoder: - wav = block(wav) - wav = wav.clamp(min=-1, max=1) - return wav + return self._conv_decode(hidden).clamp(min=-1, max=1) def _decode_xvec_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: """Initialize prefixless incremental state from the first codec chunk.""" + suffix_hidden = self._init_xvec_first_chunk_state(codes, caches) + return self._conv_decode(suffix_hidden).clamp(min=-1, max=1) + + def _decode_xvec_first_chunk_state_only(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: + """State for later chunks when the first chunk's audio was delivered upstream.""" + return self._init_xvec_first_chunk_state(codes, caches) + + # Capture hint only: the raw connector option does not guarantee that + # a request received upstream audio. Only its skip_first_audio marker does. + capture_first_audio_state_only = False + + def _decode_stream_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: + if not caches.get("skip_first_audio", False): + return self._decode_xvec_first_chunk(codes, caches) + if int(codes.shape[-1]) <= 1: + self._decode_xvec_first_chunk_state_only(codes, caches) + return codes.new_zeros((codes.shape[0], 1, 0), dtype=torch.float32) + # The Talker delivered only the first frame; a longer first chunk + # (adaptive or ramped schedules) still owes the audio of the rest. + return self._decode_xvec_first_chunk(codes, caches)[..., self.total_upsample :] + + def _init_xvec_first_chunk_state(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: hidden = self.quantizer.decode(codes) caches["decoder_prefix_frames"] = 0 caches["ref_hidden"] = hidden[:, :, :0] @@ -1172,16 +1203,8 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): ).last_hidden_state caches["past_key_values"] = empty_prefix_cache caches["prefix_hidden"] = suffix_hidden[:, :0, :] - - hidden = suffix_hidden.permute(0, 2, 1) - for blocks in self.upsample: - for block in blocks: - hidden = block(hidden) - wav = hidden - for block in self.decoder: - wav = block(wav) caches["suffix_frames"] = int(codes.shape[-1]) - return wav.clamp(min=-1, max=1) + return suffix_hidden def _decode_suffix( self, @@ -1229,14 +1252,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): hidden = torch.cat([caches["prefix_hidden"], suffix_hidden], dim=1) required_frames = min(hidden.shape[1], new_frames + _DOWNSTREAM_CONTEXT_FRAME) - hidden = hidden[:, -required_frames:, :].permute(0, 2, 1) - for blocks in self.upsample: - for block in blocks: - hidden = block(hidden) - - wav = hidden - for block in self.decoder: - wav = block(wav) + wav = self._conv_decode(hidden[:, -required_frames:, :]) wav = wav[..., -new_frames * self.total_upsample :].clamp(min=-1, max=1) return wav, next_quantized, next_conv @@ -1300,7 +1316,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}" ) if prefix_frames == 0: - return self._decode_xvec_first_chunk(codes, caches) + return self._decode_stream_first_chunk(codes, caches) return self._decode_icl_first_chunk(codes, caches, prefix_frames) else: decoder_prefix_frames = int(caches.get("decoder_prefix_frames", prefix_frames)) @@ -1546,7 +1562,11 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): return None batched_codes = torch.cat(codes_list, dim=0) batched_cache: dict[str, Any] = {} - output = self._decode_xvec_first_chunk(batched_codes, batched_cache) + skip_flags = {bool(cache.get("skip_first_audio", False)) for cache in request_caches} + if len(skip_flags) != 1: + return None + batched_cache["skip_first_audio"] = skip_flags.pop() + output = self._decode_stream_first_chunk(batched_codes, batched_cache) outputs: list[torch.Tensor] = [] for row, cache in enumerate(request_caches): for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"): diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_conv.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_conv.py new file mode 100644 index 000000000..322421af5 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_conv.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Time-major causal 1-D convolution for streaming vocoders. + +Activations are ``[B, T, C]`` so that every dense convolution is one GEMM over +``B * T`` rows (bf16 tensor cores through cuBLAS) instead of an NCT cuDNN +convolution, which for these long, thin decoder layers picks non-tensor-core +kernels and adds layout transforms around each call. + +- :func:`snake_im2col` gathers each output step's causal, dilated window into a + ``[B * T, K * C]`` row (tap-major), optionally applying ``bias`` and SnakeBeta + on the way in; ``K = 1`` is a standalone (bias +) SnakeBeta. +- :func:`causal_conv` is the dense causal (dilated) convolution as an implicit + GEMM: each K-block of the reduction reads the input rows shifted by one tap, + so the ``K * C`` im2col matrix is never materialized; bias and an optional + residual are added in the epilogue. +- :func:`transposed_overlap_add` finishes a transposed convolution of kernel + ``stride`` or ``2 * stride`` whose GEMM produced ``[B, T, K, C_out]``: the + overlap-add of consecutive inputs, the bias and the causal right trim. + +Weights come from ``nn.Conv1d``/``nn.ConvTranspose1d`` via :func:`conv_gemm_weight` +and :func:`transposed_gemm_weight`. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + + +def conv_gemm_weight(weight: torch.Tensor) -> torch.Tensor: + """``nn.Conv1d`` weight ``[C_out, C_in, K]`` -> tap-major GEMM weight ``[K * C_in, C_out]``.""" + c_out, c_in, k = weight.shape + return weight.detach().permute(2, 1, 0).reshape(k * c_in, c_out).contiguous() + + +def transposed_gemm_weight(weight: torch.Tensor) -> torch.Tensor: + """``nn.ConvTranspose1d`` weight ``[C_in, C_out, K]`` -> GEMM weight ``[C_in, K * C_out]``.""" + c_in, c_out, k = weight.shape + return weight.detach().permute(0, 2, 1).reshape(c_in, k * c_out).contiguous() + + +def snake_im2col_reference( + x: torch.Tensor, + kernel_size: int, + dilation: int = 1, + bias: torch.Tensor | None = None, + snake: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Eager equivalent of :func:`snake_im2col` (also the CPU path).""" + if bias is not None: + x = x + bias + if snake is not None: + exp_alpha, inv_beta = snake + x = x + (inv_beta.float() * torch.sin(x.float() * exp_alpha.float()).square()).to(x.dtype) + if kernel_size == 1: + return x.reshape(-1, x.shape[-1]) + t = x.shape[1] + xp = torch.nn.functional.pad(x, (0, 0, (kernel_size - 1) * dilation, 0)) + cols = torch.cat([xp[:, j * dilation : j * dilation + t] for j in range(kernel_size)], dim=-1) + return cols.reshape(-1, cols.shape[-1]) + + +def causal_conv_reference( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + kernel_size: int, + dilation: int = 1, + residual: torch.Tensor | None = None, +) -> torch.Tensor: + """Eager equivalent of :func:`causal_conv`: ``[B, T, C]`` -> ``[B, T, C_out]``.""" + bsz, t, _ = x.shape + cols = snake_im2col_reference(x, kernel_size, dilation) + out = cols @ weight if bias is None else torch.addmm(bias, cols, weight) + out = out.view(bsz, t, -1) + return out if residual is None else out + residual + + +def transposed_overlap_add_reference(z: torch.Tensor, stride: int, bias: torch.Tensor) -> torch.Tensor: + """Eager equivalent of :func:`transposed_overlap_add`.""" + bsz, t, k, c_out = z.shape + out = z[:, :, :stride] + if k > stride: + out = out + torch.nn.functional.pad(z[:, :-1, stride:], (0, 0, 0, 0, 1, 0)) + return out.reshape(bsz, t * stride, c_out) + bias + + +if HAS_TRITON: + + @triton.jit + def _snake_im2col_kernel( + x_ptr, + bias_ptr, + exp_alpha_ptr, + inv_beta_ptr, + out_ptr, + n_rows, + t_len, + c_len, + dilation, + K: tl.constexpr, # noqa: N803 + HAS_BIAS: tl.constexpr, # noqa: N803 + HAS_SNAKE: tl.constexpr, # noqa: N803 + BLOCK_R: tl.constexpr, # noqa: N803 + BLOCK_C: tl.constexpr, # noqa: N803 + ): + rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R) + chans = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + row_mask = rows < n_rows + chan_mask = chans < c_len + t = rows % t_len + if HAS_BIAS: + bias = tl.load(bias_ptr + chans, mask=chan_mask, other=0.0) + if HAS_SNAKE: + exp_alpha = tl.load(exp_alpha_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32) + inv_beta = tl.load(inv_beta_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32) + out_mask = row_mask[:, None] & chan_mask[None, :] + for j in tl.static_range(K): + shift = (K - 1 - j) * dilation + valid = row_mask & (t >= shift) + src = rows - shift + x = tl.load( + x_ptr + src[:, None].to(tl.int64) * c_len + chans[None, :], + mask=valid[:, None] & chan_mask[None, :], + other=0.0, + ) + if HAS_BIAS: + # Causal padding is applied to the conv input, after bias. + x = tl.where(valid[:, None], x + bias[None, :], 0.0).to(x_ptr.dtype.element_ty) + if HAS_SNAKE: + s = tl.sin(x.to(tl.float32) * exp_alpha[None, :]) + x = x + (inv_beta[None, :] * s * s).to(x.dtype) + tl.store(out_ptr + rows[:, None].to(tl.int64) * (K * c_len) + j * c_len + chans[None, :], x, mask=out_mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk}, num_warps=w, num_stages=st) + for bm, bn, bk, w, st in ( + (128, 128, 64, 8, 3), + (128, 64, 64, 4, 4), + (64, 128, 64, 4, 4), + (128, 32, 32, 4, 4), + (64, 64, 32, 4, 4), + (256, 32, 32, 8, 3), + (256, 64, 32, 8, 3), + (128, 16, 32, 4, 4), + ) + ], + key=["row_bucket", "n_cols", "c_len", "K", "dilation"], + ) + @triton.jit + def _causal_conv_gemm_kernel( + x_ptr, + w_ptr, + bias_ptr, + res_ptr, + out_ptr, + n_rows, + n_cols, + c_len, + t_len, + dilation, + row_bucket, # autotune key only: power-of-two bucket of n_rows + K: tl.constexpr, # noqa: N803 + HAS_BIAS: tl.constexpr, # noqa: N803 + HAS_RES: tl.constexpr, # noqa: N803 + PRECISION: tl.constexpr, # noqa: N803 + BLOCK_M: tl.constexpr, # noqa: N803 + BLOCK_N: tl.constexpr, # noqa: N803 + BLOCK_K: tl.constexpr, # noqa: N803 + ): + # Row-major tiles, N fastest: neighbouring programs share the input rows. + pid = tl.program_id(0) + num_n = tl.cdiv(n_cols, BLOCK_N) + pid_m = pid // num_n + pid_n = pid % num_n + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = rows < n_rows + col_mask = cols < n_cols + t = rows % t_len + k_blocks = tl.cdiv(c_len, BLOCK_K) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for kk in range(0, K * k_blocks): + tap = kk // k_blocks + chans = (kk % k_blocks) * BLOCK_K + tl.arange(0, BLOCK_K) + chan_mask = chans < c_len + shift = (K - 1 - tap) * dilation + a = tl.load( + x_ptr + (rows - shift)[:, None] * c_len + chans[None, :], + mask=(row_mask & (t >= shift))[:, None] & chan_mask[None, :], + other=0.0, + ) + b = tl.load( + w_ptr + (tap * c_len + chans)[:, None] * n_cols + cols[None, :], + mask=chan_mask[:, None] & col_mask[None, :], + other=0.0, + ) + acc = tl.dot(a, b, acc, input_precision=PRECISION) + if HAS_BIAS: + acc += tl.load(bias_ptr + cols, mask=col_mask, other=0.0).to(tl.float32)[None, :] + out_offsets = rows[:, None] * n_cols + cols[None, :] + out_mask = row_mask[:, None] & col_mask[None, :] + if HAS_RES: + acc += tl.load(res_ptr + out_offsets, mask=out_mask, other=0.0).to(tl.float32) + tl.store(out_ptr + out_offsets, acc.to(out_ptr.dtype.element_ty), mask=out_mask) + + @triton.jit + def _transposed_overlap_add_kernel( + z_ptr, + bias_ptr, + out_ptr, + n_out_rows, + t_len, + c_len, + STRIDE: tl.constexpr, # noqa: N803 + K: tl.constexpr, # noqa: N803 + BLOCK_R: tl.constexpr, # noqa: N803 + BLOCK_C: tl.constexpr, # noqa: N803 + ): + rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R) # over B * T * STRIDE output steps + chans = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + row_mask = rows < n_out_rows + chan_mask = chans < c_len + mask = row_mask[:, None] & chan_mask[None, :] + tap = rows % STRIDE + src = rows // STRIDE # input row b * T + t + acc = tl.load( + z_ptr + (src[:, None].to(tl.int64) * K + tap[:, None]) * c_len + chans[None, :], mask=mask, other=0.0 + ).to(tl.float32) + if K > STRIDE: + prev_valid = row_mask & (src % t_len > 0) + acc += tl.load( + z_ptr + ((src[:, None].to(tl.int64) - 1) * K + tap[:, None] + STRIDE) * c_len + chans[None, :], + mask=prev_valid[:, None] & chan_mask[None, :], + other=0.0, + ).to(tl.float32) + acc += tl.load(bias_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32)[None, :] + tl.store( + out_ptr + rows[:, None].to(tl.int64) * c_len + chans[None, :], acc.to(out_ptr.dtype.element_ty), mask=mask + ) + + +def _blocks(c_len: int) -> tuple[int, int]: + block_c = min(128, triton.next_power_of_2(c_len)) + return max(1, 4096 // block_c), block_c + + +def snake_im2col( + x: torch.Tensor, + kernel_size: int, + dilation: int = 1, + bias: torch.Tensor | None = None, + snake: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """``[B, T, C]`` -> ``[B * T, K * C]`` causal im2col with optional bias + SnakeBeta.""" + if not HAS_TRITON or not x.is_cuda: + return snake_im2col_reference(x, kernel_size, dilation, bias, snake) + x = x.contiguous() + bsz, t, c = x.shape + n_rows = bsz * t + out = torch.empty(n_rows, kernel_size * c, dtype=x.dtype, device=x.device) + block_r, block_c = _blocks(c) + exp_alpha, inv_beta = snake if snake is not None else (x, x) + _snake_im2col_kernel[(triton.cdiv(n_rows, block_r), triton.cdiv(c, block_c))]( + x, + bias if bias is not None else x, + exp_alpha, + inv_beta, + out, + n_rows, + t, + c, + dilation, + K=kernel_size, + HAS_BIAS=bias is not None, + HAS_SNAKE=snake is not None, + BLOCK_R=block_r, + BLOCK_C=block_c, + ) + return out + + +def causal_conv( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + kernel_size: int, + dilation: int = 1, + residual: torch.Tensor | None = None, +) -> torch.Tensor: + """Causal dilated conv ``[B, T, C]`` -> ``[B, T, C_out]`` with a :func:`conv_gemm_weight` weight. + + ``residual`` (``[B, T, C_out]``) is added in the epilogue. + """ + if not HAS_TRITON or not x.is_cuda: + return causal_conv_reference(x, weight, bias, kernel_size, dilation, residual) + x = x.contiguous() + bsz, t, c = x.shape + n_rows, n_cols = bsz * t, weight.shape[1] + if weight.shape[0] != kernel_size * c: + raise ValueError(f"conv weight has {weight.shape[0]} rows, expected {kernel_size} taps x {c} channels") + if max(n_rows * max(c, n_cols), weight.numel()) >= 2**31: + raise ValueError("causal_conv offsets are 32-bit; split the batch") + out = torch.empty(bsz, t, n_cols, dtype=x.dtype, device=x.device) + if residual is not None: + residual = residual.contiguous() + + def grid(meta): + return (triton.cdiv(n_rows, meta["BLOCK_M"]) * triton.cdiv(n_cols, meta["BLOCK_N"]),) + + _causal_conv_gemm_kernel[grid]( + x, + weight, + bias if bias is not None else x, + residual if residual is not None else x, + out, + n_rows, + n_cols, + c, + t, + dilation, + triton.next_power_of_2(triton.cdiv(n_rows, 1024)), + K=kernel_size, + HAS_BIAS=bias is not None, + HAS_RES=residual is not None, + PRECISION="ieee" if x.dtype == torch.float32 else "tf32", + ) + return out + + +def transposed_overlap_add(z: torch.Tensor, stride: int, bias: torch.Tensor) -> torch.Tensor: + """``[B, T, K, C_out]`` GEMM output of a kernel-``K`` transposed conv -> ``[B, T * stride, C_out]``.""" + bsz, t, k, c_out = z.shape + if k not in (stride, 2 * stride): + raise ValueError(f"transposed_overlap_add supports kernel stride or 2*stride, got kernel={k} stride={stride}") + if not HAS_TRITON or not z.is_cuda: + return transposed_overlap_add_reference(z, stride, bias) + z = z.contiguous() + n_out_rows = bsz * t * stride + out = torch.empty(bsz, t * stride, c_out, dtype=z.dtype, device=z.device) + block_r, block_c = _blocks(c_out) + _transposed_overlap_add_kernel[(triton.cdiv(n_out_rows, block_r), triton.cdiv(c_out, block_c))]( + z, bias, out, n_out_rows, t, c_out, STRIDE=stride, K=k, BLOCK_R=block_r, BLOCK_C=block_c + ) + return out diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_stack.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_stack.py new file mode 100644 index 000000000..59526140c --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_stack.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Time-major upsample + decoder conv stack for the Qwen3-TTS 12 Hz codec. + +Same math as ``Qwen3TTSTokenizerV2Decoder.upsample``/``.decoder`` on ``[B, T, C]`` +activations: dense causal convs are implicit GEMMs (tensor cores, no im2col +buffer, residual adds in the epilogue), SnakeBeta one pass per activation, +transposed convs a GEMM plus an overlap-add. See ``time_major_conv.py``. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.time_major_conv import ( + causal_conv, + conv_gemm_weight, + snake_im2col, + transposed_gemm_weight, + transposed_overlap_add, +) + + +def _snake(module) -> tuple[torch.Tensor, torch.Tensor]: + if module._exp_alpha is None: + module.precompute_exp_cache() + return module._exp_alpha, module._inv_beta + + +class _Conv: + def __init__(self, causal_conv) -> None: + conv = causal_conv.conv + if conv.stride[0] != 1 or conv.groups != 1: + raise ValueError("time-major decoder convs must be dense with stride 1") + self.weight = conv_gemm_weight(conv.weight) + self.bias = conv.bias.detach() + self.kernel_size = conv.kernel_size[0] + self.dilation = conv.dilation[0] + + def __call__(self, x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor: + return causal_conv(x, self.weight, self.bias, self.kernel_size, self.dilation, residual) + + +class _Transposed: + def __init__(self, causal_trans_conv) -> None: + conv = causal_trans_conv.conv + self.weight = transposed_gemm_weight(conv.weight) + self.bias = conv.bias.detach() + self.kernel_size = conv.kernel_size[0] + self.stride = conv.stride[0] + self.out_channels = conv.out_channels + + def __call__(self, x2d: torch.Tensor, bsz: int, t: int) -> torch.Tensor: + z = torch.mm(x2d, self.weight).view(bsz, t, self.kernel_size, self.out_channels) + return transposed_overlap_add(z, self.stride, self.bias) + + +class TimeMajorConvStack: + """Callable replacing the NCT conv loops; holds GEMM-layout copies of the weights.""" + + def __init__(self, decoder: nn.Module) -> None: + self.upsample = [(_Transposed(trans), block) for trans, block in decoder.upsample] + self.conv_in = _Conv(decoder.decoder[0]) + self.blocks = [] + for block in decoder.decoder[1:-2]: + snake, trans, *units = block.block + residuals = [(_snake(u.act1), _Conv(u.conv1), _snake(u.act2), _Conv(u.conv2)) for u in units] + self.blocks.append((_snake(snake), _Transposed(trans), residuals)) + self.snake_out = _snake(decoder.decoder[-2]) + self.conv_out = _Conv(decoder.decoder[-1]) + + @staticmethod + def _convnext(x: torch.Tensor, block) -> torch.Tensor: + dw = block.dwconv.conv + k, t = dw.kernel_size[0], x.shape[1] + xp = torch.nn.functional.pad(x, (0, 0, k - 1, 0)) + weight = dw.weight[:, 0, :].float() + h = dw.bias.float() + sum(xp[:, j : j + t].float() * weight[:, j] for j in range(k)) + h = block.pwconv2(block.act(block.pwconv1(block.norm(h.to(x.dtype))))) + return x + block.gamma * h + + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + """``hidden`` ``[B, T, latent]`` -> unclamped waveform ``[B, 1, T * upsample]``.""" + bsz = hidden.shape[0] + + def snake(x, params): + return snake_im2col(x, 1, snake=params).view(x.shape) + + for trans, block in self.upsample: + t = hidden.shape[1] + hidden = self._convnext(trans(hidden.reshape(bsz * t, -1), bsz, t), block) + hidden = self.conv_in(hidden) + for block_snake, trans, residuals in self.blocks: + t = hidden.shape[1] + hidden = trans(snake_im2col(hidden, 1, snake=block_snake), bsz, t) + for act1, conv1, act2, conv2 in residuals: + hidden = conv2(snake(conv1(snake(hidden, act1)), act2), residual=hidden) + wav = self.conv_out(snake(hidden, self.snake_out)) + return wav.view(bsz, 1, -1) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 79a6e782a..482514c04 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -407,6 +407,12 @@ def talker2code2wav_async_chunk( left_context_size=left_context_size, finished=torch.tensor(finished, dtype=torch.bool), ) + source_meta = multimodal_output.get("meta", {}) if isinstance(multimodal_output, Mapping) else {} + first_audio = source_meta.get("first_audio", False) + if isinstance(first_audio, torch.Tensor): + first_audio = bool(first_audio.numel() and first_audio.reshape(-1)[-1].item()) + if first_audio: + meta.first_audio = torch.tensor(True, dtype=torch.bool) if ref_context_size > 0 and ref_context_request_id is not None: meta.ref_context_size = ref_context_size meta.ref_context_request_id = ref_context_request_id @@ -494,11 +500,13 @@ def talker2code2wav_async_chunk_batch( request_payload[request_id] = value.contiguous() payloads: list[OmniPayloadStruct | None] = [] - for request, finished in zip(requests, is_finished): + for request, finished, pooling_output in zip(requests, is_finished, pooling_outputs): payloads.append( talker2code2wav_async_chunk( transfer_manager=transfer_manager, - multimodal_output={}, + multimodal_output={"meta": pooling_output.get("meta", {})} + if isinstance(pooling_output, Mapping) + else {}, request=request, is_finished=finished, ) diff --git a/vllm_omni/utils/device_copy.py b/vllm_omni/utils/device_copy.py new file mode 100644 index 000000000..2b271940f --- /dev/null +++ b/vllm_omni/utils/device_copy.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Host-to-device copies that do not block the host on queued GPU work. + +A copy from pageable host memory (``tensor.to("cuda")``, ``torch.tensor(..., +device="cuda")``, ``torch.as_tensor(..., device="cuda")``) makes the CPU wait +until every kernel already queued on the stream has finished. In a per-step +path under async scheduling that removes the CPU/GPU overlap the batch queue +exists for. Staging the data in +pinned memory and copying with ``non_blocking`` keeps the host running ahead; +the caching host allocator keeps the pinned source alive until the copy's +stream event completes. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +def index_to_device(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.long) -> torch.Tensor: + """A host index list as a device tensor, without a host sync.""" + if torch.device(device).type != "cuda": + return torch.tensor(values, dtype=dtype, device=device) + return torch.tensor(values, dtype=dtype, pin_memory=True).to(device, non_blocking=True) diff --git a/vllm_omni/worker_v2/first_audio_sender.py b/vllm_omni/worker_v2/first_audio_sender.py new file mode 100644 index 000000000..7e9246b5b --- /dev/null +++ b/vllm_omni/worker_v2/first_audio_sender.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Deliver in-stage first audio as soon as its device copy completes. + +The regular output of a step is published only after the whole step's +outputs are copied, materialized for every request, popped by the engine loop +and serialized together. A stream's first audio chunk does not need any of +that, so a dedicated thread waits for the chunk's own copy event and hands a +one-request output straight to the engine core's output queue. +""" + +from __future__ import annotations + +import queue +import threading +from dataclasses import dataclass +from typing import Any, Protocol + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + +logger = init_logger(__name__) + + +class FirstAudioSink(Protocol): + def prepare(self, request_ids: list[str]) -> _PreparedDelivery: + """Freeze request routes on the submitting thread before accepting delivery.""" + ... + + +class FirstAudioSender: + def __init__(self, sink: FirstAudioSink) -> None: + if not callable(getattr(sink, "prepare", None)): + raise TypeError("FirstAudioSender requires a sink with prepare(request_ids)") + self._sink = sink + self._queue: queue.SimpleQueue[ + tuple[torch.cuda.Event, torch.Tensor, list[str], torch.Tensor, torch.Tensor | None, _PreparedDelivery] + | None + ] + self._queue = queue.SimpleQueue() + self._thread = threading.Thread(target=self._run, daemon=True, name="omni-first-audio-sender") + self._thread.start() + + def submit( + self, + request_ids: list[str], + pcm: torch.Tensor, + sample_rate: torch.Tensor, + valid: torch.Tensor | None = None, + ) -> list[str]: + """Queue a D2H copy of ``pcm`` [n, samples] on the current stream and deliver it once done. + + ``valid`` [n] (on device) drops rows whose frame turned out not to be + audio (e.g. a codec EOS sample); it is read only after the copy. + + Return the request IDs whose delivery was accepted. Rows without an + engine output route must retain regular codec delivery instead of + promising the orchestrator a first frame that can never arrive. + """ + delivery = self._sink.prepare(request_ids) + accepted = list(delivery.routes) + if not accepted: + return [] + host = torch.empty(pcm.shape, dtype=pcm.dtype, pin_memory=True) + host.copy_(pcm, non_blocking=True) + host_valid = None + if valid is not None: + host_valid = torch.empty(valid.shape, dtype=torch.bool, pin_memory=True) + host_valid.copy_(valid, non_blocking=True) + copied = torch.cuda.Event() + copied.record() + self._queue.put((copied, host, list(request_ids), sample_rate, host_valid, delivery)) + return accepted + + def close(self) -> None: + self._queue.put(None) + self._thread.join(timeout=5) + + def _run(self) -> None: + while (item := self._queue.get()) is not None: + copied, host, request_ids, sample_rate, host_valid, delivery = item + pending = request_ids + try: + copied.synchronize() + rows = [row for row in range(len(request_ids)) if host_valid is None or bool(host_valid[row])] + pending = [request_ids[row] for row in rows] + if rows: + delivery(pending, [host[row].clone() for row in rows], sample_rate) + except Exception: + # If synchronization failed, validity is not readable yet. + # Otherwise only valid, not-yet-delivered rows owe first audio. + logger.exception("First-audio delivery failed for %s", pending) + delivery.fail(pending) + + +@dataclass +class _PreparedDelivery: + sink: _EngineOutputSink + # Frozen client indices for outstanding deliveries. Successful enqueues + # retire their routes so a later failure cannot abort completed deliveries. + routes: dict[str, int] + + def __call__(self, request_ids: list[str], pcm_rows: list[torch.Tensor], sample_rate: torch.Tensor) -> None: + self.sink._send(self.routes, request_ids, pcm_rows, sample_rate) + + def fail(self, request_ids: list[str]) -> None: + self.sink._fail(self.routes, request_ids) + + +class _EngineOutputSink: + def __init__(self, output_queue: Any, scheduler: Any) -> None: + self.output_queue = output_queue + self.scheduler = scheduler + + def prepare(self, request_ids: list[str]) -> _PreparedDelivery: + routes = {} + for request_id in request_ids: + request = self.scheduler.requests.get(request_id) + if request is not None: + routes[request_id] = int(getattr(request, "client_index", 0) or 0) + return _PreparedDelivery(self, routes) + + def _send(self, routes: dict[str, int], request_ids, pcm_rows, sample_rate) -> None: + from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs + + by_client: dict[int, list[OmniEngineCoreOutput]] = {} + for request_id, pcm in zip(request_ids, pcm_rows, strict=True): + if request_id not in routes: + continue + 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), + }, + ) + ) + for client_index, outputs in by_client.items(): + self.output_queue.put_nowait((client_index, OmniEngineCoreOutputs(outputs=outputs))) + for output in outputs: + routes.pop(output.request_id) + + def _fail(self, routes: dict[str, int], request_ids: list[str]) -> None: + from vllm.v1.engine import FinishReason + + from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs + + for request_id in request_ids: + if request_id in routes: + output = OmniEngineCoreOutput( + request_id=request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + stop_reason="First-audio delivery failed", + ) + self.output_queue.put_nowait((routes[request_id], OmniEngineCoreOutputs(outputs=[output]))) + routes.pop(request_id) + + +def engine_output_queue_sink(output_queue: Any, scheduler: Any) -> _EngineOutputSink: + return _EngineOutputSink(output_queue, scheduler) diff --git a/vllm_omni/worker_v2/model_states/eager_mtp.py b/vllm_omni/worker_v2/model_states/eager_mtp.py new file mode 100644 index 000000000..ba347206a --- /dev/null +++ b/vllm_omni/worker_v2/model_states/eager_mtp.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Opt-in completion of a Talker frame immediately after sampling codebook zero.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch +from vllm.logger import init_logger +from vllm.v1.worker.gpu.input_batch import InputBatch + +from vllm_omni.utils.device_copy import index_to_device + +if TYPE_CHECKING: + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +logger = init_logger(__name__) + + +def _has_ref_codes(buffer: dict[str, Any]) -> bool: + ref = buffer.get("codes", {}).get("ref") if isinstance(buffer.get("codes"), dict) else None + return isinstance(ref, torch.Tensor) and ref.numel() > 0 + + +class EagerMTPState: + """Used only by models explicitly declaring ``mtp_eager_frames``.""" + + def __init__(self, owner: OmniModelState) -> None: + self.owner = owner + self._first_audio_valid: torch.Tensor | None = None + + def set_first_audio_sink(self, sink: Any) -> None: + from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender + + self.owner._first_audio_sender = FirstAudioSender(sink) + + def _apply_eager_frames( + self, + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], + embeds: torch.Tensor, + input_batch: InputBatch, + prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None, + ) -> None: + """Decode input = codec embeddings of the previous (eager) frame + this step's text.""" + req_indices = [int(input_batch.idx_mapping_np[i]) for i, _start, _mtp in mtp_batches] + for req_idx in req_indices: + req_id = self.owner.intermediate_buffer.buffers[req_idx].get("req_id") + if self.owner._eager_ready.get(req_idx) != req_id: + # Running the deferred MTP here would re-emit or drop a frame. + raise RuntimeError(f"Eager Talker-MTP frame missing for request {req_id!r}") + assert self.owner._eager_embeds is not None + device = embeds.device + if prepacked_mtp_inputs is None: + text_step = torch.cat([step.reshape(1, -1) for _i, _start, (_hidden, step) in mtp_batches], dim=0) + else: + text_step = prepacked_mtp_inputs[1].reshape(len(mtp_batches), -1) + rows = index_to_device(req_indices, device) + offsets = self.owner._mtp_batch_offsets(mtp_batches, input_batch, device) + frame_embeds = self.owner._eager_embeds.index_select(0, rows) + embeds.index_copy_(0, offsets, (frame_embeds + text_step.to(frame_embeds.dtype)).to(embeds.dtype)) + + def _apply_settled_frames(self, settled_rows: list[tuple[int, int, int, str]], embeds: torch.Tensor) -> None: + """Settled decode rows: previous eager frame embedding plus the constant text step.""" + for _i, req_idx, _start, req_id in settled_rows: + if self.owner._eager_ready.get(req_idx) != req_id: + raise RuntimeError(f"Eager Talker-MTP frame missing for request {req_id!r}") + assert self.owner._eager_embeds is not None + device = embeds.device + rows = index_to_device([req_idx for _i, req_idx, _start, _req_id in settled_rows], device) + offsets = index_to_device([start for _i, _req_idx, start, _req_id in settled_rows], device) + text_step = self.owner.model.eager_settled_text_step().to(device=device, dtype=self.owner._eager_embeds.dtype) + frame_embeds = self.owner._eager_embeds.index_select(0, rows) + embeds.index_copy_(0, offsets, (frame_embeds + text_step).to(embeds.dtype)) + + def run_eager_mtp( + self, + input_batch: InputBatch, + text_hidden: torch.Tensor, + sampled_token_ids: torch.Tensor, + multimodal_outputs: dict[str, Any], + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None, + ) -> None: + """Complete each sampled row's frame in this step and publish it with this step's output. + + Consumes the rows recorded by ``run_preprocess``. Writes the frame codes + and validity into the last token row of each request span of the + retained multimodal output, and keeps the frame's codec embedding sum + for the next step's input. Issued on the main stream before the async + output copy, so no extra synchronization is needed. + """ + if not self.owner._eager_mtp: + return + recorded, self.owner._eager_rows = self.owner._eager_rows, None + if recorded is None or recorded[0] is not input_batch or not recorded[1]: + return + entries, input_ids = recorded[1], recorded[2] + codes_out = multimodal_outputs.get("codes", {}).get("audio") + valid_out = multimodal_outputs.get("meta", {}).get("codec_frame_valid") + if not isinstance(codes_out, torch.Tensor) or not isinstance(valid_out, torch.Tensor): + raise RuntimeError("Eager Talker-MTP requires retained codes.audio and codec_frame_valid outputs") + assert self.owner._mtp_input_ids is not None and self.owner._mtp_input_embeds is not None + assert self.owner._mtp_hidden is not None and self.owner._mtp_text_step is not None + assert self.owner._eager_embeds is not None + + bsz = len(entries) + device = text_hidden.device + batch_rows = index_to_device([i for i, _req_idx, _req_id, _prefill in entries], device) + req_indices = [req_idx for _i, req_idx, _req_id, _prefill in entries] + last_tokens = input_batch.query_start_loc.index_select(0, batch_rows + 1).long() - 1 + layer0 = sampled_token_ids.reshape(input_batch.num_reqs, -1)[:, 0].index_select(0, batch_rows) + + batch_ids = self.owner._mtp_input_ids[:bsz] + batch_emb = self.owner._mtp_input_embeds[:bsz] + batch_hidden = self.owner._mtp_hidden[:bsz] + batch_step = self.owner._mtp_text_step[:bsz] + batch_ids.copy_(layer0.to(batch_ids.dtype)) + batch_emb.copy_(self.owner.model.embed_input_ids(batch_ids.reshape(-1, 1)).reshape(bsz, -1)) + torch.index_select(text_hidden, 0, last_tokens, out=batch_hidden) + # The text step is added by the next step's preprocess. + batch_step.zero_() + frame_embeds, codes = self.owner._mtp_forward( + req_indices, + batch_ids, + batch_emb, + batch_hidden, + batch_step, + mtp_batch_descriptor_dispatcher, + ) + assert codes is not None + rows = index_to_device(req_indices, device) + self.owner._eager_embeds.index_copy_( + 0, rows, frame_embeds[:bsz].reshape(bsz, -1).to(self.owner._eager_embeds.dtype) + ) + if self.owner.vllm_config.cache_config.enable_prefix_caching: + # Prefix-hit reconstruction reads the latest frame from the buffer. + self.owner.intermediate_buffer.update_gpu_tensor_rows( + req_indices, self.owner.model.mtp_output_key, codes[:bsz] + ) + if codes_out.ndim != 2: + raise RuntimeError(f"Eager Talker-MTP expects token-major codes.audio, got {tuple(codes_out.shape)}") + codes_out.index_copy_(0, last_tokens, codes[:bsz].to(codes_out.dtype)) + # A decode row whose input CB0 was EOS belongs to a request that already + # ended (async scheduling runs it once more); its new sample must not + # become a frame. Prefill rows have no codec input. + prefill_rows = index_to_device([int(prefill) for _i, _req_idx, _req_id, prefill in entries], device).bool() + input_valid = self.owner.model.mtp_frame_valid(input_ids.index_select(0, last_tokens)) | prefill_rows + valid = self.owner.model.mtp_frame_valid(layer0) & input_valid + valid_out.index_copy_(0, last_tokens, valid.to(valid_out.dtype)) + for _i, req_idx, req_id, _prefill in entries: + self.owner._eager_ready[req_idx] = req_id + self._publish_first_audio(input_batch, entries, codes[:bsz], valid) + first_audio = multimodal_outputs.get("meta", {}).get("first_audio") + if isinstance(first_audio, torch.Tensor): + scheduled = index_to_device( + [int(req_id in self.owner._first_audio_requests) for _, _, req_id, _ in entries], device + ) + if self._first_audio_valid is None: + scheduled.zero_() + else: + scheduled *= self._first_audio_valid.index_select(0, rows) + first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype)) + + def _publish_first_audio( + self, + input_batch: InputBatch, + entries: list[tuple[int, int, str, bool]], + codes: torch.Tensor, + valid: torch.Tensor | None = None, + ) -> None: + """Decode first frames here and hand the PCM to the client output. + + Only streams without reference codes qualify: their first chunk is + decoded with no context, so the waveform equals what Code2Wav will + produce for it. ``valid`` (per entry, on device) marks rows whose frame + is real; direct delivery drops the others once their copy is done, so a + stream whose first sample is codec EOS gets no first audio from here. + """ + decoder = getattr(self.owner.model, "first_frame_decoder", None) + if decoder is None: + logger.info_once("Talker first-frame audio: no decoder on %s", type(self.owner.model).__name__) + return + first = [ + (row, i) + for row, (i, req_idx, _req_id, prefill) in enumerate(entries) + if prefill + and _req_id not in self.owner._first_audio_requests + and not _has_ref_codes(self.owner.intermediate_buffer.buffers[req_idx]) + ] + if not first: + return + # Decode on a side stream so the Talker's next step does not queue + # behind it; the async output copy waits on its event instead. + stream = self.owner._first_audio_stream + if stream is None: + _least, greatest = torch.cuda.Stream.priority_range() + stream = self.owner._first_audio_stream = torch.cuda.Stream(device=codes.device, priority=greatest) + frame_rows = index_to_device([row for row, _i in first], codes.device) + frame_codes = codes.index_select(0, frame_rows) + frame_valid = valid.index_select(0, frame_rows) if valid is not None else None + if frame_valid is not None: + frame_codes = torch.where(frame_valid[:, None], frame_codes, 0) + if self._first_audio_valid is None: + assert self.owner._eager_embeds is not None + self._first_audio_valid = torch.zeros( + self.owner._eager_embeds.shape[0], dtype=torch.bool, device=codes.device + ) + request_rows = index_to_device([entries[row][1] for row, _ in first], codes.device) + if frame_valid is not None: + self._first_audio_valid.index_copy_(0, request_rows, frame_valid) + else: + self._first_audio_valid.index_fill_(0, request_rows, True) + stream.wait_stream(torch.cuda.current_stream(codes.device)) + sr = torch.tensor(decoder.sample_rate, dtype=torch.int32) + with torch.cuda.stream(stream): + frame_codes.record_stream(stream) + pcm = decoder.decode(frame_codes) + sender = self.owner._first_audio_sender + if sender is None: + raise RuntimeError("First-frame decoder requires an engine output sink") + if frame_valid is not None: + frame_valid.record_stream(stream) + request_ids = [entries[row][2] for row, _i in first] + accepted = sender.submit(request_ids, pcm, sr, valid=frame_valid) + self.owner._first_audio_requests.update(accepted) diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 9c56790e8..ca1e76cfc 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -32,6 +32,7 @@ from vllm.v1.worker.gpu.states import RequestState from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState from vllm_omni.worker_v2.model_states.intermediate_buffer import ( OmniIntermediateBuffer, ) @@ -85,6 +86,14 @@ class OmniModelState(DefaultModelState): Model-owned preprocess and postprocess hooks supply stage-specific behavior. """ + # Eager-frame defaults for states built without __init__ (tests, adapters). + _eager_mtp = False + _eager_fastpath = False + _eager_rows: tuple[InputBatch, list[tuple[int, int, str, bool]], torch.Tensor] | None = None + _first_audio_stream: torch.cuda.Stream | None = None + # Set by the stage engine process when it can take one-request outputs directly. + _first_audio_sender: Any = None + def __init__( self, vllm_config: VllmConfig, @@ -167,6 +176,40 @@ class OmniModelState(DefaultModelState): device=device, ) self._mtp_runner = self._init_mtp_runner(model) + # Eager frames: finish MTP in the step that sampled CB0 (see + # run_eager_mtp). The next step only adds its text-step embedding. + self._eager_mtp = self._mtp_input_ids is not None and bool(getattr(model, "mtp_eager_frames", False)) + self._eager_rows = None + self._eager_state = EagerMTPState(self) + self._eager_embeds: torch.Tensor | None = None + self._eager_ready: dict[int, str] = {} + self._first_audio_requests: set[str] = set() + # Slots whose decode preprocess no longer changes state (see + # model.eager_decode_settled); they skip the per-row path. + self._eager_settled: dict[int, str] = {} + self._eager_fastpath = self._eager_mtp + if self._eager_mtp: + self._eager_embeds = torch.zeros((max_num_reqs, self._embed_dim), dtype=self.dtype, device=device) + # The model then leaves frame codes/validity and the decode MTP + # inputs to run_eager_mtp instead of building them per row. + model.eager_frames_active = True + logger.info("Eager Talker-MTP frames enabled") + + def set_first_audio_sink(self, sink: Any) -> None: + self._eager_state.set_first_audio_sink(sink) + + def run_eager_mtp( + self, + input_batch: InputBatch, + text_hidden: torch.Tensor, + sampled_token_ids: torch.Tensor, + multimodal_outputs: dict[str, Any], + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None, + ) -> None: + if self._eager_mtp: + self._eager_state.run_eager_mtp( + input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher + ) @staticmethod def _resolve_decode_preprocess(model: nn.Module) -> Callable | None: @@ -350,6 +393,9 @@ class OmniModelState(DefaultModelState): req_id = self.intermediate_buffer.buffers[req_index].get("req_id") if req_id is not None: getattr(self, "_mtp_generators", {}).pop(req_id, None) + getattr(self, "_first_audio_requests", set()).discard(req_id) + getattr(self, "_eager_ready", {}).pop(req_index, None) + getattr(self, "_eager_settled", {}).pop(req_index, None) self.intermediate_buffer.remove_request(req_index) # ------------------------------------------------------------------ @@ -555,15 +601,19 @@ class OmniModelState(DefaultModelState): self._stage_batched_preprocess_inputs(req_indices, embeds.device) preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = [] + settled = self._eager_settled if self._eager_fastpath else None + settled_rows: list[tuple[int, int, int, str]] = [] for i, req_idx in enumerate(req_indices): buf = self.intermediate_buffer.buffers[req_idx] if not buf or "req_id" not in buf: continue - if str(buf["req_id"]).startswith("_warmup_"): - continue - start = int(input_batch.query_start_loc_np[i]) n_tok = int(input_batch.num_scheduled_tokens[i]) + if settled and n_tok == 1 and settled.get(req_idx) == buf["req_id"]: + settled_rows.append((i, req_idx, start, buf["req_id"])) + continue + if str(buf["req_id"]).startswith("_warmup_"): + continue info = {key: value for key, value in buf.items() if isinstance(key, str)} prompt_len = None @@ -590,6 +640,20 @@ class OmniModelState(DefaultModelState): ) preprocess_entries.append((i, req_idx, start, n_tok, info, is_prefill)) + if self._eager_mtp: + # Rows whose sample is kept this step: decode, or the final prefill chunk. + self._eager_rows = ( + input_batch, + [ + (i, req_idx, str(info["req_id"]), is_prefill) + for i, req_idx, _start, n_tok, info, is_prefill in preprocess_entries + if not is_prefill + or int(info.get("_omni_num_computed_tokens", 0)) + n_tok >= int(info.get("_omni_prompt_len", 0)) + ] + + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows], + input_ids, + ) + preprocess_batch_mrv2 = getattr(self.model, "preprocess_batch_mrv2", None) if callable(preprocess_batch_mrv2): prefill_infos = [ @@ -620,6 +684,7 @@ class OmniModelState(DefaultModelState): batch_embeds = embeds[:batch_size] batch_offsets = None else: + # A pageable host tensor would block on the stream here. batch_offsets = torch.as_tensor(starts, device=input_ids.device, dtype=torch.long) batch_ids = input_ids.index_select(0, batch_offsets) batch_embeds = embeds.index_select(0, batch_offsets) @@ -659,13 +724,13 @@ class OmniModelState(DefaultModelState): ) for row, (entry, updates) in enumerate(zip(decode_entries, updates_by_req, strict=True)): i, req_idx, start, _n_tok, _info, _is_prefill = entry - mtp_batches.append( - ( - i, - start, - (batch_hidden[row : row + 1], batch_text_step[row : row + 1]), - ) + # Eager frames read the packed text steps; skip per-row views. + row_inputs = ( + (batch_hidden, batch_text_step) + if self._eager_mtp + else (batch_hidden[row : row + 1], batch_text_step[row : row + 1]) ) + mtp_batches.append((i, start, row_inputs)) batched_decode_indices.add(i) state_updates = [(entry[1], updates) for entry, updates in zip(decode_entries, updates_by_req, strict=True)] for req_idx, updates in state_updates: @@ -694,7 +759,18 @@ class OmniModelState(DefaultModelState): self.intermediate_buffer.update(req_idx, updates, gpu_keys) - if mtp_batches and hasattr(self.model, "mtp"): + if self._eager_mtp: + settled_check = getattr(self.model, "eager_decode_settled", None) + if self._eager_fastpath and callable(settled_check): + for _i, req_idx, _start, n_tok, info, is_prefill in preprocess_entries: + buf = self.intermediate_buffer.buffers[req_idx] + if not is_prefill and n_tok == 1 and settled_check(buf): + self._eager_settled[req_idx] = buf["req_id"] + if settled_rows: + self._eager_state._apply_settled_frames(settled_rows, embeds) + if mtp_batches and self._eager_mtp: + self._eager_state._apply_eager_frames(mtp_batches, embeds, input_batch, prepacked_mtp_inputs) + elif mtp_batches and hasattr(self.model, "mtp"): if prepacked_mtp_inputs is None: self._run_batched_mtp( mtp_batches, @@ -727,9 +803,7 @@ class OmniModelState(DefaultModelState): bsz = len(mtp_batches) if offsets is None: offsets = torch.as_tensor( - [start for _i, start, _mtp in mtp_batches], - device=input_ids.device, - dtype=torch.long, + [start for _i, start, _mtp in mtp_batches], device=input_ids.device, dtype=torch.long ) if prepacked_mtp_inputs is None: hidden_rows = [past_hidden.reshape(1, -1) for _i, _start, (past_hidden, _step) in mtp_batches] @@ -802,11 +876,7 @@ class OmniModelState(DefaultModelState): offsets.copy_(query_start_loc[:bsz], non_blocking=True) return offsets - return torch.as_tensor( - [start for _i, start, _mtp in mtp_batches], - device=device, - dtype=torch.long, - ) + return torch.as_tensor([start for _i, start, _mtp in mtp_batches], device=device, dtype=torch.long) def _run_batched_mtp( self, @@ -823,8 +893,6 @@ class OmniModelState(DefaultModelState): Uses pre-allocated static buffers to avoid per-step torch.cat memory allocations. """ - from vllm.forward_context import set_forward_context - bsz = len(mtp_batches) batch_offsets = self._mtp_batch_offsets( mtp_batches, @@ -840,6 +908,72 @@ class OmniModelState(DefaultModelState): ) req_indices = [int(input_batch.idx_mapping_np[i]) for i, _start, _mtp in mtp_batches] + new_emb, codes = self._mtp_forward( + req_indices, + batch_ids, + batch_emb, + batch_hidden, + batch_step, + mtp_batch_descriptor_dispatcher, + ) + + embeds.index_copy_(0, batch_offsets, new_emb[:bsz].reshape(bsz, -1)) + audio_key = getattr(self.model, "mtp_output_key", None) + validity_key = getattr(self.model, "mtp_validity_key", None) + valid_rows = None + if codes is not None and validity_key is not None: + valid_rows = torch.ones((bsz,), dtype=torch.bool, device=codes.device) + if codes is not None and audio_key in gpu_keys: + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + audio_key, + codes[:bsz], + ) + if valid_rows is not None: + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + validity_key, + valid_rows, + keepdim=False, + ) + return + for j, (i, _start, _) in enumerate(mtp_batches): + if codes is None: + continue + req_idx = int(input_batch.idx_mapping_np[i]) + if isinstance(audio_key, tuple) and len(audio_key) == 2: + updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}} + elif isinstance(audio_key, str): + updates = {audio_key: codes[j : j + 1]} + else: + raise TypeError( + f"mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}" + ) + if valid_rows is not None: + if isinstance(validity_key, tuple) and len(validity_key) == 2: + updates.setdefault(validity_key[0], {})[validity_key[1]] = valid_rows[j] + elif isinstance(validity_key, str): + updates[validity_key] = valid_rows[j] + else: + raise TypeError( + "mtp_validity_key must be a string or 2-tuple, " + f"got {type(validity_key).__name__}: {validity_key!r}" + ) + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + def _mtp_forward( + self, + req_indices: list[int], + batch_ids: torch.Tensor, + batch_emb: torch.Tensor, + batch_hidden: torch.Tensor, + batch_step: torch.Tensor, + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run the residual-codebook predictor on packed rows; returns graph-owned outputs.""" + from vllm.forward_context import set_forward_context + + bsz = len(req_indices) buffers = [self.intermediate_buffer.buffers[req_idx] for req_idx in req_indices] req_ids = [str(buffer.get("req_id")) for buffer in buffers] generators = [ @@ -894,50 +1028,7 @@ class OmniModelState(DefaultModelState): generators=generators, sample_uniforms=sample_uniforms, ) - - embeds.index_copy_(0, batch_offsets, new_emb[:bsz].reshape(bsz, -1)) - audio_key = getattr(self.model, "mtp_output_key", None) - validity_key = getattr(self.model, "mtp_validity_key", None) - valid_rows = None - if codes is not None and validity_key is not None: - valid_rows = torch.ones((bsz,), dtype=torch.bool, device=codes.device) - if codes is not None and audio_key in gpu_keys: - self.intermediate_buffer.update_gpu_tensor_rows( - req_indices, - audio_key, - codes[:bsz], - ) - if valid_rows is not None: - self.intermediate_buffer.update_gpu_tensor_rows( - req_indices, - validity_key, - valid_rows, - keepdim=False, - ) - return - for j, (i, _start, _) in enumerate(mtp_batches): - if codes is None: - continue - req_idx = int(input_batch.idx_mapping_np[i]) - if isinstance(audio_key, tuple) and len(audio_key) == 2: - updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}} - elif isinstance(audio_key, str): - updates = {audio_key: codes[j : j + 1]} - else: - raise TypeError( - f"mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}" - ) - if valid_rows is not None: - if isinstance(validity_key, tuple) and len(validity_key) == 2: - updates.setdefault(validity_key[0], {})[validity_key[1]] = valid_rows[j] - elif isinstance(validity_key, str): - updates[validity_key] = valid_rows[j] - else: - raise TypeError( - "mtp_validity_key must be a string or 2-tuple, " - f"got {type(validity_key).__name__}: {validity_key!r}" - ) - self.intermediate_buffer.update(req_idx, updates, gpu_keys) + return new_emb, codes def _get_mtp_base_sampling_kwargs(self) -> dict[str, Any]: sampling_params = getattr(self.model, "mtp_sampling_params", None) @@ -1107,6 +1198,9 @@ class OmniModelState(DefaultModelState): """ if not self.has_postprocess: return + if self._eager_fastpath and not getattr(self.model, "eager_frames_need_postprocess", True): + # The last hidden it stores only fed the deferred decode MTP. + return gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set()) batch_postprocess = getattr(self.model, "postprocess_batch_mrv2", None) if callable(batch_postprocess) and input_batch.num_reqs: diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 8ff6e02f3..0f020c47b 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -192,6 +192,15 @@ class OmniARModelRunner(OmniGPUModelRunner): input_batch, grammar_output, ) + run_eager_mtp = getattr(self.model_state, "run_eager_mtp", None) + if multimodal_outputs and run_eager_mtp is not None: + run_eager_mtp( + input_batch, + text_hidden, + sampler_output.sampled_token_ids, + multimodal_outputs, + self._dispatch_mtp_batch_descriptor, + ) if self.pp_handler is not None: self.pp_handler.broadcast( sampler_output.sampled_token_ids, diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 33545df36..4eed3f84c 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -186,6 +186,9 @@ class OmniGPUModelRunner(GPUModelRunner): super().add_requests(scheduler_output) def shutdown(self) -> None: + sender = getattr(getattr(self, "model_state", None), "_first_audio_sender", None) + if sender is not None: + sender.close() # Every owner must release its resources even if an earlier drain fails. try: worker = getattr(self, "_native_output_materializer", None) @@ -352,6 +355,9 @@ class OmniGPUModelRunner(GPUModelRunner): capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None) if callable(capture_mtp): capture_mtp(self._dispatch_mtp_batch_descriptor) + capture_first_frame = getattr(self.model, "capture_first_frame_graphs", None) + if callable(capture_first_frame): + capture_first_frame() return result def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any: