diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1085b4e64..2cf7327dd 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles: | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) | | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning | | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback | +| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS | | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control | +| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU | ```bash # MRV2 is the default; pass a copy with `model_runner: v1` to force V1. @@ -350,12 +352,73 @@ 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. +### Selecting one or two Qwen3-TTS stages + +Choose the deployment profile explicitly with `--deploy-config`. The default +`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to +Code2Wav, which produces audio. Its stage devices can be placed on the same +GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for +one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly, +without a Code2Wav engine or stage connector. + +```bash +# Two stages on one GPU; the two-stage pipeline remains independently selectable. +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml + +# One stage on one GPU; no inter-stage transport or Code2Wav scheduling. +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml +``` + +The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1, +an in-process worker and disabled prefix caching. Select the full profile; +adding `talker_stream_decode: true` to a two-stage Talker is rejected because +its latent output belongs to the Code2Wav input contract. Single-stage options +live in `stages[].additional_config`; two-stage transport and first-frame +options live in `connectors.*.extra`. + +The single-stage profile defaults `talker_stream_first_audio` to `false` for +throughput. To opt into a separate first-frame delivery queue, save this +overlay beside the bundled deployment files and pass its path to +`--deploy-config`: + +```yaml +base_config: qwen3_tts_fused_single_gpu.yaml +stages: + - stage_id: 0 + additional_config: + talker_stream_decode: true + talker_stream_first_audio: true + ref_code_context_frames: 72 + code_predictor_kv_cache: true + code_predictor_fused_sampling: true + code_predictor_fused: true +``` + +This queue uses PCM from the single-stage decoder. The two-stage +`talker_first_audio` option below uses its own first-frame decoder and does +not require single-stage execution. The two options cannot be enabled +together. Single-stage prefix-cache reset with +`reset_running_requests=True` returns `False` while requests are active; +wait for completion or abort them before retrying. + +Reference-cloning context is selected independently by each profile. The +single-stage profile explicitly uses 72 reference frames, matching the +`qwen3_tts.yaml` context boundary. A shorter +`additional_config.ref_code_context_frames` (for example 25) can change +timbre continuity even when WER is similar; +the throughput measurements for CustomVoice do not qualify Base voice +cloning quality. Reference layout, grouping and priming are model hooks; +the shared worker does not select Qwen's context length or codebook layout. + ### 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: +two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` 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. diff --git a/tests/config/test_config_merge.py b/tests/config/test_config_merge.py index 0bcc15b0b..54b9f4816 100644 --- a/tests/config/test_config_merge.py +++ b/tests/config/test_config_merge.py @@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path): merged = resolve_deploy_yaml(leaf) assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}} assert merged["async_chunk"] is False + + +def test_single_stage_streaming_capability_uses_central_async_resolver(): + from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled + from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE + + assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig()) + with pytest.raises(ValueError, match="requires async_chunk=True"): + _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False)) + + +@pytest.mark.parametrize( + "pipeline_key,profile,stage_count", + [ + ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2), + ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1), + ], +) +def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count): + from tests.helpers.stage_config import get_deploy_config_path + from vllm_omni.config.pipeline_registry import resolve_pipeline_config + from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy + + pipeline = resolve_pipeline_config(pipeline_key) + assert pipeline is not None + deploy = load_deploy_config(get_deploy_config_path(profile)) + stages = merge_pipeline_deploy(pipeline, deploy) + assert len(stages) == stage_count + talker = stages[0] + if stage_count == 1: + assert talker.final_output and talker.final_output_type == "audio" + assert talker.yaml_engine_args["engine_output_type"] == "audio" + assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args + assert not talker.yaml_extras.get("output_connectors") + assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False + assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72 + else: + assert not talker.final_output + assert talker.yaml_engine_args["engine_output_type"] == "latent" + assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk") + assert talker.yaml_extras["output_connectors"] + assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio" + assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True) + + +def test_single_stage_running_reset_capability_cannot_be_overridden(): + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy + from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE + + deploy = DeployConfig( + stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})] + ) + legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0] + legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True + assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False + with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"): + VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy) + structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE) + assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False + + +@pytest.fixture +def single_stage_tts_config(mocker): + from vllm.config import VllmConfig + + config = mocker.Mock(spec=VllmConfig) + config.additional_config = {"talker_stream_decode": True} + config.model_config = mocker.Mock( + stage_connector_config={}, + engine_output_type="audio", + final_output=True, + supports_running_prefix_cache_reset=False, + use_v2_model_runner=True, + async_chunk=True, + ) + config.parallel_config = mocker.Mock( + tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni" + ) + config.cache_config = mocker.Mock(enable_prefix_caching=False) + config.device_config = mocker.Mock(device="cuda") + mocker.patch( + "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True + ) + return config + + +def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.cache_config.enable_prefix_caching = True + with pytest.raises(ValueError, match="enable_prefix_caching=False"): + talker_stream_decode_enabled(config) + config.cache_config.enable_prefix_caching = False + assert talker_stream_decode_enabled(config) + + +@pytest.mark.parametrize("option_location", ["additional_config", "connector"]) +@pytest.mark.parametrize("output_type", ["latent", "audio"]) +def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.model_config.engine_output_type = output_type + config.model_config.final_output = False + if option_location == "connector": + config.model_config.stage_connector_config = {"extra": config.additional_config} + config.additional_config = {} + with pytest.raises(ValueError, match="final audio-output Talker"): + talker_stream_decode_enabled(config) + + +def test_two_stage_first_audio_remains_independent(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.model_config.engine_output_type = "latent" + config.additional_config = {} + config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}} + assert not talker_stream_decode_enabled(config) + assert talker_first_audio_enabled(config) + config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True + with pytest.raises(ValueError, match="requires the single-stage"): + talker_stream_decode_enabled(config) + + +def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.device_config.device = "cpu" + with pytest.raises(ValueError, match="requires CUDA"): + talker_stream_decode_enabled(config) + config.device_config.device = "cuda" + config.additional_config["talker_first_audio"] = True + with pytest.raises(ValueError, match="mutually exclusive"): + talker_stream_decode_enabled(config) + config.additional_config.pop("talker_first_audio") + config.additional_config["talker_stream_first_audio"] = True + assert talker_stream_decode_enabled(config) + + +def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + single_stage_tts_config.additional_config = {} + with pytest.raises(ValueError, match="requires talker_stream_decode=True"): + talker_stream_decode_enabled(single_stage_tts_config) + + +def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True + with pytest.raises(ValueError, match="must disable running prefix-cache reset"): + talker_stream_decode_enabled(single_stage_tts_config) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index b19c97144..e99da6ceb 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes(): "interleave_mm_strings", "media_io_kwargs", "final_output", + "supports_running_prefix_cache_reset", "active_stream_window", "session_mode", "duplex_max_sessions", diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index 182857ea9..b4eb839a2 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -106,7 +106,7 @@ def _make_drain_sched(session: Request) -> MagicMock: return sched -def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool: +def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool: """Feed one model-runner frame through update_from_output. Returns True when the frame's tokens were delivered (the append ran), @@ -114,7 +114,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: """ scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled} - scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {} scheduler_output.num_invalid_spec_tokens = 0 model_runner_output = MagicMock(spec=ModelRunnerOutput) @@ -134,6 +134,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: return sched._update_request_with_output.called +def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None: + from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + + session = _make_session() + session.num_computed_tokens = 6 + session.num_in_flight_tokens = 1 + session.num_output_placeholders = 1 + sched = _make_drain_sched(session) + sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id} + sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id} + sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2} + + OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched) + assert session.num_stale_output_tokens == 1 and session.drop_stale_output + assert session.num_output_placeholders == 0 + sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session) + assert _run_step(sched, session, num_scheduled=1, token=42) is False + assert session.num_stale_output_tokens == 0 + assert _run_step(sched, session, num_scheduled=1, token=43) is True + + def test_exact_drain_delivers_new_segment_frame() -> None: """One in-flight decode frame at replacement: its late output drains the counter exactly and the new segment's first frame is delivered.""" @@ -225,3 +246,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None: assert _run_step(sched, session, num_scheduled=5, token=42) is False # late prefill chunk dropped assert session.num_stale_output_tokens == 0 assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives + + +@pytest.mark.parametrize("drop", [False, True]) +def test_preemption_stale_output_follows_upstream_delivery_policy(drop): + request = _make_session() + request.status = RequestStatus.PREEMPTED + request.num_stale_output_tokens = 1 + request.num_in_flight_tokens = 1 + request.drop_stale_output = drop + sched = _make_drain_sched(request) + assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop) + assert request.num_stale_output_tokens == 0 + if not drop: + sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True) + + +def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters(): + request = _make_session() + request.num_stale_output_tokens = 2 + request.num_in_flight_tokens = 2 + request.num_computed_tokens = 3 + request.num_output_placeholders = 2 + request.drop_stale_output = False + sched = _make_drain_sched(request) + assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41]) + assert request.num_computed_tokens == 3 + assert request.num_output_placeholders == 2 + + +@pytest.mark.parametrize( + "arch,supports_reset,running,reset_running,blocked", + [ + ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True), + ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False), + ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False), + ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False), + ("AnotherStatefulAudioModel", False, True, True, True), + ("AnotherStatefulAudioModel", True, True, True, False), + ], +) +def test_reset_running_streaming_codec_is_rejected_before_preemption( + mocker, arch, supports_reset, running, reset_running, blocked +): + from vllm.config import VllmConfig + from vllm.v1.core.sched.scheduler import Scheduler + + from vllm_omni.config.model import OmniModelConfig + + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.vllm_config = mocker.Mock( + spec=VllmConfig, + model_config=mocker.Mock( + spec=OmniModelConfig, + model_arch=arch, + supports_running_prefix_cache_reset=supports_reset, + ), + ) + request = _make_session() + request.num_in_flight_tokens = 1 + request.num_output_placeholders = 1 + scheduler.running = [request] if running else [] + + def reset_cache(_self, reset_running_requests=False, reset_connector=False): + assert reset_running_requests is reset_running and reset_connector is True + return True + + reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache) + + assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked) + if blocked: + reset.assert_not_called() + assert scheduler.running == [request] + assert request.status == RequestStatus.RUNNING + assert request.num_in_flight_tokens == request.num_output_placeholders == 1 + assert not request.drop_stale_output + else: + reset.assert_called_once() diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index 5e467254b..bc99cf7fb 100644 --- a/tests/engine/test_orchestrator_first_audio.py +++ b/tests/engine/test_orchestrator_first_audio.py @@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)): ) -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) +async def _codec_output(obj, raw, stage=1): + replica = obj.stage_pools[stage].replica_id + processed = await obj._process_llm_stage_outputs(stage, replica, raw, set()) + await obj._handle_processed_outputs(stage, replica, processed) def _messages(obj): @@ -132,6 +132,57 @@ def _audio(output): return torch.cat(audio) if isinstance(audio, list) else audio +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", list(RequestOutputKind)) +@pytest.mark.parametrize("suffix_overtook", [False, True]) +async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook): + obj = _orchestrator(kind, final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + suffix = _raw(required=True, samples=(20, 21)) + terminal = _raw(required=True, terminal=True, samples=(30, 31)) + if suffix_overtook: + await _codec_output(obj, suffix, stage=0) + await _codec_output(obj, terminal, stage=0) + assert obj.output_async_queue.empty() + await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11))) + if not suffix_overtook: + await _codec_output(obj, suffix, stage=0) + await _codec_output(obj, terminal, stage=0) + messages = _messages(obj) + 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 [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind] + assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages) + assert messages[-1].finished + assert not obj.request_states["r"].pending_first_audio_outputs + + +@pytest.mark.asyncio +async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio(): + obj = _orchestrator(final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + await _codec_output(obj, _raw(required=True, terminal=True), stage=0) + assert obj.request_states["r"].pending_first_audio_outputs + await obj._cleanup_request_ids(["r"], abort=True) + await obj._route_upstream_first_audio(0, 2, _raw(first=True)) + assert not obj.request_states and obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal(): + obj = _orchestrator(final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + await _codec_output(obj, _raw(required=True, terminal=True), stage=0) + obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) + await obj._route_upstream_first_audio(0, 2, _raw(first=True)) + error = obj.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and error.stage_id == 0 + assert not obj.request_states + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", list(RequestOutputKind)) @pytest.mark.parametrize("codec_overtook", [False, True]) diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index a7104b14d..217f13e53 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -1,9 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace from unittest.mock import patch +import pytest from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.executor.uniproc_executor import UniProcExecutor -from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc +from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "first_decoder,stream_decoder,stream_first_audio,expected", + [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)], +) +def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder( + mocker, + first_decoder, + stream_decoder, + stream_first_audio, + expected, +): + executor = UniProcExecutor.__new__(UniProcExecutor) + executor.driver_worker = mocker.Mock() + runner = executor.driver_worker.worker.model_runner + runner.model.first_frame_decoder = object() if first_decoder else None + runner.model.stream_decoder = object() if stream_decoder else None + runner.model.stream_first_audio = stream_first_audio + assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected + assert runner.model_state.set_first_audio_sink.call_count == int(expected) def test_preprocess_add_request_preserves_omni_fields(): diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index 2fb76e3f6..3eda688c0 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt): assert audio_file.samplerate == 22050 else: assert len(raw) == frame_count * 2 + + +def test_float32_pcm_fast_path_matches_soundfile(): + from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes + + rng = np.random.default_rng(0) + audio = np.concatenate( + [ + rng.uniform(-1.2, 1.2, 4096), + np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]), + ] + ).astype(np.float32) + with BytesIO() as buffer: + soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16") + expected = buffer.getvalue() + assert _float32_to_pcm16_bytes(audio) == expected 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 52712a29c..8df20d87c 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 @@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators: ) +@pytest.mark.parametrize("fused_requested", [False, True]) +@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"]) +def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input): + predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes) + predictor._fused_requested = fused_requested + predictor._fused = mocker.Mock() + setup = mocker.patch.object(predictor, "_setup_compile") + uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5) + generators = [torch.Generator()] if invalid_input == "generators" else None + with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"): + predictor( + torch.zeros(2, dtype=torch.long), + torch.randn(2, talker_config.hidden_size), + torch.randn(2, talker_config.hidden_size), + sample_uniforms=uniforms, + generators=generators, + ) + setup.assert_not_called() + predictor._fused.assert_not_called() + + class TestCodePredictorModelDtype: """Test the inner model forward with different dtypes.""" @@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay: predictor = object.__new__(code_predictor_wrapper) torch.nn.Module.__init__(predictor) + predictor._fused_requested = False + predictor._fused = None predictor._num_groups = 3 predictor._model_dtype = torch.float32 predictor._setup_compile = mocker.Mock() 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 index 30994e36b..9fa3a76bb 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Supplied-uniform samples must match the existing ATen path exactly.""" +from types import SimpleNamespace + import pytest import torch @@ -12,6 +14,79 @@ pytestmark = [ ] +@pytest.mark.parametrize("top_k", [0, 1, 50, 2048]) +@torch.inference_mode() +def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k): + from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel + + vocab, hidden, groups = 2048, 64, 4 + logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16) + logits[0, 1] = -0.0 + logits[1, ::2] = -0.0 + uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :] + uniforms[0, 1] = uniforms[1, 2] = 0.99 + codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64) + table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16) + next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype) + + def sample(): + _cp_sample_kernel[(2,)]( + logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0, + V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4, + ) # fmt: skip + + scaled = logits.float() + if top_k > 0: + scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf) + expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1) + assert expected.tolist() == [1, 2] + sample() + torch.testing.assert_close(codes[:, 1], expected) + torch.testing.assert_close(next_input, table[expected]) + assert (codes[:, [0, 2, 3]] == -1).all() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + sample() + codes.fill_(-1) + graph.replay() + torch.testing.assert_close(codes[:, 1], expected) + torch.testing.assert_close(next_input, table[expected]) + + +@torch.inference_mode() +def test_fused_predictor_warmup_with_capacity_one(): + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerCodePredictorConfig, + Qwen3TTSTalkerConfig, + ) + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, + ) + + cp = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4, + ) # fmt: skip + talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4) + config = SimpleNamespace( + model_config=SimpleNamespace(stage_connector_config={}), + additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True}, + scheduler_config=SimpleNamespace(max_num_seqs=1), + ) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM( + vllm_config=config, config=cp, talker_config=talker + ).to(device="cuda", dtype=torch.bfloat16) + predictor._setup_compile() + assert predictor._fused is not None + hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16) + result = predictor( + torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden, + sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"), + ) # fmt: skip + assert result.shape == (1, 4) + assert ((result >= 0) & (result < 64)).all() + + @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]) @@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k): 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) + + +@pytest.mark.parametrize("batch", [1, 4]) +@pytest.mark.parametrize("top_k", [1, 50]) +@torch.inference_mode() +def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k): + from vllm.config import VllmConfig + + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerCodePredictorConfig, + Qwen3TTSTalkerConfig, + ) + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, + ) + + torch.manual_seed(42) + cp = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4, + ) # fmt: skip + config = mocker.Mock( + spec=VllmConfig, + model_config=mocker.Mock(stage_connector_config={}), + additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True}, + scheduler_config=mocker.Mock(max_num_seqs=4), + ) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM( + vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4) + ).to(device="cuda", dtype=torch.bfloat16) + predictor._setup_compile() + assert predictor._fused is not None + ids = torch.randint(0, 64, (batch, 1), device="cuda") + embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16) + hidden = torch.randn_like(embeds) + uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6) + actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone() + # Use identical parameters through the full re-prefill and ATen sampling + # path, bypassing both the fused predictor and its frame-local KV cache. + predictor._fused_requested = False + predictor._fused = None + predictor._frame_cache = None + expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) 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 index ce92f06a6..66fb9c10f 100644 --- a/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py +++ b/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Prefix graphs must follow per-request first-audio delivery, not a capture hint.""" +import math from types import SimpleNamespace import pytest @@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model] DEVICE = torch.device("cuda:0") +@pytest.mark.cpu +@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")]) +def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder + + captures = [] + + def capture(stream, sizes, frames=1): + captures.append((sizes, frames)) + return SimpleNamespace(sizes=sizes, frames=frames) + + monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture) + model = SimpleNamespace( + stream_decoder=SimpleNamespace(max_frames=33), + stream_graphs=None, + stream_prime_graphs=None, + stream_ref_context_frames=25, + stream_chunk_frames=25, + config=SimpleNamespace(tts_model_type=model_type), + vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)), + ) + sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity] + Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes) + expected = [(sizes, 1)] + if model_type == "base": + expected.append(([size for size in sizes if size * 25 <= 512], 25)) + assert captures == expected + Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes) + assert captures == expected + + @pytest.mark.cpu @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)]) @pytest.mark.parametrize("initial_frames", [1, 3]) @@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major 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) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_stream_decoder_preemption_restores_all_state_into_a_different_slot(): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder + + decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16) + codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE) + slot = torch.tensor([0], device=DEVICE, dtype=torch.int32) + for pos in range(7): + stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32)) + saved = stream.save_slot(0) + expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone() + # Both the freed slot and the destination can hold an unrelated request. + other_codes = (codes[:, :3] + 1) % 32 + for pos in range(3): + stream( + other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(), + torch.tensor([0, 1], device=DEVICE, dtype=torch.int32), + torch.full((2,), pos, device=DEVICE, dtype=torch.int32), + ) + without_restore = stream( + codes[:, 7:], + torch.tensor([1], device=DEVICE, dtype=torch.int32), + torch.tensor([7], device=DEVICE, dtype=torch.int32), + ).clone() + assert not torch.equal(without_restore, expected) + stream.restore_slot(1, saved) + actual = stream( + codes[:, 7:], + torch.tensor([1], device=DEVICE, dtype=torch.int32), + torch.tensor([7], device=DEVICE, dtype=torch.int32), + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("use_graph", [False, True]) +@torch.inference_mode() +def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import ( + StreamingCodecDecoder, + StreamingDecodeGraphs, + ) + + decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16) + graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None + model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + torch.nn.Module.__init__(model) + model.stream_decoder = stream + model.stream_prime_graphs = graphs + model.stream_chunk_frames = 25 + refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]] + slots = [2, 0, 3, 4] + frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE) + slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32) + positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32) + for slot, ref in zip(slots, refs, strict=True): + for t0 in range(0, len(ref), 25): + stream( + ref[None, t0 : t0 + 25].to(torch.int32).contiguous(), + torch.tensor([slot], device=DEVICE, dtype=torch.int32), + torch.tensor([t0], device=DEVICE, dtype=torch.int32), + ) + expected = stream(frame, slot_tensor, positions).clone() + untouched = stream.save_slot(1) + # Reuse dirty slots; pos=0 must reset every causal layer for each group. + model.prime_stream_decoder(list(zip(slots, refs, strict=True))) + actual = stream(frame, slot_tensor, positions) + torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5) + for got, want in zip(stream.save_slot(1), untouched, strict=True): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("use_graph", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@torch.inference_mode() +def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import ( + StreamingCodecDecoder, + StreamingDecodeGraphs, + ) + + decoder = _make_decoder().to(device=DEVICE, dtype=dtype) + # The tiny HF default initialization shrinks the deep conv stack to + # near-zero PCM. Keep the oracle signal large enough to catch lost state. + for module in decoder.modules(): + if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)): + torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5)) + if module.bias is not None: + module.bias.uniform_(-0.01, 0.01) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + reference = _make_decoder().to(device=DEVICE, dtype=torch.float32) + reference.load_state_dict(decoder.state_dict()) + reference.precompute_snake_caches() + # Cross both the attention window and KV ring wrap, checking every frame. + codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE) + expected = reference._forward_exact(codes.transpose(1, 2))[:, 0] + assert expected.square().sum() > 1e-3 + stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype) + decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream + slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32) + spf = int(stream.spf) + for frame in range(codes.shape[1]): + positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32) + actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1) + exact_frame = expected[:, frame * spf : (frame + 1) * spf] + relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt() + # BF16 kernels are approximate; compare every frame to the independent + # FP32 full decoder, not only one streaming implementation to itself. + tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4 + assert torch.all(relative_rms < tolerance), (frame, relative_rms) diff --git a/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py b/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py new file mode 100644 index 000000000..c7f677821 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Streaming col2im preserves PCM inputs and ping-pong slot state.""" + +import pytest +import torch +from vllm.triton_utils import triton + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel + +pytestmark = [pytest.mark.core_model, pytest.mark.tts] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)]) +def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int): + torch.manual_seed(23) + capacity = batch + 2 + z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16) + bias = torch.randn(channels).to(torch.bfloat16) + state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16) + expected_state = state.clone() + slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32) + positions = torch.arange(batch, dtype=torch.int32) + parity = torch.arange(capacity, dtype=torch.int32) % 2 + expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16) + for b in range(batch): + slot = int(slots[b]) + side = int(parity[slot]) + for t in range(frames): + previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels) + expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16) + expected_state[1 - side, slot] = z[b, -1, rate:] + + z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda() + slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda() + actual = torch.empty_like(expected, device="cuda") + width = min(128, triton.next_power_of_2(channels)) + rows = batch * frames + _col2im_kernel[(rows, rate, triton.cdiv(channels, width))]( + z_gpu, + bias_gpu, + state_gpu, + slots_gpu, + positions_gpu, + parity_gpu, + actual, + frames, + capacity, + R=rate, + C=channels, + BC=width, + num_warps=4, + ) + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index f929e41c9..8277ba435 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -594,6 +594,40 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes(): torch.testing.assert_close(frames[:3], ref_code[-3:]) +@pytest.mark.parametrize("tensor_frames", [False, True]) +def test_async_chunks_preserve_codebook_order_with_reference_and_short_tail(tensor_frames): + tm = _tm(left_context=3, initial_chunk_frames=1) + rid = "r-codebook-order" + quantizers = 16 + reference = [[1500 + f * quantizers + q for q in range(quantizers)] for f in range(5)] + tm.request_payload[rid] = torch.tensor(reference, dtype=torch.long) + generated = [[1 + f * quantizers + q for q in range(quantizers)] for f in range(30)] + + for start, end, finished in [(0, 1, False), (1, 26, False), (26, 30, True)]: + rows = generated[:end] + if tensor_frames: + # Non-contiguous int32 rows exercise conversion independently of + # the native producer's contiguous CPU-long representation. + storage = torch.tensor(rows, dtype=torch.int32).repeat_interleave(2, dim=1) + tm.code_prompt_token_ids[rid] = list(storage[:, ::2].unbind()) + else: + tm.code_prompt_token_ids[rid] = rows + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output=None if finished else {}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + assert payload is not None + expected = (reference[-3:] if start == 0 else []) + generated[start:end] + codes = payload.codes.audio + assert codes.dtype == torch.long and codes.device.type == "cpu" and codes.is_contiguous() + torch.testing.assert_close(codes.reshape(quantizers, -1).T, torch.tensor(expected)) + assert payload.meta.left_context_size == (3 if start == 0 else 0) + assert payload.meta.finished.item() == finished + tm.put_req_chunk[rid] += 1 + + def test_ref_code_context_can_be_buffered_before_first_emit(): tm = _tm() rid = "r-ref-buffered" diff --git a/tests/worker_v2/test_index_staging.py b/tests/worker_v2/test_index_staging.py new file mode 100644 index 000000000..81d97701d --- /dev/null +++ b/tests/worker_v2/test_index_staging.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from vllm_omni.utils.device_copy import index_to_device + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +def test_index_staging_survives_later_copies_and_graph_replay(): + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + captured = index_to_device([17] * 6, "cuda") + for _ in range(32): + index_to_device([99] * 6, "cuda") + graph.replay() + torch.accelerator.synchronize() + assert captured.tolist() == [17] * 6 diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index e425cc9c8..0bbebf491 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -25,6 +25,9 @@ class _FakeStream: def wait_stream(self, _stream) -> None: pass + def wait_event(self, _event) -> None: + pass + class _FakeEvent: def record(self, _stream) -> None: @@ -138,6 +141,7 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: runner = OmniARModelRunner.__new__(OmniARModelRunner) runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace() runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 waited: list[object] = [] @@ -307,6 +311,231 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded): assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]]) +def test_stream_audio_cancel_rejects_inflight_output_and_id_reuse() -> None: + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + cancelled = buffer.add("request") + buffer.finish({"request"}) + replacement = buffer.add("request") + assert cancelled.push(torch.ones(4), False) is None + assert not cancelled.pending + assert replacement.push(torch.full((4,), 2.0), False).tolist() == [2.0] * 4 + buffer.finish({"request"}) + assert not buffer.requests + + +def test_stream_audio_length_end_flushes_first_and_partial_chunks() -> None: + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("request") + output = StreamingAudioOutput( + torch.ones(1, 4), + torch.tensor([True]), + torch.tensor(24000), + None, + np.array([0, 1]), + [state], + np.array([True]), + ) + assert output.get_output()[0]["model_outputs"].numel() == 4 + assert not state.active + state = buffer.add("another") + assert state.push(torch.ones(4), False).numel() == 4 + assert state.push(torch.ones(4), False) is None + assert state.push(torch.ones(4), True).numel() == 8 + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, torch.bfloat16]) +def test_stream_audio_batch_views_retain_pcm_after_output_release(dtype): + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + states = [buffer.add("a"), buffer.add("b")] + expected: list[list[torch.Tensor]] = [[], []] + emitted: list[list[torch.Tensor]] = [[], []] + for step in range(5): + # Mixed prefill/decode spans, invalid rows, and a final partial chunk. + wav = (torch.arange(16).reshape(4, 4) + 100 * step).to(dtype) + valid = torch.tensor([False, True, True, step != 4]) + output = StreamingAudioOutput( + wav, + valid, + torch.tensor(24000), + None, + np.array([0, 3, 4]), + states, + np.array([step == 4, step == 4]), + ) + for i, (start, end) in enumerate([(0, 3), (3, 4)]): + expected[i].append(wav[start:end][valid[start:end]].reshape(-1).clone()) + for i, payload in enumerate(output.get_output()): + if payload is not None: + emitted[i].append(payload["model_outputs"]) + del output, wav # Pending views must keep the old host allocation alive. + for i in range(2): + assert torch.equal(torch.cat(emitted[i]), torch.cat(expected[i])) + assert torch.cat(emitted[i]).dtype == dtype + assert not states[i].active and not states[i].pending + assert states[i].pending_samples == 0 + + +def test_stream_audio_numpy_direct_first_cancel_and_mixed_fallback(): + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("r") + state.accept_first_audio() + assert state.push(np.arange(8, dtype=np.float32), False) is None + assert state.pending_samples == 4 + # A dtype fallback can arrive while a NumPy frame is pending. + chunk = state.push(torch.arange(8, 12, dtype=torch.bfloat16), True) + assert chunk.tolist() == list(range(4, 12)) + assert state.pending_samples == 0 + buffer.finish({"r"}) + replacement = buffer.add("r") + assert state.push(np.ones(4, dtype=np.float32), False) is None + assert replacement.push(np.full(4, 99, dtype=np.float32), False).tolist() == [99] * 4 + replacement.push(np.ones(4, dtype=np.float32), False) + assert replacement.pending_samples == 4 + buffer.finish({"r"}) + assert replacement.pending_samples == 0 and not replacement.pending + + +def test_stream_audio_prefill_length_end_reaches_output(mocker) -> None: + from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.states import RequestState + + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + owner = mocker.Mock() + owner.vllm_config.model_config.max_model_len = 4096 + owner._stream_decode_event = None + eager = EagerMTPState(owner) + eager._audio_buffer = StreamingAudioBuffer(4, 25) + eager._audio_buffer.add("request") + eager._frame_requests = {"request"} + batch = InputBatch.__new__(InputBatch) + batch.num_reqs, batch.req_ids = 1, ["request"] + batch.num_computed_tokens_np = np.array([0]) + batch.num_scheduled_tokens = np.array([3]) + batch.idx_mapping_np = np.array([0]) + batch.query_start_loc_np = np.array([0, 3]) + batch.is_prefilling_np = np.array([True]) + states = RequestState.__new__(RequestState) + states.max_seq_len = np.array([4]) + mm = { + "model_outputs": torch.ones(3, 4), + "meta": {"codec_frame_valid": torch.tensor([0, 0, 1])}, + "sr": [torch.tensor(24000)], + } + output = eager.prepare_audio_output(batch, states, mm) + assert output.get_output()[0]["model_outputs"].numel() == 4 + assert "model_outputs" not in mm + + +def test_stream_audio_preemption_preserves_position_rng_and_pending_pcm(mocker) -> None: + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + owner = mocker.Mock() + owner._stream_pos = {"request": 37} + owner._eager_embeds = torch.arange(18).reshape(6, 3).float() + owner._eager_ready = {} + owner.intermediate_buffer.buffers = [{} for _ in range(6)] + generator = torch.Generator().manual_seed(42) + owner._mtp_generators = {"request": generator} + saved = [torch.arange(6)] + owner.model.stream_decoder.save_slot.return_value = saved + eager = EagerMTPState(owner) + eager._talker_inputs["request"] = TalkerInputs(torch.ones(20, 3), 10) + eager._side_stream = mocker.Mock() + eager._audio_buffer = StreamingAudioBuffer(4, 25) + request = eager._audio_buffer.add("request") + request.push(torch.ones(4), False) + request.push(torch.ones(4), False) + eager.suspend_audio("request", 2) + owner._stream_pos.clear() + owner._mtp_generators.clear() + eager.resume_audio("request", 5) + assert owner._stream_pos["request"] == 37 + assert owner._mtp_generators["request"] is generator + assert eager._restore_audio["request"] is saved + assert sum(part.numel() for part in request.pending) == 4 + # A second preemption can happen before the restored slot is used. + eager.suspend_audio("request", 5) + owner.model.stream_decoder.save_slot.assert_called_once_with(2) + eager.finish_audio({"request"}) + assert not eager._suspended_audio and not eager._restore_audio + assert not eager._audio_buffer.requests and not request.active + + +def test_stream_audio_replay_uses_exact_inputs_and_consumes_last_frame_once(mocker): + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs + + owner = mocker.Mock() + owner._eager_embeds = torch.tensor([[20.0, 30.0]]) + owner.model.preprocess.return_value = (None, None, {"mtp_inputs": (None, torch.tensor([[2.0, 3.0]]))}) + owner.intermediate_buffer.buffers = [{"req_id": "request"}] + eager = EagerMTPState(owner) + eager._restore_audio["request"] = [] + history = torch.arange(8).reshape(4, 2).float() + eager._talker_inputs["request"] = TalkerInputs(history, 4) + first = torch.empty(2, 2) + assert eager.replay_inputs("request", 0, 0, torch.ones(2), first) + torch.testing.assert_close(first, history[:2]) + owner.model.preprocess.assert_not_called() + last = torch.empty(3, 2) + assert eager.replay_inputs("request", 0, 2, torch.ones(3), last) + torch.testing.assert_close(last, torch.cat((history[2:], torch.tensor([[22.0, 33.0]])))) + owner.model.preprocess.assert_called_once() + with pytest.raises(RuntimeError, match="exceeds saved inputs"): + eager.replay_inputs("request", 0, 4, torch.ones(2), torch.empty(2, 2)) + + +def test_stream_reference_priming_groups_lengths_without_padding_or_slot_aliasing(mocker): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + + primes = [(4, torch.full((26, 2), 4)), (1, torch.full((3, 2), 1)), (7, torch.full((26, 2), 7))] + calls = [] + + def stream(codes, slots, pos): + calls.append((codes.clone(), slots.clone(), pos.clone())) + + decoder = mocker.Mock(side_effect=stream, device=torch.device("cpu")) + model = mocker.Mock(stream_decoder=decoder, stream_prime_graphs=None, stream_chunk_frames=25) + Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, primes) + assert len(calls) == 3 + for (codes, slots, pos), frames, indices, position in zip( + calls, [25, 1, 3], [[4, 7], [4, 7], [1]], [0, 25, 0], strict=True + ): + assert codes.shape == (len(indices), frames, 2) + assert codes.dtype == slots.dtype == pos.dtype == torch.int32 + assert slots.tolist() == indices and pos.tolist() == [position] * len(indices) + for row, idx in enumerate(indices): + assert torch.all(codes[row] == idx) + Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, []) + assert len(calls) == 3 + + +@pytest.mark.parametrize("context_frames", [25, 72]) +def test_model_owned_reference_context_respects_configured_boundary(mocker, context_frames): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + + model = mocker.Mock(stream_ref_context_frames=context_frames, talker_config=mocker.Mock(num_code_groups=16)) + reference = torch.arange(100 * 16).reshape(100, 16) + info = {"codes": {"ref": reference}} + got = Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, info) + torch.testing.assert_close(got, reference[-context_frames:]) + assert info["codes"]["ref"] is reference + assert Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {}) is None + assert ( + Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {"codes": {"ref": torch.empty(0)}}) is None + ) + + @pytest.mark.parametrize("streaming", [False, True]) def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming): monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) @@ -340,6 +569,136 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, assert (result.pooler_output is None) == streaming +def test_stream_audio_history_replays_exact_inputs_after_batch_reordering(mocker): + from vllm.v1.worker.gpu.input_batch import InputBatch + + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + + owner = mocker.Mock() + owner.vllm_config.model_config.max_model_len = 12 + eager = EagerMTPState(owner) + first = torch.arange(9, dtype=torch.float32).reshape(3, 3) + second = torch.arange(9, 18, dtype=torch.float32).reshape(3, 3) + expected = {"a": torch.cat((first[:2], second[2:])), "b": torch.cat((first[2:], second[:2]))} + batch = InputBatch.__new__(InputBatch) + batch.req_ids = ["a", "b"] + batch.query_start_loc_np = np.array([0, 2, 3]) + batch.num_computed_tokens_np = np.array([0, 0]) + eager.record_inputs(batch, first) + batch.req_ids = ["b", "a"] + batch.num_computed_tokens_np = np.array([1, 2]) + eager.record_inputs(batch, second) + # Later input-buffer reuse must not change the conditioned replay history. + first.fill_(-1) + second.fill_(-2) + for request_id, wanted in expected.items(): + eager._restore_audio[request_id] = [] + replayed = torch.empty_like(wanted) + assert eager.replay_inputs(request_id, 0, 0, torch.zeros(3, dtype=torch.long), replayed) + torch.testing.assert_close(replayed, wanted) + + +@pytest.mark.parametrize("valid", [False, True]) +def test_stream_audio_direct_first_requires_only_real_pcm_on_terminal(valid): + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("r") + state.accept_first_audio() + first = StreamingAudioOutput( + torch.arange(4).reshape(1, 4), + torch.tensor([valid]), + torch.tensor(24000), + None, + np.array([0, 1]), + [state], + np.array([True]), + ) + payload = first.get_output()[0] + if valid: + assert payload == {FIRST_AUDIO_REQUIRED_KEY: torch.tensor(True)} + else: + assert payload is None # EOS has no promised first PCM to wait for. + assert not state.active and not state.pending + buffer.finish({"r"}) + replacement = buffer.add("r") + assert replacement.push(torch.full((4,), 99), False).tolist() == [99] * 4 + + +def test_stream_audio_direct_first_drops_one_frame_and_retains_partial_suffix(): + from vllm_omni.worker_v2.streaming_audio import AudioRequest + + state = AudioRequest(samples_per_frame=4, chunk_frames=3) + state.accept_first_audio() + assert state.push(torch.arange(4), False) is None + assert state.first_audio_required and state.emitted + assert state.push(torch.arange(4, 8), False) is None + chunk = state.push(torch.arange(8, 12), True) + assert chunk.tolist() == list(range(4, 12)) + assert not state.active and not state.pending + + +@pytest.mark.parametrize("consumer", ["pcm_only", "finalizer", "extra"]) +def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocker, consumer): + from vllm.v1.worker.gpu.input_batch import InputBatch + + from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + monkeypatch.setattr( + OmniARModelRunner, + "_build_async_chunk_outputs_from_mm", + lambda *args: pytest.fail("in-stage PCM unnecessarily partitioned code outputs"), + ) + batch = InputBatch.__new__(InputBatch) + batch.num_reqs = 3 + batch.query_start_loc_np = np.array([0, 1, 2, 4]) + batch.num_scheduled_tokens = np.array([1, 1, 2]) + batch.num_tokens_after_padding = 4 + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + states = [buffer.add(request_id) for request_id in ["frame", "eos", "prefill"]] + wav = torch.arange(16, dtype=torch.float32).reshape(4, 4) + streaming_audio = StreamingAudioOutput( + wav, + torch.tensor([True, False, True, False]), + torch.tensor(24000), + None, + batch.query_start_loc_np.copy(), + states, + np.array([False, False, True]), + ) + copy_mm = mocker.spy(omni_ar_model_runner, "_async_copy_mm") + + def finalize(payload, counts): + assert payload["codes"]["audio"].shape == (4, 2) + return RequestOutputSnapshot([None] * batch.num_reqs) + + finalizer = mocker.Mock(side_effect=finalize) if consumer == "finalizer" else None + extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None + result = _async_output( + req_ids=["frame", "eos", "prefill"], + sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None), + num_sampled_tokens=torch.tensor([1, 1, 1]), + multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}}, + input_batch=batch, + async_chunk=True, + streaming_audio=streaming_audio, + finalize_multimodal=finalizer, + extra_multimodal_outputs=extra, + ).get_output() + assert copy_mm.call_count == {"pcm_only": 0, "finalizer": 1, "extra": 2}[consumer] + if finalizer is not None: + finalizer.assert_called_once() + wav.fill_(-1) + assert result.inter_stage_outputs is None + assert result.multimodal_outputs[0]["model_outputs"].tolist() == [0, 1, 2, 3] + assert result.multimodal_outputs[1] == {} + assert result.multimodal_outputs[2]["model_outputs"].tolist() == [8, 9, 10, 11] + assert not states[1].active and not states[2].active + + @pytest.mark.parametrize("streaming", [False, True]) def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming): from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 4edfed603..b36dbb139 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -36,6 +36,7 @@ class _DummyInputBatch: def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False): state = object.__new__(OmniModelState) model = MagicMock() + model.stream_decoder = None model.has_preprocess = has_preprocess model.has_postprocess = has_postprocess model.have_multimodal_outputs = have_multimodal_outputs @@ -62,6 +63,10 @@ def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, hav state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) state._static_inputs_embeds = None + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + + state._eager_state = EagerMTPState(state) + state._stream_pos = {} state._mtp_generators = {} state._mtp_runner = None for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"): @@ -264,6 +269,7 @@ def test_seed_independence_resolve_once_and_sampling_kwargs(): # vLLM sampling seed must not produce a talker generator. cpu = torch.device("cpu") assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None + state._mtp_generators.clear() # The following rows are new requests. # Same model-local seed reproduces identical uniforms regardless of batch makeup. state._mtp_sample_uniforms = torch.empty((2, 2, 4)) assert torch.equal( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index a56e6148c..13115f00e 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -5,6 +5,7 @@ from dataclasses import MISSING, field from typing import Any from pydantic import ConfigDict, TypeAdapter +from transformers import PretrainedConfig from vllm.config import ModelConfig from vllm.config.utils import config from vllm.logger import init_logger @@ -124,9 +125,12 @@ class OmniModelConfig(ModelConfig): """ stage_id: int = 0 + # Match the upstream derived field while its value is replaced per stage. + hf_text_config: PretrainedConfig = field(init=False) async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + supports_running_prefix_cache_reset: bool = True use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. @@ -218,7 +222,7 @@ class OmniModelConfig(ModelConfig): return convertor.convert() return super().get_model_arch_config() - def draw_hf_text_config(self): + def draw_hf_text_config(self) -> PretrainedConfig: # transformers' get_text_config method is used to get the text config from thinker_config. # to handle the case that each model stage has their own text config, # we need to draw the text config from the corresponding model stage. diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d77e44194..19d5ae71c 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -482,6 +482,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): interleave_mm_strings: bool | None = None media_io_kwargs: dict[str, Any] | None = None final_output: bool = False + supports_running_prefix_cache_reset: bool = True active_stream_window: int = Field(default=0, ge=0) session_mode: str = "turn" duplex_max_sessions: int = Field(default=1, ge=1) @@ -1325,6 +1326,7 @@ _DIFFUSION_STAGE_METADATA_FIELDS = frozenset( "model_arch", "model_stage", "retains_state_across_chunks", + "supports_running_prefix_cache_reset", "scheduler_cls", "stage_connector_spec", "worker_type", @@ -1972,6 +1974,8 @@ def _build_model_config( if "active_stream_window" not in kwargs: kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window) kwargs["final_output"] = topology.final_output + if not topology.supports_running_prefix_cache_reset: + kwargs["supports_running_prefix_cache_reset"] = False if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None: kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir) stage_runner = resolve_stage_model_runner(deploy, stage_deploy) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f016b7f7d..cbcb94cdb 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -117,7 +117,7 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import ( QWEN3_OMNI_THINKER_ONLY_PIPELINE, resolve_qwen3_omni_pipeline, ) -from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE, QWEN3_TTS_PIPELINE from vllm_omni.model_executor.models.step_audio2.pipeline import ( STEP_AUDIO2_ASR_PIPELINE, STEP_AUDIO2_PIPELINE, @@ -204,6 +204,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "qwen3_omni_moe": resolve_qwen3_omni_pipeline, "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE, "qwen3_tts": QWEN3_TTS_PIPELINE, + "qwen3_tts_fused": QWEN3_TTS_FUSED_PIPELINE, "step_audio_2": STEP_AUDIO2_PIPELINE, "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5f959050a..eabea35c2 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -245,6 +245,8 @@ class StagePipelineConfig: # The model keeps per-request execution state while awaiting the next # async chunk, so the parked request continues to consume model capacity. retains_state_across_chunks: bool = False + # Some stateful audio stages cannot roll back already consumed frames. + supports_running_prefix_cache_reset: bool = True sampling_constraints: dict[str, Any] = field(default_factory=dict) custom_process_input_func: str | None = None custom_process_next_stage_input_func: str | None = None @@ -295,6 +297,9 @@ class PipelineConfig: model_type: str model_arch: str = "" stages: tuple[StagePipelineConfig, ...] = () + # A single stage that streams its own final output in async-chunk mode + # (e.g. a Talker decoding audio in-stage) keeps deploy.async_chunk. + single_stage_async_chunk: bool = False # HF architecture aliases: used by StageConfigFactory when the model's # HF config reports a generic model_type that collides with a different # model (e.g. MiMo Audio reports model_type="qwen2"). The factory @@ -1107,6 +1112,8 @@ def _build_engine_args( if ps.omni_kv_config: engine_args["omni_kv_config"] = dict(ps.omni_kv_config) engine_args["requires_full_payload_input"] = ps.requires_full_payload_input + if not ps.supports_running_prefix_cache_reset: + engine_args["supports_running_prefix_cache_reset"] = False return engine_args @@ -1172,8 +1179,11 @@ def _resolve_pipeline_async_chunk_enabled( and False otherwise. If the user tried to enable async chunk through the deploy config, but it's inapplicable or unsupported, it will be disabled with a warning. """ - # Single stage should never use async chunk if len(pipeline.stages) <= 1: + if pipeline.single_stage_async_chunk: + if deploy.async_chunk is False: + raise ValueError(f"Pipeline {pipeline.model_type!r} requires async_chunk=True") + return True if deploy.async_chunk: logger.warning( "Deploy config set async_chunk=True, but async chunk is inapplicable " @@ -1350,6 +1360,8 @@ class StageConfig: # Terminal-stage ownership comes from topology, not engine overrides. engine_args["final_output"] = self.final_output + if self.yaml_engine_args.get("supports_running_prefix_cache_reset") is False: + engine_args["supports_running_prefix_cache_reset"] = False # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0138d365e..7dc0a5d7c 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -101,6 +101,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): max_num_running_reqs: int + def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool: + model_config = self.vllm_config.model_config + if ( + reset_running_requests + and self.running + and not getattr(model_config, "supports_running_prefix_cache_reset", True) + ): + # Reset preempts and resumes in the same step while discarding + # in-flight tokens. The stateful codec and queued PCM have already + # consumed those frames and cannot roll back to that boundary. + logger.warning( + "This stage cannot reset running requests; wait for " + "completion or abort them before resetting the prefix cache." + ) + return False + return super().reset_prefix_cache(reset_running_requests, reset_connector) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Track requests that need KV cache transfer when finished @@ -517,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # domains must consume the old frame before new output passes. request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids)) - if output_is_stale or async_output_is_stale: + if async_output_is_stale or (output_is_stale and request.drop_stale_output): # Output of a step scheduled before the request's in-flight # tokens were discarded (segment stop / session replacement). # num_computed_tokens was rolled back at the discard site, so @@ -556,12 +573,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # tokens and rejections. If some tokens are rejected, # num_computed_tokens is decreased by the number of rejected # tokens. - if request.num_computed_tokens > 0: - request.num_computed_tokens -= num_rejected - # If async scheduling, num_output_placeholders also includes - # the scheduled spec tokens count and so is similarly adjusted. - if request.num_output_placeholders > 0: - request.num_output_placeholders -= num_rejected + if not output_is_stale: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + # If async scheduling, num_output_placeholders also includes + # the scheduled spec tokens count and so is similarly adjusted. + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( spec_decoding_stats, num_draft_tokens=num_draft_tokens, @@ -600,7 +618,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Check for stop and update request status. if new_token_ids: num_sampled_tokens = len(new_token_ids) - new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) + if output_is_stale: + new_token_ids, stopped = self._update_request_with_output(request, new_token_ids, is_stale=True) + else: + new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens: # A mid-step stop (e.g. spec-decode tokens sampled past # EOS) trims new_token_ids after the validation slice @@ -704,6 +725,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # a queued streaming update, and adding twice swallows the # next duplex unit's listen/speak under async scheduling. if request.num_in_flight_tokens > 0: + request.drop_stale_output = True request.num_stale_output_tokens = request.num_in_flight_tokens if outstanding_async_tokens > 0: # Discard only outputs that are already in flight and @@ -884,6 +906,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # a resumable stop applies a queued update through this helper. in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) if in_flight_tokens > 0: + session.drop_stale_output = True session.num_stale_output_tokens = in_flight_tokens if outstanding_async_tokens > 0: # Async scheduling may already have sampled the previous @@ -1405,6 +1428,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if status is not None: request.status = status if native_transfer and connector_delay_free_blocks: + assert transfer_params is not None kv_xfer_params = { **(kv_xfer_params or {}), "transfer_id": transfer_params["transfer_id"], diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 1a1a46276..77a568d14 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -222,6 +222,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # num_in_flight_tokens already includes any undrained stale share. # Assign instead of accumulating so callers that fenced the same # rollover before entering this helper do not count it twice. + session.drop_stale_output = True session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) session.num_output_placeholders = 0 session.spec_token_ids = [] @@ -316,6 +317,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # The streaming update may already have fenced this same in-flight # frame. Seed idempotently so the replacement does not count it twice. request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0) + request.drop_stale_output = True request.num_output_placeholders = 0 request.spec_token_ids = [] self._release_replaced_streaming_prompt_cache(request) diff --git a/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml b/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml new file mode 100644 index 000000000..c2492f68c --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml @@ -0,0 +1,54 @@ +# Single-stage Qwen3-TTS: the Talker decodes every frame with the stateful +# streaming codec decoder and returns PCM itself (no Code2Wav stage). +# reset_prefix_cache(reset_running_requests=True) returns False while requests +# are running; wait for completion or abort them before resetting the cache. +pipeline: qwen3_tts_fused +model_runner: v2 +async_chunk: true + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 128 + # Exact conditioned Talker inputs for preemption replay are allocated + # after KV profiling: 16 MiB per active request at 1.7B/max_model_len=4096 + # (~2 GiB at 128 requests), in addition to codec state and output buffers. + # Keep headroom when increasing gpu_memory_utilization or capacity. + gpu_memory_utilization: 0.5 + trust_remote_code: true + enable_prefix_caching: false + # Model options that multi-stage deployments put on the connector: a + # single stage has no connector edge. + additional_config: + talker_stream_decode: true + # Opt in to first-frame delivery on its own queue for lower first-audio + # latency; leave disabled when prioritizing peak throughput. + talker_stream_first_audio: false + # Match the standard two-stage profile's voice-clone context boundary. + # A shorter context (e.g. 25) can change timbre continuity. + ref_code_context_frames: 72 + code_predictor_kv_cache: true + code_predictor_fused_sampling: true + code_predictor_fused: true + # Talker decode contexts are short (~100-700 tokens): FA3 split-KV plus + # its combine kernel costs more than one pass per head. + attention_config: + flash_attn_max_num_splits_for_cuda_graph: 1 + async_scheduling: true + max_num_batched_tokens: 512 + max_model_len: 4096 + 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 + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + min_tokens: 2 + repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 458a6237d..bad555704 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -198,6 +198,7 @@ class OmniEngineArgs(EngineArgs): async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + supports_running_prefix_cache_reset: bool = True use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. @@ -430,6 +431,7 @@ class OmniEngineArgs(EngineArgs): async_chunk=self.async_chunk, session_mode=self.session_mode, retains_state_across_chunks=self.retains_state_across_chunks, + supports_running_prefix_cache_reset=self.supports_running_prefix_cache_reset, use_v2_model_runner=self.use_v2_model_runner, supports_native_mrv2_data_plane=self.supports_native_mrv2_data_plane, active_stream_window=self.active_stream_window, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4047a7147..8f55280cd 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -898,7 +898,7 @@ class OrchestratorBase: 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 + or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id ): continue audio = mm.get("model_outputs") @@ -935,7 +935,7 @@ class OrchestratorBase: 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 + codec_stage = source_stage if self.stage_pools[source_stage].final_output else 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: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index ecc9576bc..c19a5cf75 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -87,7 +87,10 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An 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"): + decodes_audio = getattr(model, "first_frame_decoder", None) is not None or ( + getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False)) + ) + if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"): return False from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f9f3849f4..92e04e946 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -249,6 +249,33 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) +def _own_reuseport_socket(sock: socket.socket, client_count: int) -> socket.socket: + """Give this API server its own SO_REUSEPORT listener on the shared address. + + All API servers accepting from the one inherited socket lets whichever + wakes first take a whole burst of connections (the event loop accepts + in a loop), and keep-alive then pins that imbalance for the connections' + lifetime. Separate reuseport listeners let the kernel spread connections + by hash. The inherited socket is bound but never listened on, so it + receives nothing. + """ + if client_count <= 1: + return sock + if sock.family not in (socket.AF_INET, socket.AF_INET6): + return sock + try: + if not sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT): + return sock + own = socket.socket(family=sock.family, type=socket.SOCK_STREAM) + own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + own.bind(sock.getsockname()[:2]) + except OSError as e: + logger.warning("Per-server reuseport listener unavailable (%s); sharing the inherited socket", e) + return sock + return own + + def run_omni_api_server_worker_proc( listen_address: str, sock: socket.socket, @@ -278,6 +305,7 @@ def run_omni_api_server_worker_proc( set_process_title("APIServer", str(client_index)) decorate_logs("APIServer", skip_if_decorated=True) + sock = _own_reuseport_socket(sock, int(omni_client_config["client_count"])) uvloop.run( omni_run_server_worker( listen_address, @@ -512,7 +540,7 @@ async def build_async_omni( # Ensures everything is shutdown and cleaned up on error/exit async with build_async_omni_from_stage_config( args, - disable_frontend_multiprocessing=disable_frontend_multiprocessing, + disable_frontend_multiprocessing=bool(disable_frontend_multiprocessing), client_config=client_config, ) as async_omni: yield async_omni diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index b82842c42..f8af4d9cd 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -46,6 +46,15 @@ except ImportError: logger = init_logger(__name__) +def _float32_to_pcm16_bytes(audio: np.ndarray) -> bytes: + """PCM_16 bytes of float32 samples, bit-identical to libsndfile's RAW writer. + + libsndfile (1.2) converts with ``floor(x * 0x8000)`` clipped to int16. + """ + scaled = np.floor(audio.astype(np.float64) * 32768.0) + return np.clip(scaled, -32768, 32767).astype(" None: + server = self.ctx.server + if server._tts_tokenizer is not None: + return + import transformers + from transformers import AutoTokenizer + + # Same construction as the engine's text tokenizer + # (Qwen3TTSPromptEmbedsBuilder.get_text_tokenizer), so the ids computed + # here can be handed to the engine. + kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True, padding_side="left") + if transformers.__version__ < "5": + kwargs["fix_mistral_regex"] = True + assert self.ctx.engine_client is not None + model_name = self.ctx.engine_client.model_config.model + server._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs) + def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: """Estimate prompt length so the placeholder matches model-side embeddings.""" try: from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, + build_assistant_text, ) server = self.ctx.server - if server._tts_tokenizer is None: - from transformers import AutoTokenizer - - model_name = self.ctx.engine_client.model_config.model - server._tts_tokenizer = AutoTokenizer.from_pretrained( - model_name, - trust_remote_code=True, - padding_side="left", - ) + self._load_text_tokenizer() + assert self.ctx.engine_client is not None hf_config = self.ctx.engine_client.model_config.hf_config talker_config = hf_config.talker_config task_type = (tts_params.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( + text = (tts_params.get("text") or [""])[0] + assistant_text = build_assistant_text(text) if isinstance(text, str) else None + assistant_ids: list[list[int]] = [] + + def tokenize_prompt(t: str) -> list[int]: + ids = server._tts_tokenizer(t, padding=False)["input_ids"] + if t == assistant_text: + assistant_ids.append(ids) + return ids + + prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=tts_params, task_type=task_type, - tokenize_prompt=lambda t: server._tts_tokenizer(t, padding=False)["input_ids"], + tokenize_prompt=tokenize_prompt, codec_language_id=getattr(talker_config, "codec_language_id", None), spk_is_dialect=getattr(talker_config, "spk_is_dialect", None), estimate_ref_code_len=self._estimate_ref_code_len, ) + if assistant_ids: + # Spares the engine tokenizing the same text on its step thread. + tts_params[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_ids[0]] + return prompt_len except Exception as e: logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e) return 2048 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 37a9c2388..88d3427bb 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -954,7 +954,13 @@ class CodePredictorWrapper(nn.Module): extra_cfg = connector_cfg.get("extra", connector_cfg) else: extra_cfg = getattr(connector_cfg, "extra", None) - return extra_cfg if isinstance(extra_cfg, dict) else {} + extra_cfg = extra_cfg if isinstance(extra_cfg, dict) else {} + # A single-stage deployment has no connector edge to carry model + # options; it sets them in the stage's ``additional_config``. + additional = getattr(vllm_config, "additional_config", None) + if isinstance(additional, dict) and additional: + return {**additional, **extra_cfg} + return extra_cfg @staticmethod def _parse_bool_config(value: object) -> bool: @@ -1332,6 +1338,21 @@ class CodePredictorWrapper(nn.Module): scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms) + def _validate_sampling_inputs( + self, + bsz: int, + generators: Sequence[torch.Generator | None] | None, + sample_uniforms: torch.Tensor | None, + ) -> None: + if generators is not None and len(generators) != bsz: + raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") + if sample_uniforms is not None: + expected_shape = (bsz, self._num_groups - 1, int(self.config.vocab_size)) + if tuple(sample_uniforms.shape) != expected_shape: + raise ValueError( + f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}" + ) + @torch.inference_mode() def forward( self, @@ -1349,14 +1370,7 @@ class CodePredictorWrapper(nn.Module): """Predict residual codebooks 1..G-1 autoregressively via re-prefill.""" bsz = int(layer0_code.shape[0]) num_groups = self._num_groups - if generators is not None and len(generators) != bsz: - raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") - if sample_uniforms is not None: - expected_shape = (bsz, num_groups - 1, int(self.config.vocab_size)) - if tuple(sample_uniforms.shape) != expected_shape: - raise ValueError( - f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}" - ) + self._validate_sampling_inputs(bsz, generators, sample_uniforms) sample_generator: _GeneratorLike = generators if generators is not None else generator device = layer0_code.device diff --git a/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py b/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py new file mode 100644 index 000000000..5133dd85a --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Fused residual-codebook predictor for Qwen3-TTS (CUDA, BF16). + +Same computation as the frame-local-KV predictor with per-call top-k Gumbel +sampling, restructured to issue few kernels per residual step, since at serving +batch sizes each of the 15 dependent steps is launch- and latency-bound: + +- Per layer: RMSNorm, one QKV GEMM, one attention kernel (q/k RMSNorm, RoPE, + this call's K/V write and attention over the frame's keys), O projection + accumulated into the residual stream (``addmm_``), RMSNorm, gate/up GEMM, + SiLU-mul, down projection accumulated into the stream. +- The embedding of each sampled code followed by ``small_to_mtp_projection`` + is linear, so it folds at load into one ``[vocab, hidden]`` table per + codebook; the sampling kernel writes the next step's input row directly. +- Sampling finds the top-k threshold by bisection over the BF16 order keys + (exact k-th largest, ties kept) instead of a sort. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, tldevice, triton + +from .tokenizer_12hz.streaming_decoder import _silu_mul_kernel + + +@triton.jit +def _rms_rows_kernel(x_ptr, x_stride, w_ptr, out_ptr, eps, N: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, N) + x = tl.load(x_ptr + row * x_stride + offs).to(tl.float32) + var = tl.sum(x * x, axis=0) / N + dt = out_ptr.dtype.element_ty + y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32) + w = tl.load(w_ptr + offs).to(tl.float32) + tl.store(out_ptr + row * N + offs, (w * y).to(dt)) + + +@triton.jit +def _bf(x): + return x.to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _head_norm_rope(x1, x2, w1, w2, c1, c2, s1, s2, eps, HD: tl.constexpr): + """HF RMSNorm (``weight * normed.to(bf16)``) then RoPE in BF16 elementwise steps, on the two halves.""" + var = (tl.sum(x1 * x1, axis=0) + tl.sum(x2 * x2, axis=0)) / HD + r = tl.rsqrt(var + eps) + n1 = _bf(w1 * _bf(x1 * r)) + n2 = _bf(w2 * _bf(x2 * r)) + # rotate_half(n) = [-n2, n1] + o1 = _bf(_bf(n1 * c1) + _bf(-n2 * s1)) + o2 = _bf(_bf(n2 * c2) + _bf(n1 * s2)) + return o1, o2 + + +@triton.jit +def _cp_attention_kernel( + qkv_ptr, # [B * NQ, (H + 2 * KVH) * HD] + kc_ptr, # [Bmax, KVH, MAXP, HD] + vc_ptr, + qw_ptr, + kw_ptr, + cos_ptr, # [MAXP, HD] bf16 + sin_ptr, + out_ptr, # [B * NQ, H * HD] + first, + eps, + scale, + NQ: tl.constexpr, + H: tl.constexpr, + KVH: tl.constexpr, + HD: tl.constexpr, + MAXP: tl.constexpr, + BK: tl.constexpr, +): + b = tl.program_id(0) + h = tl.program_id(1) + qi = tl.program_id(2) + kvh = h // (H // KVH) + W: tl.constexpr = (H + 2 * KVH) * HD + HALF: tl.constexpr = HD // 2 + d = tl.arange(0, HALF) + p = first + qi + row = (b * NQ + qi).to(tl.int64) + qw1 = tl.load(qw_ptr + d).to(tl.float32) + qw2 = tl.load(qw_ptr + HALF + d).to(tl.float32) + kw1 = tl.load(kw_ptr + d).to(tl.float32) + kw2 = tl.load(kw_ptr + HALF + d).to(tl.float32) + qb = qkv_ptr + row * W + h * HD + q1, q2 = _head_norm_rope( + tl.load(qb + d).to(tl.float32), tl.load(qb + HALF + d).to(tl.float32), qw1, qw2, + tl.load(cos_ptr + p * HD + d).to(tl.float32), tl.load(cos_ptr + p * HD + HALF + d).to(tl.float32), + tl.load(sin_ptr + p * HD + d).to(tl.float32), tl.load(sin_ptr + p * HD + HALF + d).to(tl.float32), + eps, HD, + ) # fmt: skip + j = tl.arange(0, BK) + cache = ((b * KVH + kvh) * MAXP).to(tl.int64) * HD + old = (j < first)[:, None] + kp = kc_ptr + cache + j[:, None] * HD + d[None, :] + vp = vc_ptr + cache + j[:, None] * HD + d[None, :] + k1 = tl.load(kp, mask=old, other=0.0).to(tl.float32) + k2 = tl.load(kp + HALF, mask=old, other=0.0).to(tl.float32) + v1 = tl.load(vp, mask=old, other=0.0).to(tl.float32) + v2 = tl.load(vp + HALF, mask=old, other=0.0).to(tl.float32) + for n in tl.static_range(NQ): + # this call's position first + n, computed from its own QKV row + on = first + n + kb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + kvh) * HD + vb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + KVH + kvh) * HD + kn1, kn2 = _head_norm_rope( + tl.load(kb + d).to(tl.float32), tl.load(kb + HALF + d).to(tl.float32), kw1, kw2, + tl.load(cos_ptr + on * HD + d).to(tl.float32), tl.load(cos_ptr + on * HD + HALF + d).to(tl.float32), + tl.load(sin_ptr + on * HD + d).to(tl.float32), tl.load(sin_ptr + on * HD + HALF + d).to(tl.float32), + eps, HD, + ) # fmt: skip + vn1 = tl.load(vb + d).to(tl.float32) + vn2 = tl.load(vb + HALF + d).to(tl.float32) + sel = (j == on)[:, None] + k1 = tl.where(sel, kn1[None, :], k1) + k2 = tl.where(sel, kn2[None, :], k2) + v1 = tl.where(sel, vn1[None, :], v1) + v2 = tl.where(sel, vn2[None, :], v2) + wm = (d < HALF) & (qi == n) & (h % (H // KVH) == 0) + dtc = kc_ptr.dtype.element_ty + tl.store(kc_ptr + cache + on * HD + d, kn1.to(dtc), mask=wm) + tl.store(kc_ptr + cache + on * HD + HALF + d, kn2.to(dtc), mask=wm) + tl.store(vc_ptr + cache + on * HD + d, vn1.to(dtc), mask=wm) + tl.store(vc_ptr + cache + on * HD + HALF + d, vn2.to(dtc), mask=wm) + km = j <= p + s = (tl.sum(k1 * q1[None, :], axis=1) + tl.sum(k2 * q2[None, :], axis=1)) * scale + s = tl.where(km, s, -float("inf")) + e = tl.exp(s - tl.max(s, axis=0)) + e = tl.where(km, e, 0.0) + w = e / tl.sum(e, axis=0) + ob = out_ptr + row * (H * HD) + h * HD + dto = out_ptr.dtype.element_ty + tl.store(ob + d, tl.sum(v1 * w[:, None], axis=0).to(dto)) + tl.store(ob + HALF + d, tl.sum(v2 * w[:, None], axis=0).to(dto)) + + +@triton.jit +def _cp_sample_kernel( + logits_ptr, # [B, V] bf16 + u_ptr, # uniforms, row stride u_stride, contiguous over V + u_stride, + codes_ptr, # [B, G] int64 + step, + G: tl.constexpr, + table_ptr, # [V, HID] next-step input table, or dummy + next_ptr, # [B, HID] + inv_temperature, + V: tl.constexpr, + TOPK: tl.constexpr, + HID: tl.constexpr, + HAS_NEXT: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + idx = tl.arange(0, V) + value = tl.load(logits_ptr + row * V + idx) + scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype) + if TOPK > 0: + # order-preserving 16-bit keys of the BF16 values + bits = scaled.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF + # Top-k keeps all numeric ties at the cutoff, including both zeros. + bits = tl.where((bits & 0x7FFF) == 0, 0, bits) + key = tl.where(bits >= 0x8000, 0xFFFF - bits, bits | 0x8000) + the = 0 + for bit in tl.static_range(15, -1, -1): + cand = the | (1 << bit) + cnt = tl.sum((key >= cand).to(tl.int32), axis=0) + the = tl.where(cnt >= TOPK, cand, the) + sf = tl.where(key >= the, scaled.to(tl.float32), -float("inf")) + else: + sf = scaled.to(tl.float32) + u = tl.load(u_ptr + row * u_stride + idx) + score = sf - tldevice.log(-tldevice.log(u)) + max_score = tl.max(score, 0) + chosen = tl.min(tl.where(score == max_score, idx, V), 0) + first_nan = tl.min(tl.where(score != score, idx, V), 0) + code = tl.where(first_nan < V, first_nan, chosen) + tl.store(codes_ptr + row * G + step, code.to(tl.int64)) + if HAS_NEXT: + c = tl.arange(0, HID) + tl.store(next_ptr + row * HID + c, tl.load(table_ptr + code.to(tl.int64) * HID + c)) + + +class FusedCodePredictor: + """Execution object built from a loaded ``Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM``.""" + + def __init__(self, predictor, max_batch: int) -> None: + model = predictor.model + weight = next(model.parameters()) + dev, dt = weight.device, weight.dtype + if dt != torch.bfloat16: + raise ValueError("fused code predictor requires BF16 weights") + cfg = predictor.config + self.groups = int(cfg.num_code_groups) + self.vocab = int(cfg.vocab_size) + self.hidden = int(cfg.hidden_size) + attn0 = model.layers[0].self_attn + self.heads, self.kv_heads, self.head_dim = attn0.num_heads, attn0.num_kv_heads, attn0.head_dim + self.scale = attn0.scaling + self.eps = float(model.norm.variance_epsilon) + self.max_pos = attn0.max_seq + self.max_batch = max_batch + f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32) # noqa: E731 + c = lambda t: t.to(dt).contiguous() # noqa: E731 + proj = predictor.small_to_mtp_projection + if isinstance(proj, torch.nn.Identity): + pw, pb = None, None + else: + pw, pb = f32(proj.weight), f32(proj.bias) + self.proj_wt = c(pw.t()) if pw is not None else None + self.proj_b = c(pb) if pb is not None else None + tables = [] + for emb in model.codec_embedding: + e = f32(emb.weight) + tables.append(c(e @ pw.t() + pb) if pw is not None else c(e)) + self.tables = tables + self.layers = [] + for layer in model.layers: + a = layer.self_attn + self.layers.append( + dict( + ln1=c(f32(layer.input_layernorm.weight)), + ln2=c(f32(layer.post_attention_layernorm.weight)), + eps1=float(layer.input_layernorm.variance_epsilon), + eps2=float(layer.post_attention_layernorm.variance_epsilon), + qkv=c(f32(a.qkv_proj.weight).t()), + qn=c(f32(a.q_norm.weight)), + kn=c(f32(a.k_norm.weight)), + head_eps=float(a.q_norm.variance_epsilon), + o=c(f32(a.o_proj.weight).t()), + gu=c(f32(layer.mlp.gate_up_proj.weight).t()), + down=c(f32(layer.mlp.down_proj.weight).t()), + ) + ) + if a.qkv_proj.bias is not None: + raise ValueError("fused code predictor assumes bias-free attention") + self.norm_w = c(f32(model.norm.weight)) + self.lm_heads = [c(f32(h.weight).t()) for h in predictor.lm_head] + rot = model.rotary_emb + self.cos = rot.cos_cached.to(device=dev, dtype=dt).contiguous() + self.sin = rot.sin_cached.to(device=dev, dtype=dt).contiguous() + kv = (len(self.layers), 2, max_batch, self.kv_heads, self.max_pos, self.head_dim) + self.cache = torch.zeros(kv, device=dev, dtype=dt) + self._dummy = torch.zeros(1, device=dev, dtype=dt) + + def _layer(self, x: torch.Tensor, li: int, B: int, nq: int, first: int) -> None: + L = self.layers[li] + rows = x.shape[0] + h = torch.empty_like(x) + _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln1"], h, L["eps1"], N=self.hidden) + qkv = torch.mm(h, L["qkv"]) + att = torch.empty(rows, self.heads * self.head_dim, device=x.device, dtype=x.dtype) + _cp_attention_kernel[(B, self.heads, nq)]( + qkv, self.cache[li, 0], self.cache[li, 1], L["qn"], L["kn"], self.cos, self.sin, att, + first, L["head_eps"], self.scale, + NQ=nq, H=self.heads, KVH=self.kv_heads, HD=self.head_dim, MAXP=self.max_pos, + BK=triton.next_power_of_2(self.max_pos), num_warps=1, + ) # fmt: skip + x.addmm_(att, L["o"]) + _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln2"], h, L["eps2"], N=self.hidden) + gu = torch.mm(h, L["gu"]) + inter = gu.shape[1] // 2 + a = torch.empty(rows, inter, device=x.device, dtype=x.dtype) + _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024) + x.addmm_(a, L["down"]) + + @torch.inference_mode() + def __call__( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + inv_temperature: float, + top_k: int, + sample_uniforms: torch.Tensor, + ) -> torch.Tensor: + B = int(layer0_code.shape[0]) + if B > self.max_batch: + raise ValueError("fused code predictor batch exceeds its capacity") + dev, dt = layer0_embed.device, self.norm_w.dtype + inp = torch.stack([last_talker_hidden.reshape(B, -1), layer0_embed.reshape(B, -1)], 1).to(dt) + inp = inp.reshape(2 * B, -1) + x = torch.addmm(self.proj_b, inp, self.proj_wt) if self.proj_wt is not None else inp.clone() + codes = torch.empty(B, self.groups, device=dev, dtype=torch.int64) + codes[:, 0] = layer0_code.reshape(B) + u = sample_uniforms + nq, first = 2, 0 + for step in range(1, self.groups): + for li in range(len(self.layers)): + self._layer(x, li, B, nq, first) + last = x[nq - 1 :: nq] if nq > 1 else x + hf = torch.empty(B, self.hidden, device=dev, dtype=dt) + _rms_rows_kernel[(B,)](last, last.stride(0), self.norm_w, hf, self.eps, N=self.hidden) + logits = torch.mm(hf, self.lm_heads[step - 1]) + has_next = step < self.groups - 1 + nxt = torch.empty(B, self.hidden, device=dev, dtype=dt) if has_next else self._dummy + us = u[:, step - 1] + _cp_sample_kernel[(B,)]( + logits, us, us.stride(0), codes, step, self.groups, + self.tables[step - 1] if has_next else self._dummy, nxt, inv_temperature, + V=self.vocab, TOPK=top_k, HID=self.hidden, HAS_NEXT=has_next, num_warps=4, + ) # fmt: skip + x = nxt + first = step + 1 + nq = 1 + return codes diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index dd753861d..08c1f7c3e 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -72,3 +72,33 @@ QWEN3_TTS_PIPELINE = PipelineConfig( ), ), ) + + +# Single-stage variant: the Talker decodes each frame with the stateful +# streaming codec decoder and emits PCM as the final output (enable with the +# ``talker_stream_decode`` model option). No Code2Wav stage, no connector. +QWEN3_TTS_FUSED_PIPELINE = PipelineConfig( + model_type="qwen3_tts_fused", + default_deploy_config_name="qwen3_tts_fused_single_gpu.yaml", + single_stage_async_chunk=True, + model_arch="Qwen3TTSTalkerForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="qwen3_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + final_output=True, + final_output_type="audio", + # The API output processor tags ``model_outputs`` by this modality. + engine_output_type="audio", + supports_running_prefix_cache_reset=False, + supports_native_mrv2_data_plane=True, + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2150], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index 45d887ee2..7a63a3072 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -42,6 +42,7 @@ from vllm.logger import init_logger from vllm.multimodal.audio import AudioResampler from vllm_omni.utils.audio import mel_filter_bank +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking if TYPE_CHECKING: from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig @@ -160,7 +161,7 @@ def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | N ids = ids.unsqueeze(0) if ids.ndim != 2 or ids.numel() == 0: return None - return ids.to(device=device, dtype=torch.long).contiguous() + return to_device_nonblocking(ids.to(dtype=torch.long).contiguous(), device) # --------------------------------------------------------------------------- @@ -366,6 +367,9 @@ class Qwen3TTSPromptEmbedsBuilder: # Reusing them avoids launching the embedding/projection pair once per # request during prompt construction. self._projected_token_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {} + # CustomVoice prompt pieces that do not depend on the text, per + # (device, speaker, language): bounded by the model's voices. + self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries) self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict() @@ -1024,7 +1028,9 @@ class Qwen3TTSPromptEmbedsBuilder: dev = self._device() input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev) if input_ids is None: - input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev) + input_ids = to_device_nonblocking( + tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"], dev + ) # Optional instruct prefix. instruct = (info_dict.get("instruct") or [""])[0] @@ -1411,35 +1417,44 @@ class Qwen3TTSPromptEmbedsBuilder: raise ValueError(f"Unsupported speaker: {speaker}") spk_id = spk_id_map[speaker] # Keep it at least 1D; embedding on a 0-d tensor can return 1D. - spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long) + spk_tensor = index_to_device([spk_id], input_ids.device) spk_embed = codec_embed(spk_tensor) if spk_embed.ndim in (1, 2): spk_embed = spk_embed.view(1, 1, -1) speaker_embed = spk_embed codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) - role_embed = text_projection(text_embedding(input_ids[:, :3])) - codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) - codec_prefix = codec_prefix + codec_input[:, :-1] - talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) + # The role tokens, speaker/language codec prefix, codec pad row and + # closing row do not depend on the text: cache them per speaker and + # language so a request only embeds and projects its own text. + # A text starting with whitespace could merge with the template's + # newline token, so it keeps the per-request role embedding. + cache_key = None + if non_streaming_mode and instruct_embed is None and not text[:1].isspace(): + cache_key = ( + str(input_ids.device), + int(spk_id), + codec_prefill_list[0][2] if len(codec_prefill_list[0]) == 4 else None, + ) + cv_cache = self._custom_voice_prompt_cache + cached = cv_cache.get(cache_key) if cache_key is not None else None + if cached is None: + role_embed = text_projection(text_embedding(input_ids[:, :3])) + codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) + codec_prefix = codec_prefix + codec_input[:, :-1] + prefix = torch.cat((role_embed, codec_prefix), dim=1) + pad_row = codec_embed(self._long_tensor([talker_config.codec_pad_id], input_ids.device)).view(1, 1, -1) + tail = tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)) + if cache_key is not None: + cv_cache[cache_key] = (prefix, pad_row, tail) + else: + prefix, pad_row, tail = cached + talker_prompt = prefix if non_streaming_mode: text_all = text_projection(text_embedding(input_ids[:, 3:-5])) text_all = torch.cat([text_all, tts_eos_embed], dim=1) - pad_ids = torch.full( - (1, int(text_all.shape[1])), - int(talker_config.codec_pad_id), - device=input_ids.device, - dtype=torch.long, - ) - talker_prompt = torch.cat( - [ - talker_prompt, - text_all + codec_embed(pad_ids), - tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)), - ], - dim=1, - ) + talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1) trailing_text_hidden = tts_pad_embed else: first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] 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 9620d0d86..493b3b2a2 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 @@ -65,6 +65,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp 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 + self._fused_requested = self._parse_bool_config(extra.get("code_predictor_fused")) + self._fused = None def _setup_compile(self) -> None: if self._compiled_model_fwd is not None: @@ -100,11 +102,73 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp cache(self._proj_buf, batch, step) self._synchronize_warmup(weight.device) self._frame_cache = cache + if self._fused_requested: + from .fused_code_predictor import FusedCodePredictor + + self._fused = FusedCodePredictor(self, max_batch) + # Compile every kernel variant outside any graph capture. + hidden = int(self.talker_config.hidden_size) + for batch in range(1, min(2, max_batch) + 1): + zeros = torch.zeros(batch, 1, hidden, device=weight.device, dtype=weight.dtype) + uniforms = torch.full( + (batch, self._num_groups - 1, int(self.config.vocab_size)), 0.5, device=weight.device + ) + codes = torch.zeros(batch, 1, device=weight.device, dtype=torch.long) + self._fused(codes, zeros, zeros, 1.0, 50, uniforms) + self._synchronize_warmup(weight.device) + logger.info("Qwen3-TTS fused code predictor enabled (capacity %d)", max_batch) # 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) + @torch.inference_mode() + def forward( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + do_sample: bool = True, + temperature: float = 0.9, + top_k: int = 50, + top_p: float = 1.0, + generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, + sample_uniforms: torch.Tensor | None = None, + ): + self._validate_sampling_inputs(int(layer0_code.shape[0]), generators, sample_uniforms) + if self._fused_requested: + self._setup_compile() + if ( + self._fused is not None + and do_sample + and temperature > 0 + and top_p == 1.0 + and sample_uniforms is not None + and sample_uniforms.dtype == torch.float32 + and sample_uniforms.is_contiguous() + ): + return self._fused( + layer0_code, + layer0_embed, + last_talker_hidden, + 1.0 / temperature, + int(top_k), + sample_uniforms, + ) + return super().forward( + layer0_code, + layer0_embed, + last_talker_hidden, + do_sample=do_sample, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + generators=generators, + sample_uniforms=sample_uniforms, + ) + def _predict_step_logits( self, proj_buf: torch.Tensor, 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 e05ac3871..6e199454f 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 @@ -27,6 +27,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, @@ -39,6 +40,7 @@ from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfi 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 .stream_decode import stream_ref_context_frames, talker_stream_decode_enabled from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -411,7 +413,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # 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) + self.stream_decode = talker_stream_decode_enabled(vllm_config) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM + extra = predictor._stage_connector_extra_config(vllm_config) + self.stream_first_audio = self.stream_decode and predictor._parse_bool_config( + extra.get("talker_stream_first_audio") + ) + # Model-local voice-clone priming boundary; each deployment selects + # its reference context independently of the two-stage codec profile. + self.stream_ref_context_frames = stream_ref_context_frames(vllm_config) if self.stream_decode else 0 + self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) or self.stream_decode + self.stream_decoder = None + self.stream_graphs = None + self.stream_prime_graphs = None + self.stream_sample_rate = 0 + self.stream_chunk_frames = 25 # 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. @@ -718,6 +734,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # -------------------- Omni multimodal output plumbing -------------------- + @property + def mm_outputs_fresh_per_step(self) -> bool: + """Stream-decode outputs come from :meth:`_make_eager_omni_output`, allocated per step.""" + return self.stream_decoder is not None and self.eager_frames_active + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs @@ -895,6 +916,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), }, } + if self.stream_decoder is not None: + # Filled per sampled row by the runner's eager MTP; token-major so + # each request's slice is its frames of this step. + spf = int(self.stream_decoder.spf) + mm["model_outputs"] = torch.zeros((num_tokens, spf), dtype=torch.float32, device=hidden.device) + sr = torch.tensor(self.stream_sample_rate, dtype=torch.int32) + mm["sr"] = [sr] * len(info_dicts) ref_rows = [ index for index, info in enumerate(info_dicts) @@ -1442,8 +1470,72 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): 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} + elif self.stream_decode: + # The in-stage decoder owns every frame. Optional first-audio + # delivery uses its PCM; later chunks use regular step outputs. + from .first_frame_decoder import Qwen3TTSFirstFrameDecoder + from .tokenizer_12hz.streaming_decoder import StreamingCodecDecoder + + holder = Qwen3TTSFirstFrameDecoder(self.model_path) + decoder_loaded = holder.load(self.vllm_config) + loaded = set(loaded) | {f"stream_decoder.{name}" for name in decoder_loaded} + slots = int(self.vllm_config.scheduler_config.max_num_seqs) + self.stream_decoder = StreamingCodecDecoder(holder.decoder, num_slots=slots, dtype=torch.bfloat16) + self.stream_sample_rate = int(holder.sample_rate) + logger.info("Qwen3-TTS Talker stream decode enabled (%d slots)", slots) return loaded + def capture_stream_decode_graphs(self, batch_sizes: list[int]) -> None: + if self.stream_decoder is None or self.stream_graphs is not None: + return + from .tokenizer_12hz.streaming_decoder import StreamingDecodeGraphs + + with torch.inference_mode(): + self.stream_graphs = StreamingDecodeGraphs(self.stream_decoder, batch_sizes) + frames = min(self.stream_ref_context_frames, self.stream_chunk_frames, self.stream_decoder.max_frames) + if self.config.tts_model_type == "base" and frames > 1: + # Bound the multi-frame graph workspace by the configured + # token budget; larger priming groups keep the eager path. + budget = max(1, self.vllm_config.scheduler_config.max_num_batched_tokens // frames) + prime_sizes = sorted({1, *(size for size in batch_sizes if size <= budget)}) + self.stream_prime_graphs = StreamingDecodeGraphs(self.stream_decoder, prime_sizes, frames=frames) + logger.info("Captured Talker stream decode graphs for batch sizes %s", self.stream_graphs.sizes) + if self.stream_prime_graphs is not None: + logger.info( + "Captured Talker reference priming graphs for %d frames and batch sizes %s", + self.stream_prime_graphs.frames, + self.stream_prime_graphs.sizes, + ) + + def get_stream_ref_context(self, info: OmniPayload) -> torch.Tensor | None: + """Model-owned reference-code layout and context boundary for a PCM slot.""" + codes = info.get("codes") + ref = codes.get("ref") if isinstance(codes, dict) else None + if not isinstance(ref, torch.Tensor) or not ref.numel(): + return None + return ref.reshape(-1, int(self.talker_config.num_code_groups))[-self.stream_ref_context_frames :] + + def prime_stream_decoder(self, primes: list[tuple[int, torch.Tensor]]) -> None: + """Prime equal-length reference groups without padding or aliasing live slots.""" + if not primes: + return + stream = self.stream_decoder + if stream is None: + raise RuntimeError("Reference priming requires the single-stage PCM decoder") + groups: dict[int, list[tuple[int, torch.Tensor]]] = {} + for req_idx, ref in primes: + groups.setdefault(int(ref.shape[0]), []).append((req_idx, ref)) + for items in groups.values(): + codes = to_device_nonblocking(torch.stack([ref for _idx, ref in items]), stream.device).to(torch.int32) + n = len(items) + slots = index_to_device([idx for idx, _ref in items], stream.device, dtype=torch.int32) + for t0 in range(0, int(codes.shape[1]), self.stream_chunk_frames): + pos = index_to_device([t0] * n, stream.device, dtype=torch.int32) + chunk = codes[:, t0 : t0 + self.stream_chunk_frames].contiguous() + graphs = self.stream_prime_graphs + decode = graphs if graphs is not None and chunk.shape[1] == graphs.frames else stream + decode(chunk, slots, pos) + def _build_stacked_codec_embed(self) -> None: embeds = self.code_predictor.get_input_embeddings() if not embeds: diff --git a/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py b/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py new file mode 100644 index 000000000..747e1f0b4 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Eligibility and model options for the single-stage CUDA PCM pipeline.""" + +import torch +from vllm.config import VllmConfig + +from vllm_omni.platforms import current_omni_platform + + +def talker_stream_decode_enabled(vllm_config: VllmConfig) -> bool: + """The Talker decodes every frame itself and emits PCM as its final output.""" + 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 + requested = Predictor._parse_bool_config(extra.get("talker_stream_decode")) + if not requested: + if getattr(model, "engine_output_type", None) == "audio": + raise ValueError("The single-stage Qwen3-TTS pipeline requires talker_stream_decode=True") + if Predictor._parse_bool_config(extra.get("talker_stream_first_audio")): + raise ValueError("talker_stream_first_audio requires the single-stage talker_stream_decode path") + return False + if getattr(model, "engine_output_type", None) != "audio" or not getattr(model, "final_output", False): + raise ValueError( + "talker_stream_decode requires a final audio-output Talker; " + "select qwen3_tts_fused_single_gpu.yaml instead of enabling it on the two-stage pipeline" + ) + if getattr(model, "supports_running_prefix_cache_reset", True): + raise ValueError( + "The single-stage PCM pipeline must disable running prefix-cache reset; " + "select qwen3_tts_fused_single_gpu.yaml" + ) + if not ( + 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 + ): + raise ValueError( + "talker_stream_decode requires CUDA, MRv2, async_chunk=True, TP=PP=1, " + "an in-process worker and enable_prefix_caching=False; " + "use qwen3_tts_fused_single_gpu.yaml" + ) + if Predictor._parse_bool_config(extra.get("talker_first_audio")): + raise ValueError("talker_stream_decode and talker_first_audio are mutually exclusive") + return True + + +def stream_ref_context_frames(vllm_config: VllmConfig) -> int: + """Reference-code context frames for voice-clone priming in the Talker.""" + from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor + + extra = Predictor._stage_connector_extra_config(vllm_config) + return int(extra.get("ref_code_context_frames") or extra.get("codec_left_context_frames", 25)) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py new file mode 100644 index 000000000..b1dffed00 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py @@ -0,0 +1,799 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Stateful, fused streaming decoder for the Qwen3-TTS 12 Hz codec. + +Every layer of the codec decoder is causal, so what a new frame needs from the +past is small and fixed: each causal conv's last ``(k - 1) * dilation`` input +rows, each overlapping transposed conv's last GEMM row, and each attention +layer's last ``sliding_window`` keys and values. A request owns a slot of +per-layer state holding exactly that, and a call decodes ``T`` new frames for +each of ``B`` requests, reading and advancing the state. The output equals the +whole-utterance decode (``Qwen3TTSTokenizerV2Decoder._forward_exact``) frame by +frame, up to floating point rounding: nothing is recomputed and nothing is +truncated. + +Layout and fusion: + +- Activations are time-major ``[rows, C]``; every dense conv is a tap-major + im2col (one Triton kernel that also applies the pending bias and SnakeBeta, + and reads/writes the history) followed by one cuBLAS GEMM. +- Linear chains fold at load: the RVQ output projections, ``pre_conv`` and the + transformer's ``input_proj`` become one ``k = 3`` conv over the concatenated + codebook embeddings; LayerScale folds into ``o_proj``/``down_proj``; ConvNeXt + ``gamma`` folds into ``pwconv2``. +- A bias that feeds a residual stream is carried to that stream's consumer + (the next im2col or transposed-conv GEMM) instead of being added in place. +- Residual branches accumulate into their stream with in-place ``addmm_``. + +History buffers are double-buffered per slot: a call reads half ``parity[slot]`` +and writes the other half, then flips the parity, so no kernel overwrites rows +another program of the same launch still reads. A row whose first frame index +``pos`` is 0 reads zeros, which are the convs' causal padding, and attention +only looks back ``min(pos + 1, window)`` positions, so a new request needs no +state reset. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn +from vllm.triton_utils import tl, triton + +from vllm_omni.platforms import current_omni_platform + +_KERNEL = 7 + + +# --------------------------------------------------------------------------- kernels + + +@triton.jit +def _rvq_gather_kernel(codes_ptr, books_ptr, out_ptr, NQ: tl.constexpr, CB: tl.constexpr, HALF: tl.constexpr): + row = tl.program_id(0) + offs = tl.arange(0, HALF) + c0 = tl.load(codes_ptr + row * NQ).to(tl.int64) + c0 = tl.where((c0 >= 0) & (c0 < CB), c0, 0) # EOS/special rows decode as code 0; callers drop them + first = tl.load(books_ptr + c0 * HALF + offs).to(tl.float32) + acc = tl.zeros([HALF], dtype=tl.float32) + for q in range(1, NQ): + c = tl.load(codes_ptr + row * NQ + q).to(tl.int64) + c = tl.where((c >= 0) & (c < CB), c, 0) + acc += tl.load(books_ptr + (q * CB + c) * HALF + offs).to(tl.float32) + dt = out_ptr.dtype.element_ty + tl.store(out_ptr + row * 2 * HALF + offs, first.to(dt)) + tl.store(out_ptr + row * 2 * HALF + HALF + offs, acc.to(dt)) + + +@triton.jit +def _fast_sin(y): + # Two-step reduction to [-pi, pi], then the hardware approximation + # (absolute error ~4e-7 there): accurate sin's instruction count, not + # memory, bounds every kernel that applies SnakeBeta. + k = tl.extra.cuda.libdevice.rint(y * 0.15915494309189535) + r = tl.fma(-k, -1.7484555314695172e-07, tl.fma(-k, 6.2831854820251465, y)) + return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1) + + +@triton.jit +def _snake(v, a, ib): + s = _fast_sin(v * a) + return v + ib * s * s + + +@triton.jit +def _im2col_kernel( + x_ptr, # [B * T, C] layer input rows + bias_ptr, # [C] pending bias (added to x rows only) or dummy + alpha_ptr, # [C] exp(alpha) or dummy + ibeta_ptr, # [C] 1 / (exp(beta) + eps) or dummy + hist_ptr, # [2, S, H, C] transformed history rows + slots_ptr, + pos_ptr, + par_ptr, + col_ptr, # [B * T, K * C] + T, + S, + C: tl.constexpr, + K: tl.constexpr, + D: tl.constexpr, + H: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SNAKE: tl.constexpr, + BU: tl.constexpr, + BC: tl.constexpr, + EXT: tl.constexpr = False, +): + """Source row ``u`` of ``[history (H) | x (T)]``: scatter to every tap that reads it, save the last H.""" + b = tl.program_id(0) + u = tl.program_id(1) * BU + tl.arange(0, BU) + c = tl.program_id(2) * BC + tl.arange(0, BC) + cm = c < C + n_src = H + T + um = u < n_src + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + # history rows + from_hist = u < H + hrow = tl.where(from_hist, u, 0) + hv = tl.load( + hist_ptr + ((par * S + slot) * H + hrow)[:, None] * C + c[None, :], + mask=(from_hist & um & live)[:, None] & cm[None, :], + other=0.0, + ).to(tl.float32) + xrow = tl.where(from_hist, 0, u - H) + xv = tl.load( + x_ptr + (b * T + xrow).to(tl.int64)[:, None] * C + c[None, :], + mask=((~from_hist) & um)[:, None] & cm[None, :], + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + xv += tl.load(bias_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + if HAS_SNAKE: + a = tl.load(alpha_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + ib = tl.load(ibeta_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + xv = _snake(xv, a, ib) + dt = col_ptr.dtype.element_ty + v = tl.where(from_hist[:, None], hv, xv.to(dt).to(tl.float32)).to(dt) + if EXT: + tl.store(col_ptr + (b * n_src + u).to(tl.int64)[:, None] * C + c[None, :], v, mask=um[:, None] & cm[None, :]) + else: + for j in tl.static_range(K): + t = u - j * D + tm = (t >= 0) & (t < T) & um + tl.store( + col_ptr + (b * T + tl.where(tm, t, 0)).to(tl.int64)[:, None] * (K * C) + j * C + c[None, :], + v, + mask=tm[:, None] & cm[None, :], + ) + if H > 0: + hm = (u >= T) & um + tl.store( + hist_ptr + (((1 - par) * S + slot) * H + tl.where(hm, u - T, 0))[:, None] * C + c[None, :], + v, + mask=hm[:, None] & cm[None, :], + ) + + +@triton.jit +def _rms_norm_kernel(x_ptr, w_ptr, out_ptr, eps, N: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, N) + x = tl.load(x_ptr + row * N + offs).to(tl.float32) + var = tl.sum(x * x, axis=0) / N + dt = out_ptr.dtype.element_ty + y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32) + w = tl.load(w_ptr + offs).to(tl.float32) + tl.store(out_ptr + row * N + offs, (w * y).to(dt)) + + +@triton.jit +def _rope_kv_write_kernel( + qkv_ptr, # [B * T, 3 * NH * HD] + q_ptr, # [B * T, NH * HD] + k_ring_ptr, # [S, RING, NH * HD] + v_ring_ptr, + slots_ptr, + pos_ptr, + T, + theta, + NH: tl.constexpr, + HD: tl.constexpr, + RING: tl.constexpr, +): + r = tl.program_id(0) + h = tl.program_id(1) + b = r // T + p = tl.load(pos_ptr + b) + r % T + slot = tl.load(slots_ptr + b).to(tl.int64) + half: tl.constexpr = HD // 2 + i = tl.arange(0, half) + inv_freq = tl.exp(-(2.0 * i.to(tl.float32) / HD) * tl.log(theta)) + ang = p.to(tl.float32) * inv_freq + dt = q_ptr.dtype.element_ty + cos = tl.cos(ang).to(dt).to(tl.float32) + sin = tl.sin(ang).to(dt).to(tl.float32) + W: tl.constexpr = NH * HD + base = qkv_ptr + r.to(tl.int64) * 3 * W + h * HD + q1 = tl.load(base + i).to(tl.float32) + q2 = tl.load(base + half + i).to(tl.float32) + k1 = tl.load(base + W + i).to(tl.float32) + k2 = tl.load(base + W + half + i).to(tl.float32) + vv = tl.arange(0, HD) + v = tl.load(base + 2 * W + vv) + qo = q_ptr + r.to(tl.int64) * W + h * HD + tl.store(qo + i, (q1 * cos - q2 * sin).to(dt)) + tl.store(qo + half + i, (q2 * cos + q1 * sin).to(dt)) + ring = (slot * RING + p % RING) * W + h * HD + tl.store(k_ring_ptr + ring + i, (k1 * cos - k2 * sin).to(dt)) + tl.store(k_ring_ptr + ring + half + i, (k2 * cos + k1 * sin).to(dt)) + tl.store(v_ring_ptr + ring + vv, v) + + +@triton.jit +def _attend_kernel( + q_ptr, # [B * T, NH * HD] + k_ring_ptr, + v_ring_ptr, + out_ptr, # [B * T, NH * HD] + slots_ptr, + pos_ptr, + T, + scale, + NH: tl.constexpr, + HD: tl.constexpr, + RING: tl.constexpr, + WINDOW: tl.constexpr, + BK: tl.constexpr, +): + r = tl.program_id(0) + h = tl.program_id(1) + b = r // T + p = tl.load(pos_ptr + b) + r % T + slot = tl.load(slots_ptr + b).to(tl.int64) + W: tl.constexpr = NH * HD + d = tl.arange(0, HD) + q = tl.load(q_ptr + r.to(tl.int64) * W + h * HD + d).to(tl.float32) + j = tl.arange(0, BK) + n = tl.minimum(p + 1, WINDOW) + km = j < n + kp = p - j # key positions p, p-1, ... + idx = (slot * RING + tl.where(km, kp, 0) % RING) * W + h * HD + k = tl.load(k_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32) + s = tl.sum(k * q[None, :], axis=1) * scale + s = tl.where(km, s, -float("inf")) + m = tl.max(s, axis=0) + e = tl.exp(s - m) + e = tl.where(km, e, 0.0) + w = e / tl.sum(e, axis=0) + dt = out_ptr.dtype.element_ty + w = w.to(dt).to(tl.float32) + v = tl.load(v_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32) + o = tl.sum(v * w[:, None], axis=0) + tl.store(out_ptr + r.to(tl.int64) * W + h * HD + d, o.to(dt)) + + +@triton.jit +def _silu_mul_kernel(gu_ptr, out_ptr, N: tl.constexpr, BN: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BN + tl.arange(0, BN) + m = offs < N + g = tl.load(gu_ptr + row * 2 * N + offs, mask=m).to(tl.float32) + u = tl.load(gu_ptr + row * 2 * N + N + offs, mask=m).to(tl.float32) + dt = out_ptr.dtype.element_ty + sg = (g / (1.0 + tl.exp(-g))).to(dt).to(tl.float32) + tl.store(out_ptr + row * N + offs, (sg * u).to(dt), mask=m) + + +@triton.jit +def _dwconv_ln_kernel( + z_ptr, # [B * T, C] ConvNeXt input rows (true values) + w_ptr, # [K, C] + b_ptr, # [C] + lnw_ptr, + lnb_ptr, + hist_ptr, # [2, S, K - 1, C] + slots_ptr, + pos_ptr, + par_ptr, + out_ptr, # [B * T, C] LayerNorm output + T, + S, + eps, + C: tl.constexpr, + K: tl.constexpr, +): + b = tl.program_id(0) + t = tl.program_id(1) # 0 .. max(T, K - 1) + H: tl.constexpr = K - 1 + c = tl.arange(0, C) + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + dt = out_ptr.dtype.element_ty + hbase = hist_ptr + (par * S + slot) * H * C + acc = tl.load(b_ptr + c).to(tl.float32) + for j in tl.static_range(K): + u = t + j # index into [hist (H) | z (T)] + fh = u < H + hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live, other=0.0).to(tl.float32) + zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64) + zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < T), other=0.0).to(tl.float32) + acc += tl.where(fh, hv, zv) * tl.load(w_ptr + j * C + c).to(tl.float32) + h = acc.to(dt).to(tl.float32) + mean = tl.sum(h, axis=0) / C + hc = h - mean + var = tl.sum(hc * hc, axis=0) / C + y = hc * tl.rsqrt(var + eps) * tl.load(lnw_ptr + c).to(tl.float32) + tl.load(lnb_ptr + c).to(tl.float32) + tl.store(out_ptr + (b * T + t).to(tl.int64) * C + c, y.to(dt), mask=(c < C) & (t < T)) + # new history row t: source row T + t of [hist | z] + u = T + t + fh = u < H + hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live & (t < H), other=0.0) + zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64) + zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < H), other=0.0) + tl.store( + hist_ptr + (((1 - par) * S + slot) * H + tl.where(t < H, t, 0)) * C + c, + tl.where(fh, hv, zv).to(dt), + mask=(c < C) & (t < H), + ) + + +@triton.jit +def _bias_gelu_kernel(x_ptr, b_ptr, N: tl.constexpr, BN: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BN + tl.arange(0, BN) + m = offs < N + x = tl.load(x_ptr + row * N + offs, mask=m).to(tl.float32) + tl.load(b_ptr + offs, mask=m).to(tl.float32) + dt = x_ptr.dtype.element_ty + x = x.to(dt).to(tl.float32) + y = 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476)) + tl.store(x_ptr + row * N + offs, y.to(dt), mask=m) + + +@triton.jit +def _col2im_kernel( + z_ptr, # [B * T, 2 * R, C] transposed-conv GEMM rows + bias_ptr, # [C] + prev_ptr, # [2, S, R, C] previous call's last GEMM row, second half + slots_ptr, + pos_ptr, + par_ptr, + out_ptr, # [B * T * R, C] + T, + S, + R: tl.constexpr, + C: tl.constexpr, + BC: tl.constexpr, +): + row = tl.program_id(0) # b * T + t + j = tl.program_id(1) # 0 .. R - 1 + c = tl.program_id(2) * BC + tl.arange(0, BC) + cm = c < C + b = row // T + t = row % T + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + zr = z_ptr + row.to(tl.int64) * 2 * R * C + cur = tl.load(zr + j * C + c, mask=cm).to(tl.float32) + prev_in = tl.load(zr - 2 * R * C + (R + j) * C + c, mask=cm & (t > 0), other=0.0).to(tl.float32) + prev_st = tl.load(prev_ptr + ((par * S + slot) * R + j) * C + c, mask=cm & (t == 0) & live, other=0.0).to( + tl.float32 + ) + y = cur + tl.where(t > 0, prev_in, prev_st) + tl.load(bias_ptr + c, mask=cm).to(tl.float32) + dt = out_ptr.dtype.element_ty + tl.store(out_ptr + (row.to(tl.int64) * R + j) * C + c, y.to(dt), mask=cm) + last = tl.load(zr + (R + j) * C + c, mask=cm & (t == T - 1)) + tl.store(prev_ptr + (((1 - par) * S + slot) * R + j) * C + c, last, mask=cm & (t == T - 1)) + + +@triton.jit +def _conv_out_kernel( + ext_ptr, # [B, H + T, C] transformed rows: history then this call's + w_ptr, # [K, C] + wb, # conv bias (float) + out_ptr, # [B, T] float32 waveform + T, + C: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + CP: tl.constexpr, +): + b = tl.program_id(0) + t = tl.program_id(1) * BT + tl.arange(0, BT) + c = tl.arange(0, CP) + cm = c < C + tm = t < T + base = ext_ptr + b.to(tl.int64) * (T + K - 1) * C + acc = tl.zeros([BT], dtype=tl.float32) + for j in tl.static_range(K): + v = tl.load(base + (t + j)[:, None] * C + c[None, :], mask=tm[:, None] & cm[None, :], other=0.0).to(tl.float32) + acc += tl.sum(v * tl.load(w_ptr + j * C + c, mask=cm, other=0.0).to(tl.float32)[None, :], axis=1) + dt = ext_ptr.dtype.element_ty + y = acc.to(dt).to(tl.float32) + wb + y = tl.minimum(tl.maximum(y.to(dt).to(tl.float32), -1.0), 1.0) + tl.store(out_ptr + b.to(tl.int64) * T + t, y, mask=tm) + + +@triton.jit +def _flip_kernel(par_ptr, slots_ptr): + slot = tl.load(slots_ptr + tl.program_id(0)).to(tl.int64) + tl.store(par_ptr + slot, 1 - tl.load(par_ptr + slot)) + + +# --------------------------------------------------------------------------- module + + +@dataclass +class _ConvSpec: + weight: torch.Tensor # [K * C_in, C_out] tap-major + c_in: int + k: int + d: int + hist: torch.Tensor | None = None # [2, S, H, C_in] + + +def _conv_taps(w: torch.Tensor) -> torch.Tensor: + """``[C_out, C_in, K]`` -> ``[K * C_in, C_out]``.""" + c_out, c_in, k = w.shape + return w.permute(2, 1, 0).reshape(k * c_in, c_out) + + +def _snake_params(module) -> tuple[torch.Tensor, torch.Tensor]: + # SnakeBeta: x + 1 / (exp(beta) + 1e-9) * sin(exp(alpha) * x) ** 2 + alpha = module.alpha.detach().float() + beta = module.beta.detach().float() + return torch.exp(alpha).contiguous(), (1.0 / (torch.exp(beta) + 1e-9)).contiguous() + + +class StreamingCodecDecoder(nn.Module): + """Frame-exact stateful decoder built from a loaded ``Qwen3TTSTokenizerV2Decoder``.""" + + RING = 128 # KV ring positions; a call may carry up to RING - window + 1 frames + + def __init__(self, decoder: nn.Module, num_slots: int, dtype: torch.dtype | None = None) -> None: + super().__init__() + cfg = decoder.config + dev = next(decoder.parameters()).device + dtype = dtype or next(decoder.parameters()).dtype + self.dtype, self.device = dtype, dev + self.num_slots = num_slots + 1 # last slot: scratch for padding rows + self.scratch_slot = num_slots + S = self.num_slots + self.nq = int(cfg.num_quantizers) + self.cb = int(cfg.codebook_size) + self.half = int(cfg.codebook_dim) // 2 + self.hidden = int(cfg.hidden_size) + self.heads = int(cfg.num_attention_heads) + self.head_dim = int(cfg.head_dim) + self.window = int(cfg.sliding_window) + self.theta = float(cfg.rope_theta) + self.eps = float(cfg.rms_norm_eps) + self.max_frames = self.RING - self.window + 1 + self.spf = int(decoder.total_upsample) + if self.heads != int(cfg.num_key_value_heads): + raise ValueError("streaming codec decoder assumes no grouped KV") + + f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32) # noqa: E731 + cast = lambda t: t.to(dtype).contiguous() # noqa: E731 + + def hist(h: int, c: int) -> torch.Tensor: + return torch.zeros(2, S, h, c, device=dev, dtype=dtype) + + # RVQ books, pre-divided by usage: [NQ * CB, HALF] + q = decoder.quantizer + books = [] + for vq in [q.rvq_first.vq.layers[0]] + list(q.rvq_rest.vq.layers): + cbk = vq._codebook + books.append(f32(cbk.embedding_sum) / f32(cbk.cluster_usage).clamp(min=cbk.epsilon)[:, None]) + self.books = cast(torch.cat(books)) + w_rvq = torch.cat( + [f32(q.rvq_first.output_proj.weight)[:, :, 0], f32(q.rvq_rest.output_proj.weight)[:, :, 0]], 1 + ) + # pre_conv (k=3) and input_proj folded onto the embedding concat + pre_w, pre_b = f32(decoder.pre_conv.conv.weight), f32(decoder.pre_conv.conv.bias) + pt = decoder.pre_transformer + in_w, in_b = f32(pt.input_proj.weight), f32(pt.input_proj.bias) + taps = torch.stack([in_w @ pre_w[:, :, j] @ w_rvq for j in range(pre_w.shape[2])], dim=2) # [hid, 2H, 3] + self.pre = _ConvSpec(cast(_conv_taps(taps)), 2 * self.half, 3, 1, hist(2, 2 * self.half)) + self.pre_b = cast(in_w @ pre_b + in_b) + + # transformer + self.layers = [] + qd = self.heads * self.head_dim + for layer in pt.layers: + a = layer.self_attn + qkv = torch.cat([f32(a.q_proj.weight), f32(a.k_proj.weight), f32(a.v_proj.weight)], 0) + o = f32(a.o_proj.weight) * f32(layer.self_attn_layer_scale.scale)[:, None] + gu = torch.cat([f32(layer.mlp.gate_proj.weight), f32(layer.mlp.up_proj.weight)], 0) + down = f32(layer.mlp.down_proj.weight) * f32(layer.mlp_layer_scale.scale)[:, None] + self.layers.append( + dict( + ln1=cast(f32(layer.input_layernorm.weight)), + ln2=cast(f32(layer.post_attention_layernorm.weight)), + qkv=cast(qkv.t()), + o=cast(o.t()), + gu=cast(gu.t()), + down=cast(down.t()), + k_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype), + v_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype), + ) + ) + self.norm_w = cast(f32(pt.norm.weight)) + self.out_w = cast(f32(pt.output_proj.weight).t()) + self.out_b = cast(f32(pt.output_proj.bias)) + + # upsample: transposed conv (kernel = stride, no overlap) + ConvNeXt + self.ups = [] + latent = int(cfg.latent_dim) + pending = torch.zeros(latent, device=dev, dtype=torch.float32) + for trans, block in decoder.upsample: + tw = f32(trans.conv.weight) # [C_in, C_out, r] + r = tw.shape[2] + w = tw.permute(0, 2, 1).reshape(latent, r * latent) # [C_in, r * C_out] + bias = (pending @ w).view(r, latent) + f32(trans.conv.bias)[None, :] + dw = block.dwconv.conv + gamma = f32(block.gamma) + self.ups.append( + dict( + r=r, + w=cast(w), + b=cast(bias.reshape(-1)), + dw_w=cast(f32(dw.weight)[:, 0, :].t()), + dw_b=cast(f32(dw.bias)), + ln_w=cast(f32(block.norm.weight)), + ln_b=cast(f32(block.norm.bias)), + ln_eps=float(block.norm.eps), + pw1=cast(f32(block.pwconv1.weight).t()), + pw1_b=cast(f32(block.pwconv1.bias)), + pw2=cast((f32(block.pwconv2.weight) * gamma[:, None]).t()), + hist=hist(_KERNEL - 1, latent), + ) + ) + pending = f32(block.pwconv2.bias) * gamma + + # conv_in (pending bias of the last ConvNeXt added on load) + dec = decoder.decoder + cin = dec[0].conv + self.conv_in = _ConvSpec(cast(_conv_taps(f32(cin.weight))), latent, _KERNEL, 1, hist(_KERNEL - 1, latent)) + self.conv_in_pb = cast(pending) + pending = f32(cin.bias) + + # decoder blocks + self.blocks = [] + for i, blk in enumerate(dec[1:-2]): + snake, trans, *units = blk.block + c_in = trans.conv.in_channels + c_out = trans.conv.out_channels + rate = trans.conv.stride[0] + tw = f32(trans.conv.weight) # [C_in, C_out, 2 * rate] + entry = dict( + pb=cast(pending), + snake=_snake_params(snake), + rate=rate, + c_in=c_in, + c_out=c_out, + w=cast(tw.permute(0, 2, 1).reshape(c_in, 2 * rate * c_out)), + b=cast(f32(trans.conv.bias)), + prev=torch.zeros(2, S, rate, c_out, device=dev, dtype=dtype), + units=[], + ) + pending = torch.zeros(c_out, device=dev, dtype=torch.float32) + for u in units: + c1, c2 = u.conv1.conv, u.conv2.conv + entry["units"].append( + dict( + pb=cast(pending), + act1=_snake_params(u.act1), + conv1=_ConvSpec( + cast(_conv_taps(f32(c1.weight))), + c_out, + _KERNEL, + c1.dilation[0], + hist((_KERNEL - 1) * c1.dilation[0], c_out), + ), + b1=cast(f32(c1.bias)), + conv1_w4=( + cast(f32(c1.weight)).unsqueeze(2).contiguous(memory_format=torch.channels_last) + if c1.dilation[0] == 1 + else None + ), + act2=_snake_params(u.act2), + w2=cast(f32(c2.weight)[:, :, 0].t()), + ) + ) + pending = pending + f32(c2.bias) + self.blocks.append(entry) + self.out_pb = cast(pending) + self.out_snake = _snake_params(dec[-2]) + co = dec[-1].conv + self.out_c = co.in_channels + self.conv_out_w = cast(f32(co.weight)[0].t()) # [K, C] + self.conv_out_b = float(f32(co.bias)[0]) + self.conv_out = _ConvSpec(self.conv_out_w, self.out_c, _KERNEL, 1, hist(_KERNEL - 1, self.out_c)) + self.parity = torch.zeros(S, device=dev, dtype=torch.int32) + self._dummy = torch.zeros(1, device=dev, dtype=dtype) + + def _slot_state(self, slot: int) -> list[torch.Tensor]: + state = [self.parity[slot : slot + 1]] + for layer in self.layers: + state.extend((layer["k_ring"][slot], layer["v_ring"][slot])) + for spec in (self.pre, self.conv_in, self.conv_out): + assert spec.hist is not None + state.append(spec.hist[:, slot]) + for up in self.ups: + state.append(up["hist"][:, slot]) + for block in self.blocks: + state.append(block["prev"][:, slot]) + for unit in block["units"]: + state.append(unit["conv1"].hist[:, slot]) + return state + + def save_slot(self, slot: int) -> list[torch.Tensor]: + """Offload a preempted request before its runner slot is recycled.""" + return [tensor.to(device="cpu", copy=True) for tensor in self._slot_state(slot)] + + def restore_slot(self, slot: int, state: list[torch.Tensor]) -> None: + for target, saved in zip(self._slot_state(slot), state, strict=True): + target.copy_(saved) + + # ------------------------------------------------------------------ helpers + def _im2col(self, x, spec: _ConvSpec, T, slots, pos, bias=None, snake=None, ext=False): + k = spec.k + d = spec.d + c = x.shape[1] + h = (k - 1) * d + B = slots.shape[0] + if ext: + col = torch.empty(B * (h + T), c, device=x.device, dtype=self.dtype) + else: + col = torch.empty(B * T, k * c, device=x.device, dtype=self.dtype) + bu = 32 + bc = min(128, triton.next_power_of_2(c)) + grid = (B, triton.cdiv(h + T, bu), triton.cdiv(c, bc)) + a, ib = snake if snake is not None else (self._dummy, self._dummy) + _im2col_kernel[grid]( + x, bias if bias is not None else self._dummy, a, ib, + spec.hist if h > 0 else self._dummy, slots, pos, self.parity, col, + T, self.num_slots, C=c, K=k, D=d, H=h, + HAS_BIAS=bias is not None, HAS_SNAKE=snake is not None, BU=bu, BC=bc, EXT=ext, + ) # fmt: skip + return col + + def _conv_cudnn(self, x, u, T, slots, pos): + """Undilated unit conv: one source-row pass, then cuDNN on its channels-last view. + + The im2col path writes a ``K``-times patch matrix; cuDNN reads the + ``[history | new]`` rows directly (same result up to summation order). + """ + spec = u["conv1"] + ext = self._im2col(x, spec, T, slots, pos, bias=u["pb"], snake=u["act1"], ext=True) + b = slots.shape[0] + c = x.shape[1] + x4 = ext.view(b, -1, c).permute(0, 2, 1).unsqueeze(2) # [B, C, 1, H + T], channels-last memory + y = torch.nn.functional.conv2d(x4, u["conv1_w4"]) # [B, N, 1, T], channels-last memory + return y.permute(0, 2, 3, 1).reshape(b * T, -1) + + def _snake(self, x, snake, bias): + n, c = x.shape + out = torch.empty_like(x) + bu, bc = 32, min(128, triton.next_power_of_2(c)) + one = torch.zeros(1, device=x.device, dtype=torch.int32) + _im2col_kernel[(1, triton.cdiv(n, bu), triton.cdiv(c, bc))]( + x, bias, snake[0], snake[1], self._dummy, one, one, self.parity, out, + n, self.num_slots, C=c, K=1, D=1, H=0, HAS_BIAS=True, HAS_SNAKE=True, BU=bu, BC=bc, + ) # fmt: skip + return out + + # ------------------------------------------------------------------ forward + @torch.no_grad() + def forward(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + """``codes`` [B, T, NQ] int, ``slots``/``pos`` [B] int32 -> waveform [B, T * spf] float32. + + ``pos[b]`` is the index of the row's first frame in its stream; 0 starts + a new stream in ``slots[b]``. Frames beyond ``max_frames`` per call must + be split by the caller. + """ + B, T, _ = codes.shape + if T > self.max_frames: + raise ValueError(f"at most {self.max_frames} frames per call, got {T}") + dev, dt = codes.device, self.dtype + codes = codes.reshape(B * T, self.nq).to(torch.int32).contiguous() + emb = torch.empty(B * T, 2 * self.half, device=dev, dtype=dt) + _rvq_gather_kernel[(B * T,)](codes, self.books, emb, NQ=self.nq, CB=self.cb, HALF=self.half) + x = torch.addmm(self.pre_b, self._im2col(emb, self.pre, T, slots, pos), self.pre.weight) + + hid, qd = self.hidden, self.heads * self.head_dim + rows = B * T + h = torch.empty(rows, hid, device=dev, dtype=dt) + q = torch.empty(rows, qd, device=dev, dtype=dt) + att = torch.empty(rows, qd, device=dev, dtype=dt) + a = torch.empty(rows, self.layers[0]["down"].shape[0], device=dev, dtype=dt) + inter = a.shape[1] + scale = self.head_dim**-0.5 + for L in self.layers: + _rms_norm_kernel[(rows,)](x, L["ln1"], h, self.eps, N=hid) + qkv = torch.mm(h, L["qkv"]) + _rope_kv_write_kernel[(rows, self.heads)]( + qkv, q, L["k_ring"], L["v_ring"], slots, pos, T, self.theta, + NH=self.heads, HD=self.head_dim, RING=self.RING, + ) # fmt: skip + _attend_kernel[(rows, self.heads)]( + q, L["k_ring"], L["v_ring"], att, slots, pos, T, scale, + NH=self.heads, HD=self.head_dim, RING=self.RING, WINDOW=self.window, + BK=triton.next_power_of_2(self.window), + ) # fmt: skip + x.addmm_(att, L["o"]) + _rms_norm_kernel[(rows,)](x, L["ln2"], h, self.eps, N=hid) + gu = torch.mm(h, L["gu"]) + _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024) + x.addmm_(a, L["down"]) + _rms_norm_kernel[(rows,)](x, self.norm_w, h, self.eps, N=hid) + y = torch.addmm(self.out_b, h, self.out_w) # [rows, latent] + + t_rows = T + for U in self.ups: + r = U["r"] + z = torch.addmm(U["b"], y, U["w"]).view(rows * r, -1) + rows, t_rows = rows * r, t_rows * r + c = z.shape[1] + n = torch.empty_like(z) + _dwconv_ln_kernel[(B, max(t_rows, _KERNEL - 1))]( + z, U["dw_w"], U["dw_b"], U["ln_w"], U["ln_b"], U["hist"], slots, pos, self.parity, n, + t_rows, self.num_slots, U["ln_eps"], C=c, K=_KERNEL, + ) # fmt: skip + p1 = torch.mm(n, U["pw1"]) + _bias_gelu_kernel[(rows, triton.cdiv(p1.shape[1], 1024))](p1, U["pw1_b"], N=p1.shape[1], BN=1024) + z.addmm_(p1, U["pw2"]) + y = z + + x = torch.mm(self._im2col(y, self.conv_in, t_rows, slots, pos, bias=self.conv_in_pb), self.conv_in.weight) + for blk in self.blocks: + s = self._snake(x, blk["snake"], blk["pb"]) + zt = torch.mm(s, blk["w"]) # [rows, 2 * rate * c_out] + r, c_out = blk["rate"], blk["c_out"] + x = torch.empty(rows * r, c_out, device=dev, dtype=dt) + bc = min(128, triton.next_power_of_2(c_out)) + _col2im_kernel[(rows, r, triton.cdiv(c_out, bc))]( + zt, blk["b"], blk["prev"], slots, pos, self.parity, x, t_rows, self.num_slots, + R=r, C=c_out, BC=bc, + ) # fmt: skip + rows, t_rows = rows * r, t_rows * r + for u in blk["units"]: + if u.get("conv1_w4") is not None: + c1 = self._conv_cudnn(x, u, t_rows, slots, pos) + else: + col = self._im2col(x, u["conv1"], t_rows, slots, pos, bias=u["pb"], snake=u["act1"]) + c1 = torch.mm(col, u["conv1"].weight) + s2 = self._snake(c1, u["act2"], u["b1"]) + x.addmm_(s2, u["w2"]) + + ext = self._im2col(x, self.conv_out, t_rows, slots, pos, bias=self.out_pb, snake=self.out_snake, ext=True) + wav = torch.empty(B, t_rows, device=dev, dtype=torch.float32) + _conv_out_kernel[(B, triton.cdiv(t_rows, 256))]( + ext, self.conv_out_w, self.conv_out_b, wav, t_rows, + C=self.out_c, K=_KERNEL, BT=256, CP=triton.next_power_of_2(self.out_c), + ) # fmt: skip + _flip_kernel[(B,)](self.parity, slots) + return wav + + +class StreamingDecodeGraphs: + """CUDA graphs of ``StreamingCodecDecoder`` for a fixed frame count per batch bucket.""" + + def __init__(self, sd: StreamingCodecDecoder, batch_sizes: list[int], frames: int = 1) -> None: + self.sd, self.frames = sd, frames + self.graphs: dict[int, tuple] = {} + dev = sd.device + pool = torch.cuda.graph_pool_handle() + for bsz in sorted(set(batch_sizes)): + codes = torch.zeros(bsz, frames, sd.nq, dtype=torch.int32, device=dev) + slots = torch.full((bsz,), sd.scratch_slot, dtype=torch.int32, device=dev) + pos = torch.zeros(bsz, dtype=torch.int32, device=dev) + sd(codes, slots, pos) + current_omni_platform.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = sd(codes, slots, pos) + self.graphs[bsz] = (graph, codes, slots, pos, out) + self.sizes = sorted(self.graphs) + + def __call__(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + """``codes`` [n, frames, NQ], ``slots``/``pos`` [n] on device -> waveform view [n, frames * spf].""" + n = int(codes.shape[0]) + size = next((b for b in self.sizes if b >= n), None) + if size is None: + return self.sd(codes, slots.to(torch.int32), pos.to(torch.int32)) + graph, s_codes, s_slots, s_pos, out = self.graphs[size] + s_codes[:n].copy_(codes) + s_slots[:n].copy_(slots) + s_pos[:n].copy_(pos) + if n < size: + s_slots[n:].fill_(self.sd.scratch_slot) + s_pos[n:].zero_() + graph.replay() + return out[:n] 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 482514c04..33c0da7bd 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -390,16 +390,17 @@ def talker2code2wav_async_chunk( if ref_context_size > 0: if emitted_chunks <= 0: ref_context_request_id = request_id - ref_frames = ref_context.tolist() + ref_frames = list(ref_context.to(device="cpu").unbind(0)) window_frames = ref_frames + window_frames ref_context_included = True left_context_size = ref_context_size - num_quantizers = len(window_frames[0]) - num_frames = len(window_frames) - code_predictor_codes = torch.tensor( - [window_frames[f][q] for q in range(num_quantizers) for f in range(num_frames)], - dtype=torch.long, + # Keep the codebook-major CPU wire format without extracting every + # tensor element as a Python scalar. Legacy buffered frames can be lists. + code_predictor_codes = ( + torch.stack([torch.as_tensor(frame, dtype=torch.long, device="cpu") for frame in window_frames]) + .transpose(0, 1) + .reshape(-1) ) meta = MetaStruct( diff --git a/vllm_omni/worker_v2/model_states/eager_mtp.py b/vllm_omni/worker_v2/model_states/eager_mtp.py index ba347206a..03fb8c0f1 100644 --- a/vllm_omni/worker_v2/model_states/eager_mtp.py +++ b/vllm_omni/worker_v2/model_states/eager_mtp.py @@ -5,13 +5,17 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Any +import numpy as np import torch from vllm.logger import init_logger from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.states import RequestState from vllm_omni.utils.device_copy import index_to_device +from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput if TYPE_CHECKING: from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState @@ -24,12 +28,175 @@ def _has_ref_codes(buffer: dict[str, Any]) -> bool: return isinstance(ref, torch.Tensor) and ref.numel() > 0 +def _snapshot_runtime(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.to("cpu", copy=True) + if isinstance(value, dict): + return {key: _snapshot_runtime(item) for key, item in value.items()} + if isinstance(value, list): + return [_snapshot_runtime(item) for item in value] + if isinstance(value, tuple): + return tuple(_snapshot_runtime(item) for item in value) + return value + + +@dataclass +class TalkerInputs: + embeds: torch.Tensor + length: int = 0 + + +@dataclass +class SuspendedAudio: + position: int + decoder_state: list[torch.Tensor] + generator: torch.Generator | None + runtime: dict[str, Any] + frame_embed: torch.Tensor + + 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 + self._side_stream: torch.cuda.Stream | None = None + self._fast_ok: bool | None = None + self._audio_buffer: StreamingAudioBuffer | None = None + self._frame_requests: set[str] = set() + self._suspended_audio: dict[str, SuspendedAudio] = {} + self._restore_audio: dict[str, list[torch.Tensor]] = {} + self._talker_inputs: dict[str, TalkerInputs] = {} + + def suspend_audio(self, req_id: str, req_idx: int) -> None: + decoder = getattr(self.owner.model, "stream_decoder", None) + position = self.owner._stream_pos.get(req_id) + if decoder is None or position is None: + return + # Slot reuse is rare and must wait for the last side-stream write. + saved = self._restore_audio.pop(req_id, None) + if saved is None: + self._decode_stream.synchronize() + saved = decoder.save_slot(req_idx) + history = self._talker_inputs[req_id] + history.embeds = history.embeds[: history.length].to("cpu", copy=True) + runtime = self.owner.intermediate_buffer.buffers[req_idx] + assert self.owner._eager_embeds is not None + self._suspended_audio[req_id] = SuspendedAudio( + position, + saved, + self.owner._mtp_generators.get(req_id), + _snapshot_runtime(runtime), + self.owner._eager_embeds[req_idx].to("cpu", copy=True), + ) + + def resume_audio(self, req_id: str, req_idx: int) -> None: + saved = self._suspended_audio.pop(req_id, None) + if saved is not None: + self.owner._stream_pos[req_id] = saved.position + self.owner._mtp_generators[req_id] = saved.generator + self._restore_audio[req_id] = saved.decoder_state + self.owner.intermediate_buffer.buffers[req_idx] = saved.runtime + assert self.owner._eager_embeds is not None + self.owner._eager_embeds[req_idx].copy_(saved.frame_embed) + self.owner._eager_ready[req_idx] = req_id + + def finish_audio(self, req_ids: set[str]) -> None: + for req_id in req_ids: + self._suspended_audio.pop(req_id, None) + self._restore_audio.pop(req_id, None) + self.owner._stream_pos.pop(req_id, None) + self._talker_inputs.pop(req_id, None) + if self._audio_buffer is not None: + self._audio_buffer.finish(req_ids) + + def replay_inputs(self, req_id: str, req_idx: int, offset: int, ids: torch.Tensor, embeds: torch.Tensor) -> bool: + """Rebuild freed Talker KV from the exact conditioned inputs, without rerunning MTP.""" + if req_id not in self._restore_audio: + return False + history = self._talker_inputs[req_id] + count = min(len(ids), max(0, history.length - offset)) + if count: + embeds[:count].copy_(history.embeds[offset : offset + count]) + if count < len(ids): + # Recompute ends with the last accepted CB0, whose full eager frame + # was saved but has not yet been consumed as a Talker input. + if len(ids) - count != 1 or offset + count != history.length: + raise RuntimeError(f"Talker replay exceeds saved inputs for {req_id}") + info = self.owner.intermediate_buffer.buffers[req_idx] + _, _, updates = self.owner.model.preprocess(ids[count:], embeds[count:], **info, _omni_is_prefill=False) + mtp_inputs = updates.pop("mtp_inputs", None) + if mtp_inputs is None: + raise RuntimeError(f"Talker replay has no text conditioning for {req_id}") + assert self.owner._eager_embeds is not None + embeds[count:].copy_(self.owner._eager_embeds[req_idx] + mtp_inputs[1].reshape(1, -1)) + self.owner.intermediate_buffer.update(req_idx, updates, self.owner.model.gpu_resident_buffer_keys) + return True + + def record_inputs(self, input_batch: InputBatch, embeds: torch.Tensor) -> None: + if getattr(self.owner.model, "stream_decoder", None) is None: + return + history_slices: list[torch.Tensor] = [] + input_slices: list[torch.Tensor] = [] + for i, req_id in enumerate(input_batch.req_ids): + start, end = input_batch.query_start_loc_np[i : i + 2] + offset = int(input_batch.num_computed_tokens_np[i]) + length = offset + int(end - start) + history = self._talker_inputs.get(req_id) + if history is None or history.embeds.device != embeds.device: + storage = torch.empty( + self.owner.vllm_config.model_config.max_model_len, + embeds.shape[-1], + device=embeds.device, + dtype=embeds.dtype, + ) + if history is not None: + storage[: history.length].copy_(history.embeds) + history = TalkerInputs(storage, history.length if history else 0) + self._talker_inputs[req_id] = history + history_slices.append(history.embeds[offset:length]) + input_slices.append(embeds[start:end]) + history.length = max(history.length, length) + if history_slices: + # Keep exact inputs for replay without launching one copy per row. + torch._foreach_copy_(history_slices, input_slices) + + def prepare_audio_output( + self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any] + ) -> StreamingAudioOutput | None: + if self._audio_buffer is None: + return None + wav = outputs.pop("model_outputs", None) + if not isinstance(wav, torch.Tensor): + return None + n = input_batch.num_reqs + total_after = input_batch.num_computed_tokens_np[:n] + np.asarray(input_batch.num_scheduled_tokens[:n]) + 1 + limits = np.minimum( + req_states.max_seq_len[input_batch.idx_mapping_np[:n]], self.owner.vllm_config.model_config.max_model_len + ) + requests = [ + self._audio_buffer.requests.get(rid) if rid in self._frame_requests else None for rid in input_batch.req_ids + ] + self._frame_requests.clear() + event = self.owner._stream_decode_event + self.owner._stream_decode_event = None + return StreamingAudioOutput( + wav, + outputs["meta"]["codec_frame_valid"], + outputs["sr"][0], + event, + input_batch.query_start_loc_np[: n + 1].copy(), + requests, + total_after >= limits, + ) + + @property + def _decode_stream(self) -> torch.cuda.Stream: + """Side stream for Talker stream decode, so the codec overlaps the next step.""" + if self._side_stream is None: + self._side_stream = torch.cuda.Stream() + return self._side_stream def set_first_audio_sink(self, sink: Any) -> None: from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender @@ -106,6 +273,14 @@ class EagerMTPState: bsz = len(entries) device = text_hidden.device + if self._fast_path_ok(codes_out, text_hidden): + self._run_eager_mtp_fused( + input_batch, text_hidden, sampled_token_ids, multimodal_outputs, entries, input_ids, + codes_out, valid_out, mtp_batch_descriptor_dispatcher, + ) # fmt: skip + return + if getattr(self.owner.model, "stream_decoder", None) is not None: + raise RuntimeError("Talker stream decode requires the fused eager-MTP path") 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 @@ -162,6 +337,168 @@ class EagerMTPState: scheduled *= self._first_audio_valid.index_select(0, rows) first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype)) + def _embedding_weight(self) -> torch.Tensor | None: + inner = getattr(self.owner.model, "model", None) + embed = getattr(inner, "embed_tokens", None) + weight = getattr(embed, "weight", None) + method = getattr(embed, "quant_method", None) + unquantized = method is None or type(method).__name__ == "UnquantizedEmbeddingMethod" + tp = getattr(embed, "tp_size", 1) + if isinstance(weight, torch.Tensor) and weight.ndim == 2 and unquantized and tp == 1: + return weight + return None + + def _fast_path_ok(self, codes_out: torch.Tensor, text_hidden: torch.Tensor) -> bool: + if self._fast_ok is None: + self._fast_ok = ( + text_hidden.is_cuda + and not self.owner.vllm_config.cache_config.enable_prefix_caching + and self._embedding_weight() is not None + and hasattr(self.owner.model, "_codebook_vocab_size") + ) + return self._fast_ok and codes_out.ndim == 2 and text_hidden.is_contiguous() + + def _run_eager_mtp_fused( + self, + input_batch: InputBatch, + text_hidden: torch.Tensor, + sampled_token_ids: torch.Tensor, + multimodal_outputs: dict[str, Any], + entries: list[tuple[int, int, str, bool]], + input_ids: torch.Tensor, + codes_out: torch.Tensor, + valid_out: torch.Tensor, + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None, + ) -> None: + from vllm_omni.worker_v2.model_states.eager_mtp_kernels import eager_post, eager_pre + + owner = self.owner + bsz = len(entries) + device = text_hidden.device + first_audio = multimodal_outputs.get("meta", {}).get("first_audio") + fa_requests = owner._first_audio_requests + model = owner.model + stream = getattr(model, "stream_decoder", None) + stream_out = multimodal_outputs.get("model_outputs") if stream is not None else None + pos_list: list[int] = [] + first_rows: list[int] = [] + primes: list[tuple[int, torch.Tensor]] = [] + if isinstance(stream_out, torch.Tensor): + assert stream is not None + # Frame index of each row's frame in its request's stream; 0 starts + # a fresh decoder state in the request's slot. A voice-clone + # request's stream starts after its reference codes, which prime + # the decoder state first (as Code2Wav's first-chunk context does). + if self._audio_buffer is None: + self._audio_buffer = StreamingAudioBuffer(int(stream.spf), int(model.stream_chunk_frames)) + self._frame_requests = {req_id for _i, _idx, req_id, _p in entries} + positions = owner._stream_pos + for row, (_i, req_idx, req_id, _p) in enumerate(entries): + self._audio_buffer.add(req_id) + p = positions.get(req_id) + if p is None: + first_rows.append(row) + ref = model.get_stream_ref_context(owner.intermediate_buffer.buffers[req_idx]) + p = 0 + if ref is not None: + primes.append((req_idx, ref)) + p = int(ref.shape[0]) + pos_list.append(p) + positions[req_id] = p + 1 + meta_list = ( + [i for i, _r, _q, _p in entries] + + [r for _i, r, _q, _p in entries] + + [int(p) for _i, _r, _q, p in entries] + + [int(q in fa_requests) for _i, _r, q, _p in entries] + + pos_list + ) + meta = index_to_device(meta_list, device) + rows = meta[bsz : 2 * bsz] + sampled = sampled_token_ids.reshape(input_batch.num_reqs, -1) + assert owner._mtp_input_ids is not None and owner._mtp_input_embeds is not None + assert owner._mtp_hidden is not None and owner._mtp_text_step is not None + batch_ids = owner._mtp_input_ids[:bsz] + batch_emb = owner._mtp_input_embeds[:bsz] + batch_hidden = owner._mtp_hidden[:bsz] + batch_step = owner._mtp_text_step[:bsz] + last_tokens, layer0 = eager_pre( + meta, bsz, input_batch.query_start_loc, sampled, text_hidden, self._embedding_weight(), + batch_ids, batch_emb, batch_hidden, batch_step, + ) # fmt: skip + req_indices = meta_list[bsz : 2 * bsz] + frame_embeds, codes = owner._mtp_forward( + req_indices, batch_ids, batch_emb, batch_hidden, batch_step, mtp_batch_descriptor_dispatcher + ) + assert codes is not None + has_prefill = any(p for _i, _r, _q, p in entries) + fa_in_kernel = isinstance(first_audio, torch.Tensor) and not has_prefill + valid = eager_post( + meta, bsz, last_tokens, layer0, frame_embeds.reshape(frame_embeds.shape[0], -1).contiguous(), + owner._eager_embeds, codes.contiguous(), codes_out, input_ids, valid_out, + first_audio if fa_in_kernel else None, self._first_audio_valid if fa_in_kernel else None, + owner.model._codebook_vocab_size, + ) # fmt: skip + for _i, req_idx, req_id, _prefill in entries: + owner._eager_ready[req_idx] = req_id + if pos_list: + assert stream is not None + assert isinstance(stream_out, torch.Tensor) + decode = model.stream_graphs if model.stream_graphs is not None else stream + frame_codes = codes[:bsz].reshape(bsz, 1, -1).to(torch.int32) + # The next Talker step does not read the PCM, so the codec runs + # beside it; only the output copies wait for it. + side = self._decode_stream + side.wait_stream(torch.cuda.current_stream()) + for t in (frame_codes, meta, last_tokens, valid, stream_out): + t.record_stream(side) + with torch.cuda.stream(side): + for _i, req_idx, req_id, _p in entries: + saved = self._restore_audio.pop(req_id, None) + if saved is not None: + stream.restore_slot(req_idx, saved) + model.prime_stream_decoder(primes) + pcm = decode(frame_codes, rows, meta[4 * bsz : 5 * bsz]) + stream_out.index_copy_(0, last_tokens, pcm.reshape(bsz, -1).to(stream_out.dtype)) + self._send_stream_first_frames(entries, pcm, valid, first_rows) + done = torch.cuda.Event() + done.record(side) + owner._stream_decode_event = done + if has_prefill: + self._publish_first_audio(input_batch, entries, codes[:bsz], valid) + if isinstance(first_audio, torch.Tensor): + # Requests accepted by this step's publish count as scheduled. + scheduled = index_to_device( + [int(q in owner._first_audio_requests) for _i, _r, q, _p 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 _send_stream_first_frames( + self, + entries: list[tuple[int, int, str, bool]], + pcm: torch.Tensor, + valid: torch.Tensor, + first_rows: list[int], + ) -> None: + sender = self.owner._first_audio_sender + if sender is None or not first_rows: + return + rows = index_to_device(first_rows, pcm.device) + request_ids = [entries[row][2] for row in first_rows] + sample_rate = torch.tensor(int(self.owner.model.stream_sample_rate), dtype=torch.int32) + accepted = sender.submit( + request_ids, + pcm.reshape(pcm.shape[0], -1).index_select(0, rows).float(), + sample_rate, + valid=valid.index_select(0, rows), + ) + assert self._audio_buffer is not None + for request_id in accepted: + self._audio_buffer.requests[request_id].accept_first_audio() + def _publish_first_audio( self, input_batch: InputBatch, diff --git a/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py b/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py new file mode 100644 index 000000000..c48f87875 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Gather/scatter kernels around the eager Talker-MTP call. + +Each engine step, the rows sampled this step need their CB0, codec embedding +and last hidden state gathered into the MTP inputs, and the MTP results +scattered back into the eager-frame embeddings and the retained outputs. As +separate torch ops that is a few dozen small launches per step, and the engine +loop is host-bound, so each costs the whole batch. These two kernels do the +same indexing in one launch each; the row metadata arrives in one H2D copy. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _pre_kernel( + meta_ptr, # int64 [4, n]: batch row, request index, prefill flag, first-audio scheduled + n, + qsl_ptr, # query_start_loc + sampled_ptr, # [num_reqs, s_stride] + s_stride, + hidden_ptr, # [T, H] + emb_w_ptr, # [V, H] + ids_ptr, # [n] MTP input ids + emb_ptr, # [n, H] + mtp_hidden_ptr, # [n, H] + step_ptr, # [n, H] + last_ptr, # int64 [n] out + layer0_ptr, # int64 [n] out + H: tl.constexpr, + BH: tl.constexpr, +): + b = tl.program_id(0) + row = tl.load(meta_ptr + b) + last = tl.load(qsl_ptr + row + 1).to(tl.int64) - 1 + l0 = tl.load(sampled_ptr + row * s_stride).to(tl.int64) + tl.store(ids_ptr + b, l0.to(ids_ptr.dtype.element_ty)) + tl.store(last_ptr + b, last) + tl.store(layer0_ptr + b, l0) + for h0 in tl.static_range(0, H, BH): + c = h0 + tl.arange(0, BH) + m = c < H + tl.store(emb_ptr + b * H + c, tl.load(emb_w_ptr + l0 * H + c, mask=m).to(emb_ptr.dtype.element_ty), mask=m) + tl.store(mtp_hidden_ptr + b * H + c, tl.load(hidden_ptr + last * H + c, mask=m), mask=m) + tl.store(step_ptr + b * H + c, tl.zeros([BH], dtype=step_ptr.dtype.element_ty), mask=m) + + +@triton.jit +def _post_kernel( + meta_ptr, + n, + last_ptr, + layer0_ptr, + frame_ptr, # [>=n, H] MTP frame embeddings + eager_ptr, # [R, H] + codes_ptr, # [>=n, Q] + codes_out_ptr, # [T, Q] + input_ids_ptr, # [T] this step's input ids + valid_out_ptr, # [T] + valid_ptr, # [n] bool out + fa_ptr, # [T] first_audio output or dummy + fa_valid_ptr, # [R] bool or dummy + vocab, + H: tl.constexpr, + Q: tl.constexpr, + QP: tl.constexpr, + BH: tl.constexpr, + WRITE_FA: tl.constexpr, + HAS_FA_VALID: tl.constexpr, +): + b = tl.program_id(0) + req = tl.load(meta_ptr + n + b) + prefill = tl.load(meta_ptr + 2 * n + b) != 0 + last = tl.load(last_ptr + b) + l0 = tl.load(layer0_ptr + b) + for h0 in tl.static_range(0, H, BH): + c = h0 + tl.arange(0, BH) + m = c < H + v = tl.load(frame_ptr + b * H + c, mask=m) + tl.store(eager_ptr + req * H + c, v.to(eager_ptr.dtype.element_ty), mask=m) + q = tl.arange(0, QP) + qm = q < Q + cv = tl.load(codes_ptr + b * Q + q, mask=qm) + tl.store(codes_out_ptr + last * Q + q, cv.to(codes_out_ptr.dtype.element_ty), mask=qm) + inp = tl.load(input_ids_ptr + last).to(tl.int64) + in_valid = ((inp >= 0) & (inp < vocab)) | prefill + valid = (l0 >= 0) & (l0 < vocab) & in_valid + tl.store(valid_out_ptr + last, valid.to(valid_out_ptr.dtype.element_ty)) + tl.store(valid_ptr + b, valid) + if WRITE_FA: + scheduled = tl.load(meta_ptr + 3 * n + b) != 0 + if HAS_FA_VALID: + scheduled = scheduled & (tl.load(fa_valid_ptr + req) != 0) + else: + scheduled = scheduled & False + tl.store(fa_ptr + last, scheduled.to(fa_ptr.dtype.element_ty)) + + +def eager_pre(meta, n, qsl, sampled, hidden, emb_w, ids, emb, mtp_hidden, step): + last = torch.empty(n, dtype=torch.long, device=hidden.device) + layer0 = torch.empty(n, dtype=torch.long, device=hidden.device) + H = hidden.shape[-1] + _pre_kernel[(n,)]( + meta, n, qsl, sampled, sampled.stride(0), hidden, emb_w, ids, emb, mtp_hidden, step, last, layer0, + H=H, BH=min(1024, triton.next_power_of_2(H)), + ) # fmt: skip + return last, layer0 + + +def eager_post(meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, fa, fa_valid, vocab): + valid = torch.empty(n, dtype=torch.bool, device=last.device) + H = eager.shape[-1] + Q = codes_out.shape[-1] + dummy = valid + _post_kernel[(n,)]( + meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, valid, + fa if fa is not None else dummy, fa_valid if fa_valid is not None else dummy, int(vocab), + H=H, Q=Q, QP=triton.next_power_of_2(Q), BH=min(1024, triton.next_power_of_2(H)), + WRITE_FA=fa is not None, HAS_FA_VALID=fa_valid is not None, + ) # fmt: skip + return valid 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 f20a0f2db..f10149f46 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -41,6 +41,7 @@ from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState from vllm_omni.worker_v2.model_states.intermediate_buffer import ( OmniIntermediateBuffer, ) +from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput logger = init_logger(__name__) @@ -85,6 +86,9 @@ def _make_safe_get_rope(orig_get_rope): return _safe_get_rope +_NO_SEED = object() + + class OmniModelState(DefaultModelState): """Generic Omni ``ModelState`` — works for **all** Omni model stages. @@ -146,7 +150,7 @@ class OmniModelState(DefaultModelState): self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False) self._decode_preprocess = self._resolve_decode_preprocess(model) self._decode_preprocess_is_identity = bool(getattr(model, "mrv2_decode_preprocess_is_identity", False)) - self._mtp_generators: dict[str, torch.Generator] = {} + self._mtp_generators: dict[str, torch.Generator | None] = {} # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim. self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0 @@ -167,6 +171,10 @@ class OmniModelState(DefaultModelState): self._mtp_offsets: torch.Tensor | None = None self._mtp_sample_uniforms: torch.Tensor | None = None self._mtp_runner: Any | None = None + # Talker stream decode: each request's next frame position, and the + # side-stream event the step's PCM output copy must wait on. + self._stream_pos: dict[str, int] = {} + self._stream_decode_event: torch.cuda.Event | None = None if self._embed_dim > 0 and hasattr(model, "mtp"): max_bs = max_num_reqs self._mtp_input_ids = torch.zeros(max_bs, dtype=torch.long, device=device) @@ -218,6 +226,17 @@ class OmniModelState(DefaultModelState): input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher ) + def prepare_streaming_audio_output( + self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any] + ) -> StreamingAudioOutput | None: + return self._eager_state.prepare_audio_output(input_batch, req_states, outputs) + + def on_requests_finished(self, req_ids: set[str]) -> None: + self._eager_state.finish_audio(req_ids) + + def on_request_preempted(self, req_id: str, req_index: int) -> None: + self._eager_state.suspend_audio(req_id, req_index) + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: """Let the model replace or wrap the MRv2 sampler (``mrv2_custom_sampler``).""" hook = getattr(self.model, "mrv2_custom_sampler", None) @@ -372,6 +391,7 @@ class OmniModelState(DefaultModelState): def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) self.intermediate_buffer.add_request(req_index, new_req_data) + self._eager_state.resume_audio(new_req_data.req_id, req_index) self._initialize_upstream_warmup_buffer(req_index, new_req_data.req_id) def _initialize_upstream_warmup_buffer(self, req_index: int, req_id: str) -> None: @@ -407,6 +427,7 @@ 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) + self._stream_pos.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) @@ -673,6 +694,7 @@ class OmniModelState(DefaultModelState): 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]] = [] + replay_rows: list[tuple[int, int, str, bool]] = [] # Declared identity decode has no per-row updates or hook work. skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None @@ -684,6 +706,20 @@ class OmniModelState(DefaultModelState): continue start = int(input_batch.query_start_loc_np[i]) n_tok = int(input_batch.num_scheduled_tokens[i]) + if ( + self._eager_mtp + and self.model.stream_decoder is not None + and self._eager_state.replay_inputs( + buf["req_id"], + req_idx, + int(input_batch.num_computed_tokens_np[i]), + input_ids[start : start + n_tok], + embeds[start : start + n_tok], + ) + ): + if input_batch.num_computed_prefill_tokens_np[i] + n_tok >= input_batch.prefill_len_np[i]: + replay_rows.append((i, req_idx, buf["req_id"], True)) + continue 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 @@ -725,7 +761,8 @@ class OmniModelState(DefaultModelState): 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], + + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows] + + replay_rows, input_ids, ) @@ -760,7 +797,7 @@ class OmniModelState(DefaultModelState): 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_offsets = index_to_device(starts, input_ids.device) batch_ids = input_ids.index_select(0, batch_offsets) batch_embeds = embeds.index_select(0, batch_offsets) @@ -866,6 +903,9 @@ class OmniModelState(DefaultModelState): prepacked_mtp_inputs=prepacked_mtp_inputs, ) + if self._eager_mtp: + self._eager_state.record_inputs(input_batch, embeds) + def publish_sampled_embeddings( self, input_batch: InputBatch, sampled_token_ids: torch.Tensor ) -> tuple[dict[str, Any], torch.cuda.Event] | None: @@ -989,7 +1029,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 index_to_device([start for _i, start, _mtp in mtp_batches], device) def _run_batched_mtp( self, @@ -1180,15 +1220,21 @@ class OmniModelState(DefaultModelState): sampling_params: Any, device: torch.device, ) -> torch.Generator | None: - resolve_seed = getattr(self.model, "get_mtp_seed", None) - seed = resolve_seed(sampling_params) if resolve_seed is not None else None - if seed is None: - return None - cache = getattr(self, "_mtp_generators", None) if cache is None: cache = {} self._mtp_generators = cache + generator = cache.get(req_id, _NO_SEED) + if generator is None: + # Resolved earlier: this request has no seed. + return None + if generator is not _NO_SEED and generator.device == device: + return generator + resolve_seed = getattr(self.model, "get_mtp_seed", None) + seed = resolve_seed(sampling_params) if resolve_seed is not None else None + if seed is None: + cache[req_id] = None + return None generator = cache.get(req_id) if generator is None or generator.device != device: generator = torch.Generator(device=device) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 32622e8ef..73434b8a1 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -14,6 +14,7 @@ Extends ``OmniGPUModelRunner`` with: from __future__ import annotations from contextlib import nullcontext +from functools import partial from typing import Any, cast import numpy as np @@ -37,6 +38,7 @@ from vllm_omni.utils.mm_outputs import partition_flat_payload from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner from vllm_omni.worker_v2.omni_sampler import sample_with_output from vllm_omni.worker_v2.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot, pack_output_snapshot +from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput logger = init_logger(__name__) _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64 @@ -250,6 +252,13 @@ class OmniARModelRunner(OmniGPUModelRunner): model_runner_output.kv_extracted_req_ids = kv_extracted model_runner_output._async_chunk = bool(getattr(self.model_config, "async_chunk", False)) + prepare_streaming = getattr(self.model_state, "prepare_streaming_audio_output", None) + streaming_audio = ( + prepare_streaming(input_batch, self.req_states, multimodal_outputs) + if need_pooler and multimodal_outputs and callable(prepare_streaming) + else None + ) + # --- Async D2H via OmniAsyncOutput --- materialize_native = self._uses_native_output_materializer() async_output = OmniAsyncOutput( @@ -266,6 +275,7 @@ class OmniARModelRunner(OmniGPUModelRunner): finalize_multimodal=sampling_output.finalize_multimodal, check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, + streaming_audio=streaming_audio, extra_multimodal_outputs=extra_outputs, ) self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event) @@ -297,6 +307,10 @@ class OmniARModelRunner(OmniGPUModelRunner): def _retain_multimodal_outputs(self, outputs: dict[str, Any]) -> dict[str, Any]: if not bool(getattr(self.model_config, "async_chunk", False)) or not outputs: return outputs + if getattr(self.model, "mm_outputs_fresh_per_step", False): + # Freshly allocated each step and filled after sampling; nothing + # overwrites them before the host copy. + return outputs slot_index = self._async_mm_snapshot_cursor if self._async_mm_snapshot_pending[slot_index]: raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().") @@ -516,6 +530,12 @@ def _async_copy_mm_value( for key, val in value.items() } if isinstance(value, list): + if len(value) > 1 and isinstance(value[0], torch.Tensor) and value[0].device.type == "cpu": + first = value[0] + if all(val is first for val in value): + # One shared per-step host tensor (e.g. every request's sample + # rate): one copy serves all entries. + return [_async_copy_tensor(first)] * len(value) return [ _async_copy_mm_value( val, @@ -679,6 +699,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): finalize_multimodal: Any | None = None, check_ep_fault: bool = False, routed_experts: RoutedExpertsTensors | None = None, + streaming_audio: StreamingAudioOutput | None = None, extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None, ): self.model_runner_output = model_runner_output @@ -768,17 +789,31 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._hidden_cpu: torch.Tensor | None = None self._mm_cpu: dict[str, Any] = {} self._mm_snapshot: dict[str, Any] = {} + self._streaming_audio = ( + streaming_audio.to_cpu( + copy_stream, partial(_async_copy_tensor, copy_stream=copy_stream, pin_memory=pin_memory) + ) + if streaming_audio is not None + else None + ) if self._need_pooler and (self._async_chunk or self._finalize_multimodal is not None): # CUDA graph replay reuses the model's output buffers. Take # ownership directly in pinned host memory on the output copy # stream so deferred finalization never performs a blocking # D2H copy on the runner thread. - self._mm_snapshot = _async_copy_mm( - multimodal_outputs, - self._total_tokens, - copy_stream=copy_stream, - pin_memory=pin_memory, - ) + if ( + streaming_audio is None + or self._finalize_multimodal is not None + or extra_multimodal_outputs is not None + ): + # PCM owns its request partition. Generic codes/meta are + # discarded unless a finalizer or extra payload needs them. + self._mm_snapshot = _async_copy_mm( + multimodal_outputs, + self._total_tokens, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) if extra_multimodal_outputs: # Produced after sampling on the producer stream; copy only # once its completion event has been observed. @@ -856,7 +891,9 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): pooler_client is not None and len(pooler_client) != self._num_reqs ): raise ValueError("Model-owned output snapshot does not match the request batch") - else: + elif self._streaming_audio is None: + # In-stage PCM already owns its request partition. Building + # generic code payloads here would immediately discard them. pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm( self._mm_snapshot, self._query_start_loc_np, @@ -865,6 +902,8 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._total_tokens, self._padded_total_tokens, ) + if self._streaming_audio is not None: + pooler_inter, pooler_client = None, self._streaming_audio.get_output() self.model_runner_output.pooler_output = None if self._async_chunk else pooler_inter self.model_runner_output.inter_stage_outputs = pooler_inter self.model_runner_output.multimodal_outputs = ( diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 10f4053cf..9e73201e5 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -437,6 +437,10 @@ class OmniGPUModelRunner(GPUModelRunner): capture_first_frame = getattr(self.model, "capture_first_frame_graphs", None) if callable(capture_first_frame): capture_first_frame() + capture_stream = getattr(self.model, "capture_stream_decode_graphs", None) + if callable(capture_stream): + sizes = [int(s) for s in self.model_state._get_mtp_capture_sizes()] + capture_stream(sorted(set(sizes + [1]))) return result def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any: @@ -746,7 +750,16 @@ class OmniGPUModelRunner(GPUModelRunner): # Request-owned state outlives generation slots and chunk boundaries. # Notify even if the last chunk already released its runner slot. on_finished(finished) + finish_state = getattr(self.model_state, "on_requests_finished", None) + if finished and callable(finish_state): + finish_state(finished) preempted = scheduler_output.preempted_req_ids + suspend_state = getattr(self.model_state, "on_request_preempted", None) + if preempted and callable(suspend_state): + for req_id in preempted - finished: + idx = self.req_states.req_id_to_index.get(req_id) + if idx is not None: + suspend_state(req_id, idx) all_done = finished | preempted if preempted else finished for req_id in all_done: idx = self.req_states.req_id_to_index.get(req_id) diff --git a/vllm_omni/worker_v2/streaming_audio.py b/vllm_omni/worker_v2/streaming_audio.py new file mode 100644 index 000000000..bf78de678 --- /dev/null +++ b/vllm_omni/worker_v2/streaming_audio.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request-owned PCM accumulation for a stage that decodes audio in place.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field + +import numpy as np +import torch + +from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + +@dataclass +class AudioRequest: + samples_per_frame: int + chunk_frames: int + active: bool = True + emitted: bool = False + direct_first_pending: bool = False + first_audio_required: bool = False + pending: list[torch.Tensor | np.ndarray] = field(default_factory=list) + pending_samples: int = 0 + lock: threading.Lock = field(default_factory=threading.Lock) + + def accept_first_audio(self) -> None: + with self.lock: + if self.active: + self.direct_first_pending = True + + def push(self, pcm: torch.Tensor | np.ndarray, ended: bool) -> torch.Tensor | None: + with self.lock: + if not self.active: + return None + size = pcm.numel() if isinstance(pcm, torch.Tensor) else pcm.size + if self.direct_first_pending: + self.direct_first_pending = False + if size: + pcm = pcm[self.samples_per_frame :] + size = max(0, size - self.samples_per_frame) + self.emitted = True + self.first_audio_required = True + if size: + self.pending.append(pcm) + self.pending_samples += size + frames = self.pending_samples // self.samples_per_frame + out = None + if frames and (ended or not self.emitted or frames >= self.chunk_frames): + if all(isinstance(part, np.ndarray) for part in self.pending): + array = np.concatenate(self.pending) if len(self.pending) > 1 else self.pending[0] + out = torch.from_numpy(array) + else: + parts = [torch.from_numpy(part) if isinstance(part, np.ndarray) else part for part in self.pending] + out = torch.cat(parts) if len(parts) > 1 else parts[0] + self.pending.clear() + self.pending_samples = 0 + self.emitted = True + if ended: + self.active = False + self.pending.clear() + self.pending_samples = 0 + return out + + def finish(self) -> None: + with self.lock: + self.active = False + self.pending.clear() + self.pending_samples = 0 + + +class StreamingAudioBuffer: + def __init__(self, samples_per_frame: int, chunk_frames: int) -> None: + self.samples_per_frame = samples_per_frame + self.chunk_frames = chunk_frames + self.requests: dict[str, AudioRequest] = {} + + def add(self, request_id: str) -> AudioRequest: + if request_id not in self.requests: + self.requests[request_id] = AudioRequest(self.samples_per_frame, self.chunk_frames) + return self.requests[request_id] + + def finish(self, request_ids: Iterable[str]) -> None: + for request_id in request_ids: + state = self.requests.pop(request_id, None) + if state is not None: + # Already queued outputs retain this object, never a fresh lookup + # by ID. A late copy cannot resurrect a cancelled request. + state.finish() + + +@dataclass +class StreamingAudioOutput: + wav: torch.Tensor + valid: torch.Tensor + sample_rate: torch.Tensor + event: torch.cuda.Event | None + query_start_loc: np.ndarray + requests: list[AudioRequest | None] + length_end: np.ndarray + + def to_cpu( + self, + copy_stream: torch.cuda.Stream, + copy_tensor: Callable[[torch.Tensor], torch.Tensor], + ) -> StreamingAudioOutput: + if self.event is not None: + copy_stream.wait_event(self.event) + return StreamingAudioOutput( + copy_tensor(self.wav), + copy_tensor(self.valid), + self.sample_rate, + None, + self.query_start_loc, + self.requests, + self.length_end, + ) + + def get_output(self) -> list[dict[str, torch.Tensor] | None]: + valid = self.valid.numpy().astype(bool) + # Each D2H batch owns its host allocation. NumPy slices retain that + # backing tensor while a request accumulates PCM across steps; no + # per-request Torch views are needed until an audio chunk is emitted. + # Keep the tensor path for dtypes NumPy cannot represent (e.g. BF16). + wav = self.wav.numpy() if self.wav.dtype in (torch.float16, torch.float32, torch.float64) else None + output: list[dict[str, torch.Tensor] | None] = [] + required_marker = None + for i, state in enumerate(self.requests): + payload = None + start, end = self.query_start_loc[i : i + 2] + if state is not None and end > start: + if wav is not None: + if end - start == 1: + pcm = wav[start:end].reshape(-1) if valid[start] else wav[0:0].reshape(-1) + else: + pcm = wav[start:end][valid[start:end]].reshape(-1) + elif end - start == 1: + pcm = self.wav[start:end].reshape(-1) if valid[start] else self.wav.new_empty(0) + else: + pcm = self.wav[start:end][torch.from_numpy(valid[start:end])].reshape(-1) + ended = bool(self.length_end[i]) or not valid[end - 1] + chunk = state.push(pcm, ended) + if chunk is not None: + payload = {"model_outputs": chunk, "sr": self.sample_rate} + if state.first_audio_required and (chunk is not None or ended): + if required_marker is None: + required_marker = torch.tensor(True) + if payload is None: + payload = {} + payload[FIRST_AUDIO_REQUIRED_KEY] = required_marker + output.append(payload) + return output