From a9e2117f5e18a882332501a3d96c6eaa076a9310 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sun, 20 Sep 2026 23:03:08 -0700 Subject: [PATCH] [Core][Model] Enable Qwen3-TTS MRV2 and optimize the TTS pipeline (#7781) Signed-off-by: Sy03 <1370724210@qq.com> --- .../common/ci_source_file_dependencies.yml | 10 + .buildkite/cuda/test-ready.yml | 2 +- docs/configuration/stage_configs.md | 87 ++ tests/config/test_config_factory.py | 53 + tests/config/test_omni_config.py | 66 + .../test_generation_scheduler_restore.py | 157 ++- ...t_omni_ar_scheduler_aborted_queue_sweep.py | 4 + .../sched/test_omni_ar_scheduler_logprobs.py | 22 +- .../test_omni_ar_scheduler_output_policy.py | 43 + .../sched/test_omni_ar_scheduler_streaming.py | 45 + ...ni_ar_scheduler_unified_decode_deferral.py | 4 + .../test_omni_sched_deferred_free_fence.py | 3 + .../sched/test_omni_scheduler_mixin_shared.py | 10 +- .../sched/test_omni_scheduler_ready_inbox.py | 38 + ..._omni_scheduler_streaming_input_counter.py | 4 + .../sched/test_omni_scheduling_coordinator.py | 72 +- .../offline_inference/test_qwen3_tts_base.py | 8 +- .../e2e/online_serving/test_qwen3_tts_base.py | 12 +- tests/engine/test_orchestrator.py | 10 +- tests/engine/test_output_processor.py | 14 + .../engine/test_output_processor_mrv2_text.py | 137 ++ tests/engine/test_request_snapshot.py | 21 + tests/engine/test_single_stage_mode.py | 13 +- tests/engine/test_stage_engine_args.py | 15 +- .../test_terminal_stage_metrics_order.py | 82 ++ .../openai_api/test_ref_audio_cache.py | 95 ++ .../qwen3_tts/test_code_predictor_dtype.py | 36 + .../qwen3_tts/test_cuda_graph_decoder.py | 24 +- .../test_qwen3_tts_output_async_cuda.py | 55 + .../test_qwen3_tts_stateless_capture.py | 100 ++ .../test_qwen3_tts_talker_preprocess.py | 86 +- .../test_qwen3_tts_talker_ref_codes.py | 56 +- .../test_qwen3_tts_talker_silence_ban.py | 24 +- .../test_qwen3_tts_async_chunk.py | 203 ++- .../worker/test_min_tokens_stop_ids_guard.py | 21 +- tests/worker/test_omni_connector_mixin.py | 13 + tests/worker/test_omni_gpu_model_runner.py | 6 +- tests/worker/test_omni_v2_profile_dispatch.py | 60 + tests/worker/test_tts_sampling_seed.py | 34 + tests/worker/test_worker_base.py | 2 +- tests/worker_v2/__init__.py | 0 tests/worker_v2/test_intermediate_buffer.py | 96 ++ ...test_intermediate_buffer_cuda_residency.py | 28 + tests/worker_v2/test_omni_ar_model_runner.py | 253 ++++ tests/worker_v2/test_omni_data_plane.py | 481 +++++++ .../test_omni_generation_model_runner.py | 184 +++ tests/worker_v2/test_omni_gpu_model_runner.py | 166 +++ tests/worker_v2/test_omni_model_state.py | 361 +++++ vllm_omni/config/model.py | 6 + vllm_omni/config/omni_config.py | 12 + vllm_omni/config/stage_config.py | 67 +- vllm_omni/core/sched/omni_ar_scheduler.py | 55 +- .../core/sched/omni_generation_scheduler.py | 186 ++- vllm_omni/core/sched/omni_scheduler_mixin.py | 163 ++- .../core/sched/omni_scheduling_coordinator.py | 240 +++- vllm_omni/core/sched/output.py | 9 +- vllm_omni/data_entry_keys.py | 12 + .../qwen3_tts_high_concurrency_mrv2.yaml | 28 + .../qwen3_tts_high_concurrency_mrv2_b4.yaml | 8 + ..._tts_high_concurrency_mrv2_single_gpu.yaml | 23 + vllm_omni/deploy/qwen3_tts_mrv2.yaml | 18 + .../connectors/shm_connector.py | 2 +- .../omni_connector_payload_transport.py | 454 ++++-- .../model_runner/omni_connector_runtime.py | 134 +- .../chunk_transfer_adapter.py | 4 + vllm_omni/engine/__init__.py | 12 + vllm_omni/engine/arg_utils.py | 6 + vllm_omni/engine/orchestrator.py | 105 +- vllm_omni/engine/request_snapshot.py | 42 + vllm_omni/engine/stage_engine_core_client.py | 5 +- vllm_omni/engine/stage_engine_core_proc.py | 27 + vllm_omni/engine/stage_pool.py | 24 + .../entrypoints/openai/serving_speech.py | 16 +- .../models/common/qwen3_code_predictor.py | 150 +- .../model_executor/models/output_templates.py | 15 + .../models/qwen3_omni/qwen3_omni.py | 1 + .../qwen3_tts/configuration_qwen3_tts.py | 7 + .../qwen3_tts/cuda_graph_decoder_wrapper.py | 131 +- .../models/qwen3_tts/pipeline.py | 4 +- .../models/qwen3_tts/prompt_embeds_builder.py | 112 +- .../models/qwen3_tts/qwen3_tts_code2wav.py | 32 +- .../models/qwen3_tts/qwen3_tts_talker.py | 404 +++++- .../modeling_qwen3_tts_tokenizer_v2.py | 98 +- .../stage_input_processors/qwen3_tts.py | 135 +- vllm_omni/outputs/__init__.py | 4 + vllm_omni/outputs/output_processor.py | 53 +- vllm_omni/worker/base.py | 3 + vllm_omni/worker/gpu_ar_worker.py | 27 +- vllm_omni/worker/gpu_generation_worker.py | 31 +- vllm_omni/worker/gpu_model_runner.py | 10 + vllm_omni/worker/payload_span.py | 48 + vllm_omni/worker/sampling_utils.py | 58 +- vllm_omni/worker_v2/__init__.py | 0 vllm_omni/worker_v2/delivery.py | 218 +++ vllm_omni/worker_v2/model_states/__init__.py | 43 + .../model_states/intermediate_buffer.py | 265 ++++ .../model_states/omni_model_state.py | 1223 +++++++++++++++++ vllm_omni/worker_v2/native_output_worker.py | 99 ++ vllm_omni/worker_v2/omni_ar_model_runner.py | 818 +++++++++++ vllm_omni/worker_v2/omni_data_plane.py | 604 ++++++++ .../worker_v2/omni_generation_model_runner.py | 582 ++++++++ vllm_omni/worker_v2/omni_model_runner.py | 666 +++++++++ vllm_omni/worker_v2/output_snapshot.py | 78 ++ 103 files changed, 10232 insertions(+), 661 deletions(-) create mode 100644 tests/core/sched/test_omni_ar_scheduler_output_policy.py create mode 100644 tests/core/sched/test_omni_scheduler_ready_inbox.py create mode 100644 tests/engine/test_output_processor_mrv2_text.py create mode 100644 tests/engine/test_request_snapshot.py create mode 100644 tests/engine/test_terminal_stage_metrics_order.py create mode 100644 tests/entrypoints/openai_api/test_ref_audio_cache.py create mode 100644 tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py create mode 100644 tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py create mode 100644 tests/worker/test_omni_v2_profile_dispatch.py create mode 100644 tests/worker/test_tts_sampling_seed.py create mode 100644 tests/worker_v2/__init__.py create mode 100644 tests/worker_v2/test_intermediate_buffer.py create mode 100644 tests/worker_v2/test_intermediate_buffer_cuda_residency.py create mode 100644 tests/worker_v2/test_omni_ar_model_runner.py create mode 100644 tests/worker_v2/test_omni_data_plane.py create mode 100644 tests/worker_v2/test_omni_generation_model_runner.py create mode 100644 tests/worker_v2/test_omni_gpu_model_runner.py create mode 100644 tests/worker_v2/test_omni_model_state.py create mode 100644 vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml create mode 100644 vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml create mode 100644 vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml create mode 100644 vllm_omni/deploy/qwen3_tts_mrv2.yaml create mode 100644 vllm_omni/engine/request_snapshot.py create mode 100644 vllm_omni/worker/payload_span.py create mode 100644 vllm_omni/worker_v2/__init__.py create mode 100644 vllm_omni/worker_v2/delivery.py create mode 100644 vllm_omni/worker_v2/model_states/__init__.py create mode 100644 vllm_omni/worker_v2/model_states/intermediate_buffer.py create mode 100644 vllm_omni/worker_v2/model_states/omni_model_state.py create mode 100644 vllm_omni/worker_v2/native_output_worker.py create mode 100644 vllm_omni/worker_v2/omni_ar_model_runner.py create mode 100644 vllm_omni/worker_v2/omni_data_plane.py create mode 100644 vllm_omni/worker_v2/omni_generation_model_runner.py create mode 100644 vllm_omni/worker_v2/omni_model_runner.py create mode 100644 vllm_omni/worker_v2/output_snapshot.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 067af9fa6..5dfbf8da4 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -23,6 +23,16 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/entrypoints/openai/serving_speech.py - vllm_omni/entrypoints/openai/tts_adapters/ - vllm_omni/engine/stage_pool.py + - vllm_omni/worker_v2/ + - vllm_omni/core/sched/ + - vllm_omni/distributed/omni_connectors/ + - vllm_omni/config/stage_config.py + - vllm_omni/config/omni_config.py + - vllm_omni/engine/orchestrator.py + - vllm_omni/engine/stage_engine_core_proc.py + - vllm_omni/deploy/qwen3_tts_mrv2.yaml + - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml + - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml x-minicpmo-4-5: &minicpmo_4_5 - vllm_omni/engine/async_omni_engine.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 53b9bc4d0..04e5d07b0 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -102,7 +102,7 @@ steps: commands: - | timeout 15m bash -c " - pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model' + pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model' " mirror_hardwares: l4_1 diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 199241cf6..2540b19d9 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -290,3 +290,90 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/depl !!! important We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers. + +## Qwen3-TTS with Model Runner V2 + +Qwen3-TTS can opt into the native CUDA Model Runner V2 pipeline on vLLM +0.29.0. Select one of these deployment profiles: + +| Profile | Runner | Code2Wav graph batches | Intended use | +| --- | --- | --- | --- | +| `qwen3_tts.yaml` | V1 | Existing defaults | Existing deployment / regression control | +| `qwen3_tts_mrv2.yaml` | V2 | B1 | Native runner validation | +| `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 | + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_mrv2.yaml +``` + +The V2 profiles bound Talker prefill to 512 tokens per step and select the +Talker AR runner and the Code2Wav generation +runner together. Native inter-stage delivery carries codec payloads directly; +request-owned snapshots preserve buffers through asynchronous completion and +CUDA graph reuse. Terminal completion waits for upstream stage metrics before +releasing request state. Platform sections retain V1 on NPU, XPU, ROCm and MUSA; +this change does not qualify MRV2 on those backends or enable other model families. +Selecting V2 for a stage without `supports_native_mrv2_data_plane` emits a warning +when pipeline and deployment settings are merged. That stage retains the legacy +transport path; the warning does not establish support for that combination. + +MRV2 model hooks use capability declarations rather than architecture or stage +names. Omni lifecycle flags select the model state, and `_returns_tuple` declares +the capture output contract. The optional batched predictor hook is `mtp`, with +an explicit `mtp_output_key` (a string or two-part payload key), optional +`mtp_validity_key`, and `mtp_graph_safe`/`mtp_disable_graph` capture controls. +Its inputs are token IDs, embeddings, previous hidden states and per-row +conditioning; it returns updated embeddings and prediction codes. Models may +supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local explicit seeds, and declare +`mtp_accepts_per_row_generators`, `mtp_accepts_req_infos`, or `mtp_sample_uniforms` +(with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains +its existing `talker_mtp` entry point for V1. + +### Included performance work + +- Request snapshots have a fast path for immutable scalar leaves, including + waveform lists. Mutable containers, aliases and cycles keep deep-copy + semantics. This is not the historical shallow `dict(prompt)` experiment. +- The API reference cache holds owned float32 arrays with a default capacity + of 1024 entries and a 512 MiB waveform-payload budget. Cache hits return + independent lists. Qwen3-TTS model artifact caches also default to 1024 + entries. These caches have different owners and eviction policies; equal + entry limits do not make them a single coherent cache. +- Talker state stays on GPU where its lifecycle allows it. BOS/EOS projections + are cached in their projection dtype and invalidated on weight loading. +- Code2Wav packs CPU codec inputs before device transfer. Optional B2/B4 graph + buckets batch compatible requests without sharing their per-request state. +- Native output materialization runs in a bounded worker and drains before + closing the data plane. The existing control-only shortcut avoids launching + model work for a step that has only lifecycle events. + +MTP prefix re-prefill remains disabled in the high-concurrency V2 profile. +Adding more graph shapes increases compilation cost and has not established a +stable end-to-end gain for this extracted version. + +### Decoder batches and playback + +Changing `decode_cudagraph_batch_sizes` selects the captured batch buckets. +Keep `decode_batch_max_size` consistent with the intended maximum too; the +stateful graph path currently groups according to the captured buckets, while +the stateless path also uses the explicit maximum. + +B4 remains experimental. First-packet latency alone does not establish +uninterrupted playback. Validate inter-chunk arrival times, buffering, WER and +speaker similarity before adopting either batching preset for a production +workload. Floating-point decoder outputs can differ across batch sizes; this PR +does not claim bitwise or quality equivalence. + +### Optional MPS deployment + +NVIDIA MPS is an optional operator setting for colocated CUDA processes, not a +YAML option or a library default. This PR does not establish a throughput or +first-packet latency benefit from MPS. Measure the exact deployment with and +without MPS before enabling it. + +Use only assigned GPUs and an independent MPS pipe directory. A private MPS +server does not provide exclusive GPU ownership or MIG isolation. For a +single-GPU deployment, explicitly place both stages on that GPU; the supplied +high-concurrency profile places its two stages on different GPUs by default. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 2ea871890..fdba7e90a 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1429,6 +1429,9 @@ stages: assert deploy.async_chunk is True assert deploy.connectors is not None assert deploy.platforms is not None + connector = deploy.connectors["connector_of_shared_memory"] + assert "async_chunk_batch_min_size" not in connector["extra"] + assert "async_chunk_batch_max_wait_ms" not in connector["extra"] @pytest.mark.parametrize( ("hf_config", "model"), @@ -2704,6 +2707,56 @@ class TestPlatformOverrides: assert deploy.stages[0].engine_extras["block_size"] == 128 + @pytest.mark.parametrize("platform", ["cuda", "npu"]) + def test_recommended_native_runner_platform_defaults(self, platform): + filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts" + deploy = load_deploy_config(Path(get_deploy_config_path(filename))) + deploy = _apply_platform_overrides(deploy, platform=platform) + expect_v2 = platform == "cuda" + assert deploy.model_runner == ("v2" if expect_v2 else "v1") + + # The selection must reach the live worker-dispatch consumer, not just + # the transport-level DeployConfig field. + pipeline = resolve_pipeline_config(pipeline_key) + stages = merge_pipeline_deploy(pipeline, deploy) + assert [stage.yaml_engine_args["use_v2_model_runner"] for stage in stages] == [expect_v2] * len(stages) + + if expect_v2 and pipeline.stages and deploy.async_chunk: + # v2 only engages the native plane on stages declaring support. + assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages) + + @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)]) + def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native): + from unittest.mock import Mock + + warning = Mock() + monkeypatch.setattr("vllm_omni.config.stage_config.logger.warning", warning) + pipeline = PipelineConfig( + model_type="custom_pipeline", + stages=[ + StagePipelineConfig( + stage_id=0, model_stage="custom_ar", final_output=True, supports_native_mrv2_data_plane=native + ) + ], + ) + merge_pipeline_deploy(pipeline, DeployConfig(model_runner=runner)) + assert warning.call_count == int(runner == "v2" and not native) + if warning.called: + assert "legacy transport path" in warning.call_args.args[0] + + def test_runner_selection_rejects_engine_extras_override(self): + pipeline = PipelineConfig( + model_type="runner_selection_reserved", + stages=SINGLE_STAGE_PIPE_CFG, + ) + for reserved in ("use_v2_model_runner", "supports_native_mrv2_data_plane"): + deploy = DeployConfig( + model_runner="v2", + stages=[StageDeployConfig(stage_id=0, engine_extras={reserved: False})], + ) + with pytest.raises(ValueError, match=f"{reserved!r} must not be set"): + merge_pipeline_deploy(pipeline, deploy) + def test_npu_overrides(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index cdb303316..331717ba6 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -50,8 +50,10 @@ from vllm_omni.config.stage_config import ( StageDeployConfig, StageExecutionType, StagePipelineConfig, + _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy, + resolve_deploy_yaml, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.engine.stage_engine_startup import _serialize_stage_config @@ -783,6 +785,67 @@ def test_joyai_code2wav_waits_for_full_payload(): assert code2wav.model_config.requires_full_payload_input is True +@pytest.mark.parametrize("model_runner", ["v1", "v2"]) +def test_deploy_model_runner_selection_propagates_to_every_stage(tmp_path: Path, model_runner: str): + deploy_path = tmp_path / "qwen3_tts_runner.yaml" + deploy_path.write_text( + f"""\ +pipeline: qwen3_tts +model_runner: {model_runner} +async_chunk: true +stages: + - stage_id: 0 + - stage_id: 1 +""" + ) + + deploy = load_deploy_config(deploy_path) + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + legacy_stages = merge_pipeline_deploy(pipeline, deploy) + structured = _from_pipeline_key("qwen3_tts", deploy_config_path=str(deploy_path)) + expect_v2 = model_runner == "v2" + + assert deploy.model_runner == model_runner + assert all(stage.yaml_engine_args["use_v2_model_runner"] is expect_v2 for stage in legacy_stages) + assert all(stage.model_config.use_v2_model_runner is expect_v2 for stage in structured.stage_configs) + + +@pytest.mark.parametrize("platform", ["npu", "xpu"]) +def test_mrv2_fails_fast_on_platforms_without_native_workers(platform: str): + with pytest.raises(NotImplementedError, match="Model Runner V2"): + _apply_platform_overrides(DeployConfig(model_runner="v2"), platform=platform) + + +@pytest.mark.parametrize( + ("default_name", "mrv2_name"), + [ + ("qwen3_tts.yaml", "qwen3_tts_mrv2.yaml"), + ( + "qwen3_tts_high_concurrency.yaml", + "qwen3_tts_high_concurrency_mrv2.yaml", + ), + ], +) +def test_qwen3_mrv2_profiles_are_explicit_opt_in(default_name: str, mrv2_name: str): + assert load_deploy_config(_DEPLOY_DIR / default_name).model_runner == "v1" + assert load_deploy_config(_DEPLOY_DIR / mrv2_name).model_runner == "v2" + + +def test_qwen3_tts_mrv2_retunes_do_not_change_default_mrv1_profile(): + default = resolve_deploy_yaml(_DEPLOY_DIR / "qwen3_tts_high_concurrency.yaml") + mrv2 = resolve_deploy_yaml(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2.yaml") + default_extra = default["connectors"]["connector_of_shared_memory"]["extra"] + mrv2_extra = mrv2["connectors"]["connector_of_shared_memory"]["extra"] + + # B2 batch stays consistent with the graph buckets; V1 profile untouched. + assert default_extra["decode_batch_max_size"] == 1 + assert mrv2_extra["decode_batch_max_size"] == 2 + assert mrv2_extra["decode_cudagraph_batch_sizes"] == [1, 2] + assert mrv2_extra["code_predictor_prefix_graphs"] is False + assert mrv2["stages"][1]["enforce_eager"] is False + assert default["stages"][1]["enforce_eager"] is True + + def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config") assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config") @@ -879,9 +942,11 @@ def test_sub_config_fields_match_structured_scopes(): "limit_mm_per_prompt", "interleave_mm_strings", "media_io_kwargs", + "final_output", "active_stream_window", "session_mode", "duplex_max_sessions", + "use_v2_model_runner", "enable_sleep_mode", "default_sampling_params", "subtalker_sampling_params", @@ -890,6 +955,7 @@ def test_sub_config_fields_match_structured_scopes(): "custom_voice_dir", "task_type", "codec_frame_rate_hz", + "supports_native_mrv2_data_plane", "enforce_eager", "max_cudagraph_capture_size", "enable_flashinfer_autotune", diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 096471e5b..377d6265f 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/tests/core/sched/test_generation_scheduler_restore.py @@ -1,14 +1,22 @@ -"""Test that OmniGenerationScheduler restores chunk-waiting requests -even when the OmniNewRequestData rewrapping fails. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -Regression test: if process_pending_chunks() moves requests into -internal deques but restore_queues() is not called due to an exception, -those requests are permanently orphaned. -""" +"""Chunk-lifecycle coverage for OmniGenerationScheduler: restore on error, +native terminal state, in-flight no-resubmission, slot release on finish, +and single scheduling of a terminal empty-prompt chunk.""" from collections import deque +from types import SimpleNamespace import pytest +import torch +from vllm import SamplingParams +from vllm.v1.core.sched.interface import PauseState +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import Request, RequestStatus + +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, _has_async_chunk_payload_to_run +from vllm_omni.engine.serialization import serialize_additional_information pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -20,16 +28,20 @@ class FakeAdapter: self.waiting_for_chunk_waiting_requests = deque() self.waiting_for_chunk_running_requests = deque() self.restore_called = False + self.done_request_ids = set() - def process_pending_chunks(self, waiting, running): - """Simulate moving requests out of the scheduler queues.""" - # Move one request from running into internal deque + def process_pending_chunks(self, waiting, running, scheduler_requests=None): if running: req = running.pop() self.waiting_for_chunk_running_requests.append(req) + def is_done_receiving_chunks(self, request_id): + return request_id in self.done_request_ids + + def collect_failed_send_request_ids(self): + return {} + def restore_queues(self, waiting, running, scheduler_requests=None): - """Put requests back.""" self.restore_called = True running.extend(self.waiting_for_chunk_running_requests) self.waiting_for_chunk_running_requests = deque() @@ -38,6 +50,131 @@ class FakeAdapter: pass +def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False): + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.max_num_scheduled_tokens = 8 + scheduler.max_num_running_reqs = 1 + scheduler._pause_state = PauseState.UNPAUSED + scheduler.running = [] + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.waiting.add_request(waiting_request) + scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.requests = {waiting_request.request_id: waiting_request} + scheduler.policy = SchedulingPolicy.FCFS + scheduler.chunk_transfer_adapter = FakeAdapter() + scheduler.input_coordinator = None + scheduler.log_stats = False + scheduler.scheduler_config = SimpleNamespace(enable_chunked_prefill=True) + scheduler.num_lookahead_tokens = 0 + scheduler.kv_cache_manager = SimpleNamespace( + new_step_starts=lambda: None, + allocate_slots=lambda *args, **kwargs: SimpleNamespace(get_block_ids=lambda: ([1],)), + get_num_common_prefix_blocks=lambda request_id: [0], + take_new_block_ids=lambda: None, + ) + scheduler.kv_cache_config = SimpleNamespace(kv_cache_groups=[object()]) + scheduler.use_v2_model_runner = use_v2_model_runner + scheduler._retains_state_across_chunks = False + scheduler.needs_kv_cache_zeroing = False + scheduler.finished_req_ids = set() + scheduler.encoder_cache_manager = SimpleNamespace(get_freed_mm_hashes=lambda: []) + scheduler.connector = None + scheduler.ec_connector = None + scheduler.prev_step_scheduled_req_ids = set() + scheduler._pending_finish_reqs = [] + scheduler._consume_pending_connector_output = lambda model_mode: None + scheduler._process_pending_input_timeouts = lambda: None + scheduler._make_cached_request_data = lambda **kwargs: SimpleNamespace( + req_ids=[], + resumed_req_ids=[], + new_token_ids=[], + all_token_ids=[], + new_block_ids=[], + num_computed_tokens=[], + num_output_tokens=[], + ) + scheduler._update_after_schedule = lambda output: None + scheduler._wrap_omni_scheduler_output = lambda output: output + return scheduler + + +def _chunk_request(request_id, **kwargs): + defaults = dict( + request_id=request_id, + prompt_token_ids=[1], + num_computed_tokens=0, + num_in_flight_tokens=0, + status=None, + sampling_params=None, + pooling_params=None, + mm_features=None, + lora_request=None, + prompt_is_token_ids=True, + additional_information=None, + external_req_id=request_id, + prefill_stats=None, + record_event=lambda *args, **kwargs: None, + ) + defaults.update(kwargs) + return SimpleNamespace(**defaults) + + +def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None: + # Native plane: a completed payload without a new chunk is never executed + # again; requeued requests keep WAITING_FOR_CHUNK and their token counts. + scheduler = _make_generation_scheduler(_chunk_request("unused")) + monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue) + scheduler._native_data_plane = True + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None) + scheduler.running = [] + scheduler.waiting = create_request_queue(scheduler.policy) + scheduler.skipped_waiting = create_request_queue(scheduler.policy) + + completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None) + completed.status = RequestStatus.RUNNING + completed.num_computed_tokens = 2 + scheduler.running = [completed] + scheduler.requests = {"completed": completed} + scheduler.input_coordinator.finished_requests.add("completed") # terminal payload, no new chunk + + assert scheduler._is_done_receiving_chunks("completed") + assert not scheduler._is_done_receiving_chunks("unknown") + output = scheduler.schedule() + assert not output.num_scheduled_tokens + assert scheduler._pending_finish_reqs == [completed] + + completed.status = RequestStatus.WAITING_FOR_CHUNK + scheduler._pending_finish_reqs = [] + scheduler._requeue_completed_native_chunks() + assert completed.status == RequestStatus.WAITING_FOR_CHUNK + assert completed.num_computed_tokens == 2 + assert not scheduler.running + + +def test_generation_scheduler_schedules_terminal_empty_prompt_chunk_once(monkeypatch: pytest.MonkeyPatch) -> None: + waiting = _chunk_request( + "terminal", + prompt_token_ids=[], + num_prompt_tokens=0, + additional_information=serialize_additional_information( + {"codes": {"audio": torch.tensor([[1, 2, 3]], dtype=torch.long)}} + ), + ) + scheduler = _make_generation_scheduler(waiting) + scheduler.chunk_transfer_adapter.done_request_ids.add("terminal") + scheduler.requests = {"terminal": waiting} + assert _has_async_chunk_payload_to_run(waiting) + # A terminal codec request with an empty prompt is still scheduled exactly once. + monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue) + + output = OmniGenerationScheduler.schedule(scheduler) + + assert output.num_scheduled_tokens == {"terminal": 1} + assert output.scheduled_new_reqs[0].req_id == "terminal" + assert scheduler._pending_finish_reqs == [] + + class TestRestoreQueuesOnError: """Verify that restore_queues is called even when rewrapping raises.""" diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py index d9fbc452b..f8b9d9b4c 100644 --- a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py +++ b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for ``OmniARScheduler._drop_aborted_queued_requests``. ``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before @@ -68,6 +71,7 @@ def test_schedule_sweeps_skipped_waiting_before_upstream_selection( aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED) scheduler = _make_scheduler(policy, skipped_waiting=[aborted]) scheduler.chunk_transfer_adapter = None + scheduler.max_num_running_reqs = 8 scheduler.input_coordinator = None scheduler.num_waiting_for_streaming_input = 0 scheduler._process_pending_omni_inputs = lambda model_mode: None diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 810c506c7..13594b3bd 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for the AR sampled-token logprob contract.""" from __future__ import annotations @@ -93,6 +96,7 @@ class _Request: self.num_computed_tokens = 0 self.num_in_flight_tokens = 0 self.num_output_placeholders = 0 + self.output_token_ids: list[int] = [] # vLLM 0.27 (a0c092ee72): Request gained num_stale_output_tokens to # track in-flight outputs discarded at preemption/streaming-stop. self.num_stale_output_tokens = 0 @@ -145,6 +149,14 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: kv_cache_manager=SimpleNamespace(take_events=lambda: None), kv_event_publisher=SimpleNamespace(publish=lambda _events: None), recompute_kv_load_failures=False, + _native_data_plane=False, + vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + use_v2_model_runner=False, + async_chunk=False, + final_output=True, + ) + ), ) for name in _MIXIN_UPDATE_HELPERS: setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler)) @@ -167,10 +179,15 @@ def _bind_request_lifecycle( scheduler.requests.pop(request.request_id, None) scheduler.finished_req_ids.add(request.request_id) scheduler.finished_req_ids_dict[request.client_index].add(request.request_id) - # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params) + # Current vLLM contract: (kv_xfer_params, ec_xfer_params) return None, None - scheduler._update_request_with_output = update_request + def update_with_history(request, token_ids): + accepted, stopped = update_request(request, token_ids) + request.output_token_ids.extend(accepted) + return accepted, stopped + + scheduler._update_request_with_output = update_with_history scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False scheduler._handle_stopped_request = handle_stopped or (lambda _request: True) scheduler._free_request = free_request @@ -219,6 +236,7 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None: (output,) = outputs[0].outputs assert output.new_token_ids == [7] + assert output.num_generation_tokens == 1 assert output.finish_reason is FinishReason.STOP token_rows = np.asarray(output.new_logprobs.logprob_token_ids) value_rows = np.asarray(output.new_logprobs.logprobs) diff --git a/tests/core/sched/test_omni_ar_scheduler_output_policy.py b/tests/core/sched/test_omni_ar_scheduler_output_policy.py new file mode 100644 index 000000000..89c08bf38 --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_output_policy.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.core.sched.omni_ar_scheduler import _should_emit_engine_output + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("async_chunk", "final_output", "use_v2", "stopped", "has_control", "expected"), + [ + (True, False, True, False, False, False), # async+v2 suppresses data + (True, False, False, False, False, True), # v1 continues to emit + (True, False, True, True, False, True), # stopped always emits + (True, False, True, False, True, True), # control-only still emits + ], +) +def test_async_chunk_intermediate_stage_emits_only_control_outputs( + async_chunk: bool, + final_output: bool, + use_v2: bool, + stopped: bool, + has_control: bool, + expected: bool, +) -> None: + model_config = SimpleNamespace( + async_chunk=async_chunk, + final_output=final_output, + use_v2_model_runner=use_v2, + ) + + assert ( + _should_emit_engine_output( + model_config, + stopped=stopped, + has_control=has_control, + ) + is expected + ) diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 154644fb2..83c7038f2 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -17,6 +17,7 @@ import pytest import vllm_omni # noqa: F401 - import for side effects (patch vLLM) from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.engine import FinishReason from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue @@ -1261,3 +1262,47 @@ def test_context_overflow_emits_an_error_output_with_the_reason() -> None: sched._emit_streaming_context_overflow_outputs(outputs) assert len(outputs[2]) == 1 + + +@pytest.mark.parametrize("native", [True, False]) +def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, native) -> None: + sched = _make_scheduler(stage_id=1) + parked = SimpleNamespace(request_id="parked") + sched.requests = {"parked": parked} + sched.waiting = [] + sched.running = [] + sched._native_data_plane = native + sched.use_v2_model_runner = native + sched.max_num_running_reqs = 8 + sched.input_coordinator = ( + SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None + ) + sched.chunk_transfer_adapter = ( + None + if native + else SimpleNamespace( + waiting_for_chunk_running_requests=[parked], + _held_non_active=[], + process_pending_chunks=lambda *_a, **_kw: None, + collect_failed_send_request_ids=lambda: {}, + restore_queues=lambda *_a, **_kw: None, + postprocess_scheduler_output=lambda *_a, **_kw: None, + ) + ) + sched._consume_pending_connector_output = lambda model_mode: None + sched._process_pending_input_timeouts = lambda: None + sched._should_defer_waiting_admission = lambda: False + sched.get_finished_requests_needing_kv_transfer = lambda: {} + sched._wrap_omni_scheduler_output = lambda output, **_kwargs: output + observed_limits: list[int] = [] + + def fake_schedule(self, _throttle_prefills=False): + observed_limits.append(self.max_num_running_reqs) + return SimpleNamespace(scheduled_new_reqs=[]) + + monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule) + + sched.schedule() + + assert observed_limits == [7 if native else 8] + assert sched.max_num_running_reqs == 8 diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index 0e0ed04ad..de68f7d5a 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from types import SimpleNamespace @@ -141,6 +144,7 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None: scheduler = _make_scheduler() + scheduler.max_num_running_reqs = 8 scheduler.running = [_MockRequest("decode")] original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING) deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING) diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index 36b084f7c..9088c6bc2 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -123,6 +123,9 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request sched = MagicMock() sched.max_num_scheduled_tokens = 8 + sched.max_num_running_reqs = 8 + sched._native_data_plane = False + sched._retains_state_across_chunks = False sched._pause_state = PauseState.UNPAUSED sched.requests = {request.request_id: request} sched.running = [request] diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index bb9fd8859..ba4b2454c 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py @@ -1,5 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections import defaultdict from types import SimpleNamespace +from typing import Any import pytest from vllm.v1.engine import FinishReason @@ -87,7 +91,7 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): - calls = [] + calls: list[tuple[Any, ...]] = [] scheduler = _Scheduler() scheduler.waiting = ["waiting"] scheduler.running = ["running"] @@ -115,7 +119,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): collect_failed_send_request_ids=_collect_failed_sends, ) scheduler.input_coordinator = SimpleNamespace( - restore_queues=lambda waiting: calls.append(("restore-full", waiting)) + restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running)) ) scheduler._process_pending_omni_inputs("ar") @@ -130,7 +134,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S), ("failed-sends",), ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests), - ("restore-full", scheduler.waiting), + ("restore-full", scheduler.waiting, scheduler.running), ] diff --git a/tests/core/sched/test_omni_scheduler_ready_inbox.py b/tests/core/sched/test_omni_scheduler_ready_inbox.py new file mode 100644 index 000000000..483255f21 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_ready_inbox.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Drain ready events once, merging live requests and filtering cancellations.""" + +from types import SimpleNamespace + +import pytest + +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin +from vllm_omni.outputs import OmniConnectorOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker): + scheduler = OmniSchedulerMixin() + scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], [] + coordinator = SimpleNamespace( + _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock() + ) + scheduler.input_coordinator = coordinator + scheduler._init_omni_connector_output_inbox() + scheduler.enqueue_omni_connector_output(OmniConnectorOutput(chunk_ready_req_ids={"r"})) + scheduler.enqueue_omni_connector_output( + OmniConnectorOutput( + chunk_ready_req_ids={"r", "aborted"}, + chunk_finished_req_ids={"r", "aborted"}, + request_metadata={"r": {"decode_token_end": 2}, "aborted": {"decode_token_end": 99}}, + ) + ) + scheduler._consume_pending_connector_output(model_mode="ar") + coordinator.update_request_metadata.assert_called_once_with( + scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar" + ) + coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"}) + scheduler._consume_pending_connector_output(model_mode="ar") + assert coordinator.update_request_metadata.call_count == 1 + coordinator.process_pending_chunks.assert_called_with([], [], set(), set()) diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py index 6f44aadac..bc392abca 100644 --- a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py +++ b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for the ``num_waiting_for_streaming_input`` resync placed in ``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests`` and in ``OmniARScheduler.schedule``. @@ -128,6 +131,7 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk scheduler.skipped_waiting = [] scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite scheduler.policy = "fcfs" + scheduler.max_num_running_reqs = 8 seen: dict[str, int] = {} diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index f37aac29f..8a62cbbfd 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniSchedulingCoordinator. These tests use mock request objects and mock queues. They do not require @@ -20,7 +20,9 @@ import torch import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, + uses_native_mrv2_data_plane, ) +from vllm_omni.core.sched.output import OmniChunkRecvHandle pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -56,6 +58,7 @@ def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace: prompt_token_ids=[], num_prompt_tokens=0, num_computed_tokens=0, + num_output_placeholders=0, _all_token_ids=[], _output_token_ids=[], payload_sender_info=None, @@ -83,7 +86,7 @@ class MockQueue: def prepend_requests(self, requests): self._items = list(requests) + self._items - def remove(self, request): + def remove_request(self, request): self._items.remove(request) def remove_requests(self, requests): @@ -96,7 +99,47 @@ class MockQueue: # ------------------------------------------------------------------ # -class TestCoordinatorUpdateRequestMetadata(unittest.TestCase): +class TestNativeMRV2DataPlaneSelection(unittest.TestCase): + def test_native_plane_requires_v2_and_async_chunk_capability(self): + native = SimpleNamespace(async_chunk=True, supports_native_mrv2_data_plane=True) + + self.assertTrue(uses_native_mrv2_data_plane(native, use_v2_model_runner=True)) + self.assertFalse(uses_native_mrv2_data_plane(native, use_v2_model_runner=False)) + self.assertFalse(uses_native_mrv2_data_plane(SimpleNamespace(async_chunk=False), use_v2_model_runner=True)) + + +def test_chunk_registration_ready_and_terminal_lifecycle(): + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) + req = _make_request("internal", status=RequestStatus.WAITING) + req.external_req_id = "external" + waiting = MockQueue([req]) + coord.process_pending_chunks(waiting, [], set(), set()) + assert req.status == RequestStatus.WAITING_FOR_CHUNK + [handle] = coord.pending_chunk_registrations + assert isinstance(handle, OmniChunkRecvHandle) + assert (handle.request_id, handle.external_req_id) == ("internal", "external") + coord.restore_queues(waiting, []) + coord.process_pending_chunks(waiting, [], {"internal"}, set()) + assert req.status == RequestStatus.WAITING + assert "internal" in coord.requests_with_ready_chunks + req.status = RequestStatus.WAITING_FOR_CHUNK + coord.process_pending_chunks(waiting, [], set(), {"internal"}) + assert "internal" in coord.finished_requests + + +def test_chunk_waiting_removes_request_from_running_list(): + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) + req = _make_request("running", status=RequestStatus.RUNNING) + running = [req] + + coord.process_pending_chunks(MockQueue(), running, set(), set()) + + assert running == [] + assert req.status == RequestStatus.WAITING_FOR_CHUNK + assert list(coord._waiting_for_chunk_running) == [req] + + +class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase): """Test update_request_metadata applies scheduling metadata to requests.""" def test_ar_mode_no_longer_sets_additional_information(self): @@ -382,17 +425,18 @@ class TestTimeoutDetection(unittest.TestCase): self.assertEqual(result, {"r1"}) self.assertEqual(len(coord._waiting_for_input), 0) - def test_free_finished_request_clears_waiting_since(self): - """free_finished_request clears coordinator lifecycle markers.""" - coord = OmniSchedulingCoordinator(stage_id=1) - coord._waiting_since["r1"] = 0.0 - coord._full_payload_input_received.add("r1") + def test_free_finished_request_clears_all_lifecycle_state(self): + """free_finished_request makes stale connector events harmless.""" + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + r1, r2 = _make_request("r1"), _make_request("r2") coord.finished_requests.add("r1") + coord.requests_with_ready_chunks.add("r1") + coord._waiting_for_chunk_running.extend([r1, r2]) + coord.pending_chunk_registrations = [OmniChunkRecvHandle(request_id="r1"), OmniChunkRecvHandle(request_id="r2")] + coord.free_finished_request("r1") - self.assertNotIn("r1", coord._waiting_since) - self.assertNotIn("r1", coord._full_payload_input_received) + self.assertNotIn("r1", coord.finished_requests) - - -if __name__ == "__main__": - unittest.main() + self.assertNotIn("r1", coord.requests_with_ready_chunks) + self.assertEqual([r.request_id for r in coord._waiting_for_chunk_running], ["r2"]) + self.assertEqual([h.request_id for h in coord.pending_chunk_registrations], ["r2"]) diff --git a/tests/e2e/offline_inference/test_qwen3_tts_base.py b/tests/e2e/offline_inference/test_qwen3_tts_base.py index f800a1163..6515b2e58 100644 --- a/tests/e2e/offline_inference/test_qwen3_tts_base.py +++ b/tests/e2e/offline_inference/test_qwen3_tts_base.py @@ -13,6 +13,7 @@ import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest +from vllm.platforms import current_platform from tests.helpers.mark import hardware_test from tests.helpers.media import get_asset_path @@ -57,7 +58,12 @@ tts_server_params = [ pytest.param( (MODEL, get_cuda_graph_config()), id="no_cuda_graph", - ) + ), + pytest.param( + (MODEL, get_deploy_config_path("qwen3_tts_mrv2.yaml")), + id="async_chunk_mrv2", + marks=pytest.mark.skipif(not current_platform.is_cuda(), reason="MRV2 validation is CUDA-only"), + ), ] diff --git a/tests/e2e/online_serving/test_qwen3_tts_base.py b/tests/e2e/online_serving/test_qwen3_tts_base.py index dd25520f9..338655925 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base.py @@ -12,6 +12,7 @@ import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest +from vllm.platforms import current_platform from tests.helpers.mark import hardware_test from tests.helpers.media import get_asset_path @@ -50,7 +51,16 @@ tts_server_params = [ server_args=["--trust-remote-code"], ), id="async_chunk", - ) + ), + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("qwen3_tts_mrv2.yaml"), + server_args=["--trust-remote-code"], + ), + id="async_chunk_mrv2", + marks=pytest.mark.skipif(not current_platform.is_cuda(), reason="MRV2 validation is CUDA-only"), + ), ] diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 775eeb86b..4bbca2515 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -311,7 +311,7 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace: return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None) -def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs: +def _terminal_engine_core_outputs(request_id: str, timestamp: float = 1.0) -> EngineCoreOutputs: return EngineCoreOutputs( outputs=[ EngineCoreOutput( @@ -320,7 +320,7 @@ def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs: finish_reason=FinishReason.STOP, ) ], - timestamp=1.0, + timestamp=timestamp, finished_requests={request_id}, ) @@ -520,6 +520,7 @@ async def _enqueue_add_request( original_prompt, sampling_params_list, final_stage_id: int, + final_output_stage_ids: list[int] | None = None, ) -> None: orchestrator_fixture.request_sync_q.put_nowait( StageSubmissionMessage( @@ -530,6 +531,7 @@ async def _enqueue_add_request( output_prompt_text=None, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, preprocess_ms=0.0, request_timestamp=time.time(), enqueue_ts=time.perf_counter(), @@ -789,6 +791,10 @@ async def test_run_async_chunk(orchestrator_factory) -> None: stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0)) + await _wait_for( + lambda: orchestrator_fixture.orchestrator.request_states["req-async"].pending_final_output is not None + ) + stage0.push_engine_core_outputs(_engine_core_outputs("stage0-final", 3.1)) output_msg = await _get_output_message(orchestrator_fixture) assert output_msg.request_id == "req-async" diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 419be9305..9fbecc1c1 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -259,6 +259,17 @@ def test_native_text_tpot_only_accepts_positive_finished_metric( assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms +def test_suppressed_native_tokens_use_authoritative_count_even_without_stats(monkeypatch): + monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None) + processor = object.__new__(MultimodalOutputProcessor) + processor._native_text_metrics_by_request = {} + state = _make_state(RequestOutputKind.FINAL_ONLY) + processor._update_stats_from_output( + state, SimpleNamespace(new_token_ids=[2150], num_generation_tokens=2048), None, None + ) + assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 2048 + + def test_delta_drains_output_modality_per_step(): """DELTA drains the mm_type key (output modality) but preserves hidden-state keys.""" s = _make_state(RequestOutputKind.DELTA) @@ -616,6 +627,7 @@ def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypa finish_reason=None, stop_reason=None, kv_transfer_params=None, + ec_transfer_params=None, routed_experts=None, num_cached_tokens=0, ) @@ -631,6 +643,7 @@ def _make_mm_only_output_processor(monkeypatch): processor = object.__new__(MultimodalOutputProcessor) processor.output_modality = OutputModality.AUDIO processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)} + processor.tracing_enabled = False monkeypatch.setattr( VLLMOutputProcessor, "process_outputs", @@ -660,6 +673,7 @@ def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool): finish_reason=FinishReason.STOP, stop_reason=None, kv_transfer_params=None, + ec_transfer_params=None, routed_experts=None, num_cached_tokens=0, is_segment_finished=is_segment_finished, diff --git a/tests/engine/test_output_processor_mrv2_text.py b/tests/engine/test_output_processor_mrv2_text.py new file mode 100644 index 000000000..b9dbecb1b --- /dev/null +++ b/tests/engine/test_output_processor_mrv2_text.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for MRv2 Omni AR text detach and generation payload outputs.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import FinishReason + +from vllm_omni.engine import OmniEngineCoreOutput +from vllm_omni.engine.output_modality import OutputModality +from vllm_omni.engine.output_processor import ( + MultimodalOutputProcessor, + OmniRequestState, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Detokenizer: + def __init__(self): + self.output_token_ids = [] + self.output_text = "" + self.update_calls = [] + + def update(self, token_ids, stop): + self.update_calls.append((list(token_ids), stop)) + self.output_token_ids.extend(token_ids) + self.output_text += "X" * len(token_ids) + return None + + def get_next_output_text(self, finished, delta): + return self.output_text + + def num_output_tokens(self): + return len(self.output_token_ids) + + +def _make_processor(output_modality, detokenizer=None): + state = OmniRequestState( + request_id="r", + external_req_id="r", + parent_req=None, + request_index=0, + lora_request=None, + output_kind=RequestOutputKind.FINAL_ONLY, + prompt="prompt" if detokenizer else None, + prompt_token_ids=[1] if detokenizer else [], + prompt_embeds=None, + logprobs_processor=MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None) + if detokenizer + else None, + detokenizer=detokenizer, + max_tokens_param=None, + arrival_time=0.0, + queue=None, + log_stats=False, + stream_interval=1, + ) + processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=output_modality) + processor.request_states["r"] = state + processor.external_req_ids["r"].append("r") + return processor, state + + +def _rehydrated_output(*, pooling_output=None, finish_reason=FinishReason.STOP): + """Build the post-StagePool shape without violating vLLM's wire schema.""" + return SimpleNamespace( + request_id="r", + new_token_ids=[42], + pooling_output=pooling_output, + finish_reason=finish_reason, + stop_reason=None, + kv_transfer_params=None, + ec_transfer_params=None, + routed_experts=None, + trace_headers=None, + prefill_stats=None, + spec_decode_metrics=None, + new_sampling_mask=None, + num_nans_in_logits=0, + multimodal_output=None, + is_segment_finished=False, + is_non_final_audio_chunk=False, + output_type=None, + num_generation_tokens=None, + finished=True, + ) + + +def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload(): + detokenizer = _Detokenizer() + processor, state = _make_processor(OutputModality.LATENT, detokenizer) + + # StagePool has already rehydrated the MRv2 carrier. Use a structural + # output here because upstream EngineCoreOutput strictly validates + # pooling_output as torch.Tensor | None on construction. + output = _rehydrated_output(pooling_output={"hidden": torch.ones(1, 4)}) + + processed = processor.process_outputs([output]) + + assert detokenizer.update_calls == [([42], True)] + completion = processed.request_outputs[0].outputs[0] + assert list(completion.token_ids) == [42] + assert completion.text == "X" + assert not state.mm_accumulated.is_empty + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_generation_stage_mm_delivered_only_with_terminal(streaming): + # Non-terminal chunks accumulate; the terminal output carries the audio. + # When the first output is already terminal it is delivered directly. + processor, _ = _make_processor(OutputModality.AUDIO) + audio = torch.ones(1, 320) + sr = torch.tensor(24000, dtype=torch.int32) + mm_output = OmniEngineCoreOutput( + request_id="r", + new_token_ids=[], + multimodal_output={"model_outputs": audio, "sr": sr}, + finish_reason=None if streaming else FinishReason.STOP, + ) + + processed = processor.process_outputs([mm_output]) + if streaming: + assert processed.request_outputs == [] + processed = processor.process_outputs( + [OmniEngineCoreOutput(request_id="r", new_token_ids=[], finish_reason=FinishReason.STOP)] + ) + + assert len(processed.request_outputs) == 1 + completion = processed.request_outputs[0].outputs[0] + assert completion.text == "" + assert torch.equal(completion.multimodal_output["audio"], audio) + assert completion.multimodal_output["sr"].item() == 24000 diff --git a/tests/engine/test_request_snapshot.py b/tests/engine/test_request_snapshot.py new file mode 100644 index 000000000..51d61ae3e --- /dev/null +++ b/tests/engine/test_request_snapshot.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy + +import pytest + +from vllm_omni.engine.request_snapshot import copy_request_snapshot + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_full_snapshot_preserves_fields_aliases_and_isolation(): + samples = [0.25, 0.5] + source = {"prompt_token_ids": [1, 2], "additional_information": {"samples": samples}, "alias": samples} + result = copy_request_snapshot(source) + assert result == copy.deepcopy(source) + assert result["alias"] is result["additional_information"]["samples"] + result["alias"].append(0.75) + result["prompt_token_ids"][0] = 9 + assert samples == [0.25, 0.5] + assert source["prompt_token_ids"] == [1, 2] diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 9ce707dbc..2e01e910b 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1393,7 +1393,7 @@ class TestConnectRemoteEngineCoresCoordinator: yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True) with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, @@ -1428,7 +1428,7 @@ class TestConnectRemoteEngineCoresCoordinator: yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True) with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, @@ -1454,7 +1454,7 @@ class TestConnectRemoteEngineCoresCoordinator: yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True) with connect_remote_diffusion_proc( omni_master_server=omni_master_server, @@ -1466,7 +1466,8 @@ class TestConnectRemoteEngineCoresCoordinator: omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) mock_wait.assert_called_once() - _, core_engines, parallel_config, *_ = mock_wait.call_args.args + _, core_engines, parallel_config, _, _, launch = mock_wait.call_args.args + assert launch.addresses is omni_master_server.get_zmq_addresses.return_value assert core_engines[0].local is False assert parallel_config.data_parallel_size_local == 0 @@ -1507,7 +1508,7 @@ class TestLaunchOmniCoreEngines: "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", return_value=local_engine_manager, ) - mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True) with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), @@ -1584,7 +1585,7 @@ class TestLaunchOmniCoreEngines: "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", return_value=mocker.Mock(), ) - mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True) with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 6d28280e0..c27292be4 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -8,6 +8,7 @@ import sys import types from dataclasses import fields, replace from pathlib import Path +from unittest.mock import patch import pytest from pydantic.fields import FieldInfo @@ -61,9 +62,14 @@ _DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy" def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]: backend_fields = fields(config_cls) backend_field_names = {backend_field.name for backend_field in backend_fields} - backend_config = config_cls( - **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names} - ) + # This is a schema/normalization comparison, not a Hub download test. + with ( + patch("vllm.engine.arg_utils.get_model_path", side_effect=lambda model, *_: model), + patch("vllm_omni.diffusion.data.get_model_path", side_effect=lambda model, *_: model), + ): + backend_config = config_cls( + **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names} + ) effective_values: dict[str, object] = {} for backend_field in backend_fields: value = getattr(backend_config, backend_field.name) @@ -78,6 +84,9 @@ _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusi _TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"}) _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( { + "final_output", + "use_v2_model_runner", + "supports_native_mrv2_data_plane", "active_stream_window", "codec_frame_rate_hz", "custom_voice_dir", diff --git a/tests/engine/test_terminal_stage_metrics_order.py b/tests/engine/test_terminal_stage_metrics_order.py new file mode 100644 index 000000000..a64fe8cc3 --- /dev/null +++ b/tests/engine/test_terminal_stage_metrics_order.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Terminal audio must not retire a request before upstream terminal metrics.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def build(): + o = Orchestrator.__new__(Orchestrator) + o.stage_pools = [ + SimpleNamespace( + final_output=False, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=0, num_tokens_out=42) + ), + SimpleNamespace( + final_output=True, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=1, num_tokens_out=0) + ), + ] + state = OrchestratorRequestState( + request_id="r", + final_stage_id=1, + final_output_stage_ids={1}, + stage_submit_ts={0: 0, 1: 0}, + sampling_params_list=[None, None], + ) + o.request_states = {"r": state} + o.output_async_queue = asyncio.Queue() + o.async_chunk = True + o._pd_pair = None + o._cfg_tracker = MagicMock() + o._cfg_tracker.is_companion.return_value = False + o._cfg_tracker.cleanup_parent.return_value = [] + o._is_duplex_session_request = lambda s: False + o._duplex_output_decision = lambda *a: None + o._stage_receives_async_chunks = lambda s: True + cleaned = [] + + async def cleanup(ids): + cleaned.extend(ids) + for r in ids: + o.request_states.pop(r, None) + + o._cleanup_request_ids = cleanup + return o, state, cleaned + + +@pytest.mark.asyncio +@pytest.mark.parametrize("order", [(0, 1), (1, 0)]) +async def test_terminal_metrics_precede_final_audio_for_either_arrival_order(order): + o, state, cleaned = build() + for stage in order: + await o._handle_processed_outputs(stage, 0, [SimpleNamespace(request_id="r", finished=True, error=None)]) + if stage == 1 and not state.finished_stage_ids.issuperset({0, 1}): + assert o.output_async_queue.empty() + assert "r" in o.request_states + assert cleaned == [] + metrics = o.output_async_queue.get_nowait() + final = o.output_async_queue.get_nowait() + assert type(metrics).__name__ == "StageMetricsMessage" + assert metrics.stage_id == 0 and metrics.metrics.num_tokens_out == 42 + assert type(final).__name__ == "OutputMessage" and final.finished + assert o.output_async_queue.empty() + assert cleaned == ["r"] + assert state.pending_final_output is None + + +@pytest.mark.asyncio +async def test_cancelled_pending_output_is_not_published_by_late_upstream(): + o, state, cleaned = build() + await o._handle_processed_outputs(1, 0, [SimpleNamespace(request_id="r", finished=True, error=None)]) + assert state.pending_final_output is not None + await o._cleanup_request_ids(["r"]) + await o._handle_processed_outputs(0, 0, [SimpleNamespace(request_id="r", finished=True, error=None)]) + assert o.output_async_queue.empty() + assert cleaned == ["r"] diff --git a/tests/entrypoints/openai_api/test_ref_audio_cache.py b/tests/entrypoints/openai_api/test_ref_audio_cache.py new file mode 100644 index 000000000..8026864b0 --- /dev/null +++ b/tests/entrypoints/openai_api/test_ref_audio_cache.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import hashlib +from collections import OrderedDict + +import numpy as np +import pytest + +from vllm_omni.entrypoints.openai.serving_speech import ( + _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES, + _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES, + OmniOpenAIServingSpeech, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def cache(): + server = object.__new__(OmniOpenAIServingSpeech) + server._diffusion_mode = False + server.model_config = None + server._ref_audio_resolve_cache = OrderedDict() + server._ref_audio_resolve_cache_bytes = 0 + server._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES + server._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES + server._ref_audio_model_artifact_ready = set() + server._request_ref_audio_artifact_keys = {} + return server + + +def _put(cache, key, samples, artifact=None): + cache._put_resolved_ref_audio(key, samples, 24000, artifact or key) + + +def test_compact_storage_and_default_capacity(cache): + assert cache._ref_audio_resolve_cache_max_entries == 1024 + assert cache._ref_audio_resolve_cache_max_bytes == 512 * 1024 * 1024 + _put(cache, "a", [0.0, 0.25, -0.5, 1.0]) + value = cache._ref_audio_resolve_cache["a"] + assert value[0].dtype == np.float32 + assert value[2] == cache._ref_audio_resolve_cache_bytes == 16 + + +def test_cached_float32_roundtrip_and_mutation_isolation(cache): + source = "data:audio/wav;base64,already-cached" + key = hashlib.sha1(source.encode()).hexdigest() + values = np.random.default_rng(42).standard_normal(48000).astype(np.float32).tolist() + expected = values.copy() + _put(cache, key, values) + values[0] = 99.0 + first, sample_rate, _ = asyncio.run(cache._resolve_ref_audio(source)) + assert first == expected + first[0] = 100.0 + second, _, _ = asyncio.run(cache._resolve_ref_audio(source)) + assert sample_rate == 24000 + assert second == expected + + +def test_entry_eviction_invalidates_artifact_readiness(cache): + cache._ref_audio_resolve_cache_max_entries = 2 + _put(cache, "a", [0.0]) + cache._ref_audio_model_artifact_ready.add(("a", False)) + _put(cache, "b", [0.0]) + _put(cache, "c", [0.0]) + assert list(cache._ref_audio_resolve_cache) == ["b", "c"] + assert ("a", False) not in cache._ref_audio_model_artifact_ready + + +def test_byte_budget_evicts_and_rejects_oversized_entries(cache): + cache._ref_audio_resolve_cache_max_bytes = 16 + _put(cache, "a", [0.0] * 3) + _put(cache, "b", [0.0] * 3) + assert list(cache._ref_audio_resolve_cache) == ["b"] + assert cache._ref_audio_resolve_cache_bytes == 12 + _put(cache, "too-large", [0.0] * 5) + assert "too-large" not in cache._ref_audio_resolve_cache + + +def test_replacement_accounts_bytes_and_preserves_referenced_artifact(cache): + _put(cache, "a", [0.0] * 3, "same") + _put(cache, "alias", [0.0] * 2, "same") + cache._ref_audio_model_artifact_ready.add(("same", False)) + _put(cache, "a", [1.0] * 4, "new") + assert cache._ref_audio_resolve_cache_bytes == 24 + assert ("same", False) in cache._ref_audio_model_artifact_ready + + +def test_pool_accepts_666_references(cache): + for index in range(666): + _put(cache, str(index), [0.0] * 8) + assert len(cache._ref_audio_resolve_cache) == 666 + assert cache._ref_audio_resolve_cache_bytes == 666 * 8 * 4 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 1998a4dd4..587993fca 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 @@ -682,6 +682,7 @@ class TestCodePredictorGraphReplay: return_proj_buf=False, ) predictor._prefix_graphs_enabled = False + predictor._prefix_reprefill_enabled = False predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)} graph = mocker.Mock() graph_output = torch.zeros(1, 4, 4) @@ -1362,3 +1363,38 @@ class TestCodePredictorFusedProjections: with pytest.raises(RuntimeError, match="missing fused parameters"): model.load_weights(weights) + + +class TestMTPExecutionBuckets: + """Outer MRv2 MTP graph buckets must align with the predictor bucket set.""" + + def test_configure_mtp_execution_buckets(self, mocker: MockerFixture, loaded_target_classes) -> None: + _ = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) + + cp_config, _ = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker, max_num_seqs=4) + wrapper = common_mod.CodePredictorWrapper( + vllm_config=vllm_config, + cp_config=cp_config, + wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False), + talker_hidden_size=cp_config.hidden_size, + ) + + # Default keeps the legacy power-of-two derivation. + assert wrapper._execution_batch_buckets is None + assert wrapper._batch_bucket_sizes() == [1, 2, 4] + + wrapper.configure_mtp_execution_buckets([3, 1, 3, 0, 9]) + assert wrapper._batch_bucket_sizes() == [1, 3, 4] # deduplicate/filter declaration and retain max + + with pytest.raises(RuntimeError, match="before the first warmup"): + wrapper2 = common_mod.CodePredictorWrapper( + vllm_config=vllm_config, + cp_config=cp_config, + wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False), + talker_hidden_size=cp_config.hidden_size, + ) + wrapper2._bucket_sizes = [1, 2, 4] # stand-in for warmed-up state + wrapper2.configure_mtp_execution_buckets([3]) diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py index f3e8d0166..23eec9c8f 100644 --- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py +++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for CUDA Graph decoder wrapper numerical equivalence. @@ -142,6 +142,28 @@ def test_padded_output_shape_and_length(decoder, wrapper, seq_len): assert graph_out.shape[-1] == expected_len +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_padded_batch_size_uses_larger_capture_bucket(decoder): + """Non-exact batch sizes should replay the next captured batch bucket.""" + w = CUDAGraphDecoderWrapper( + decoder=decoder, + capture_sizes=[25], + capture_batch_sizes=[1, 4], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + w.warmup(DEVICE) + codes = _random_codes(25, batch_size=3) + + with torch.no_grad(): + eager_out = decoder(codes) + graph_out = w.decode(codes) + + assert (4, 25) in w.graphs + assert graph_out.shape == eager_out.shape + torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) + + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) def test_padded_interior_positions_close(decoder, wrapper, seq_len): diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py new file mode 100644 index 000000000..cd45288aa --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Output metadata must not wait for the preceding model forward on CUDA.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, +) +from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async + +pytestmark = [pytest.mark.core_model] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("validity_device", ["cpu", "cuda"]) +@torch.inference_mode() +def test_output_validity_preserves_values_without_waiting_for_forward(validity_device): + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)) + hidden = torch.zeros((3, 8), device="cuda") + infos = [ + { + "codes": {"audio": torch.zeros((frames, 16), device="cuda", dtype=torch.long)}, + "meta": {"codec_frame_valid": torch.tensor(valid, device=validity_device)}, + } + for frames, valid in [(2, True), (1, False)] + ] + # Warm allocation paths before checking the stream-ordering contract. + talker.make_omni_output(hidden, model_intermediate_buffer=infos) + copy_stream = torch.cuda.Stream() + torch.accelerator.synchronize() + + preceding_forward = torch.cuda.Event() + torch.cuda._sleep(500_000_000) + preceding_forward.record() + output = talker.make_omni_output(hidden, model_intermediate_buffer=infos) + waited_for_forward = preceding_forward.query() + validity = output.multimodal_outputs["meta"]["codec_frame_valid"] + snapshot = _snapshot_tensor_payload_to_cpu_async( + output.multimodal_outputs, copy_stream=copy_stream, pin_memory=True + ) + snapshot.wait() + + assert not waited_for_forward, "output construction synchronized the preceding forward" + assert validity.device.type == "cuda" + assert validity.dtype == torch.int8 + cpu_validity = snapshot.payload["meta"]["codec_frame_valid"] + assert cpu_validity.device.type == "cpu" + assert cpu_validity.tolist() == [1, 1, 0] + assert all(info["meta"]["codec_frame_valid"].device.type == validity_device for info in infos) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py new file mode 100644 index 000000000..8e4d5e094 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise actual stateless decoder capture, not a mocked capture indicator.""" + +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2DecoderConfig, +) +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2DecoderTransformerModel, +) + +pytestmark = [pytest.mark.core_model, *hardware_marks(res={"cuda": "L4"}, num_cards=1)] + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("input_mode", ["implicit", "positions"]) +@torch.inference_mode() +def test_stateless_decoder_capture_replays_new_inputs(batch_size, input_mode): + sequence_length = 8 + torch.manual_seed(42) + config = Qwen3TTSTokenizerV2DecoderConfig( + hidden_size=32, + latent_dim=16, + max_position_embeddings=512, + num_attention_heads=4, + num_key_value_heads=4, + intermediate_size=64, + num_hidden_layers=1, + num_quantizers=2, + sliding_window=16, + decoder_dim=32, + upsample_rates=(2,), + upsampling_ratios=(2,), + ) + config._attn_implementation = "sdpa" + model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config).cuda().eval() + inputs = torch.randn(batch_size, sequence_length, config.latent_dim, device="cuda") + positions = {} + if input_mode == "positions": + cache_position = torch.arange(sequence_length, device="cuda") + positions = {"cache_position": cache_position, "position_ids": cache_position.unsqueeze(0)} + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + model(inputs_embeds=inputs, use_cache=False, **positions) + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize() + + # No wrapper or eager fallback: an illegal capture must fail this test. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = model(inputs_embeds=inputs, use_cache=False, **positions).last_hidden_state + + outputs = [] + for _ in range(3): + new_inputs = torch.randn_like(inputs) + expected = model(inputs_embeds=new_inputs, use_cache=False, **positions).last_hidden_state + if input_mode == "implicit": + # Compare the cached mask against the independent live-mask path, + # not only against eager execution using the same cached mask. + reference = model( + inputs_embeds=new_inputs, + use_cache=False, + position_ids=torch.arange(sequence_length, device="cuda").unsqueeze(0), + ).last_hidden_state + torch.testing.assert_close(expected, reference, atol=1e-5, rtol=1e-4) + inputs.copy_(new_inputs) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(captured_output, expected, atol=1e-5, rtol=1e-4) + outputs.append(captured_output.clone()) + assert not torch.equal(outputs[0], outputs[1]), "replay must consume the new input" + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@torch.inference_mode() +def test_stateless_codec_wrapper_captures_and_replays_waveform(batch_size): + from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _make_small_decoder + from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper + + torch.manual_seed(42) + decoder = _make_small_decoder().cuda().eval() + wrapper = CUDAGraphDecoderWrapper(decoder, num_quantizers=2, async_chunk=False) + # Exercise the production capture method, including quantization and + # waveform convolutions, without accepting warmup's eager fallback. + wrapper._capture_stateless(batch_size, 8, torch.device("cuda"), torch.long) + state = wrapper.stateless_states[(batch_size, 8)] + for _ in range(3): + codes = torch.randint(0, 32, (batch_size, 2, 8), device="cuda") + expected = decoder._forward_exact(codes) + state["input"]["codes"].copy_(codes) + state["graph"].replay() + torch.accelerator.synchronize() + torch.testing.assert_close(state["output"], expected, atol=1e-5, rtol=1e-4) + assert torch.isfinite(state["output"]).all() diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index d176a1ff0..39e3fdce5 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -1,5 +1,9 @@ -from collections import OrderedDict +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections import OrderedDict, UserDict from types import SimpleNamespace +from typing import Any import numpy as np import pytest @@ -36,6 +40,7 @@ def _make_minimal_talker( """ model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16) + model._codebook_vocab_size = 2048 model._embedding_dtype = torch.bfloat16 if tts_pad_embed is None: tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16) @@ -51,6 +56,22 @@ def _make_minimal_talker( return model +def test_postprocess_batch_gathers_each_request_tail(): + from vllm_omni.model_executor.models.output_templates import OwnedBatchTensor + + model = _make_minimal_talker() + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]) + + key, values = model.postprocess_batch_mrv2( + hidden_states=hidden, + last_token_indices=torch.tensor([1, 3]), + ) + + assert key == ("hidden_states", "last") + assert isinstance(values, OwnedBatchTensor) # freshly allocated rows transfer ownership + assert torch.equal(values.tensor, torch.tensor([[3.0, 4.0], [7.0, 8.0]])) + + def _make_minimal_builder( *, config: SimpleNamespace | None = None, @@ -95,6 +116,7 @@ def _make_minimal_builder( ) builder._speaker_cache = None builder._text_tokenizer = None + builder._projected_token_cache = {} builder._embedding_dtype = torch.bfloat16 builder._ref_audio_artifact_cache_max_entries = 256 builder._ref_audio_artifact_cache = OrderedDict() @@ -259,6 +281,30 @@ def test_decode_compacts_long_trailing_text_after_large_offset(): assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:]) +def test_decode_marks_codec_frame_validity_from_processed_input_token(): + model = _make_minimal_talker() + model.embed_input_ids = lambda input_ids: input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4) + common = { + "input_embeds": None, + "text": ["hello"], + "task_type": ["CustomVoice"], + "hidden_states": { + "trailing_text": torch.ones((2, 4)), + "last": torch.ones(4), + }, + "meta": {"talker_text_offset": 0}, + "_omni_is_prefill": False, + "_omni_num_computed_tokens": 2, + "_omni_prompt_len": 2, + } + + _, _, valid_update = model.preprocess(input_ids=torch.tensor([123]), **common) + _, _, eos_update = model.preprocess(input_ids=torch.tensor([4198]), **common) + + assert valid_update["meta"]["codec_frame_valid"].item() is True + assert eos_update["meta"]["codec_frame_valid"].item() is False + + def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state(): model = _make_minimal_talker() @@ -286,14 +332,17 @@ def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state(): out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16), ) - assert update == {"meta": {"codec_streaming": True}} + assert update["meta"]["codec_streaming"] is True + assert update["meta"]["codec_frame_valid"].item() is False -def test_decode_batch_preprocess_matches_decode_state_updates(): +@pytest.mark.parametrize("mrv2", [False, True]) +def test_decode_batch_preprocess_matches_decode_state_updates(mrv2): tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) model = _make_minimal_talker(tts_pad_embed=tts_pad) def fake_embed_input_ids(input_ids): + assert not mrv2, "MRV2 must reuse the prepared token embeddings" return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4) model.embed_input_ids = fake_embed_input_ids @@ -302,7 +351,10 @@ def test_decode_batch_preprocess_matches_decode_state_updates(): last_a = torch.full((4,), 2.0, dtype=torch.float32) last_b = torch.full((4,), 3.0, dtype=torch.float32) - out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch( + preprocess = model.preprocess_decode_batch_mrv2 if mrv2 else model.preprocess_decode_batch + prepared = torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16) + out_ids, out_embeds, past_hidden, text_step, updates = preprocess( + **({"input_embeds": prepared} if mrv2 else {}), input_ids=torch.tensor([101, 202], dtype=torch.long), req_infos=[ { @@ -321,15 +373,19 @@ def test_decode_batch_preprocess_matches_decode_state_updates(): ) assert out_ids.tolist() == [101, 202] + if mrv2: + assert out_embeds.data_ptr() == prepared.data_ptr() assert torch.equal(out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16)) assert torch.equal(past_hidden.cpu(), torch.stack([last_a, last_b]).to(torch.bfloat16)) assert torch.equal(text_step[0].cpu(), trailing_a[1].to(torch.bfloat16)) assert torch.equal(text_step[1].cpu(), tts_pad.reshape(-1).to(torch.bfloat16)) assert updates[0]["meta"]["talker_text_offset"] == 2 assert updates[0]["meta"]["codec_streaming"] is True + assert updates[0]["meta"]["codec_frame_valid"].item() is True assert "hidden_states" not in updates[0] assert updates[1]["meta"]["talker_text_offset"] == 0 assert updates[1]["meta"]["codec_streaming"] is False + assert updates[1]["meta"]["codec_frame_valid"].item() is True assert updates[1]["hidden_states"]["trailing_text"].numel() == 0 @@ -449,11 +505,11 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): def __call__(self, texts, *, padding=False): self.calls.append((texts, padding)) - return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]} + return UserDict({"input_ids": [[1, 2, 3], [2, 3, 4, 5]]}) text_tok = FakeTextTokenizer() builder._text_tokenizer = text_tok - buf = { + buf: dict[str, dict[str, Any]] = { "r1": { "task_type": ["Base"], "text": ["one"], @@ -488,9 +544,17 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate(): assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2 assert len(text_tok.calls) == 2 assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3])) - assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4])) + assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4, 5])) assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3])) - assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4])) + assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4, 5])) + assert ( + buf["r1"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr() + == buf["r2"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr() + ) + assert ( + buf["r1"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr() + == buf["r2"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr() + ) def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer(): @@ -506,7 +570,7 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit return {"input_ids": [[7, 8, 9] for _ in texts]} builder._text_tokenizer = FakeTextTokenizer() - buf = { + buf: dict[str, dict[str, Any]] = { "r1": { "task_type": ["Base"], "text": ["one"], @@ -539,7 +603,7 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw( AssertionError("speech tokenizer not expected") ) - buf = { + buf: dict[str, dict[str, Any]] = { "r1": { "task_type": ["Base"], "text": ["one"], @@ -574,7 +638,7 @@ def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_witho return {"input_ids": [[7, 8, 9] for _ in texts]} builder._text_tokenizer = FakeTextTokenizer() - buf = { + buf: dict[str, dict[str, Any]] = { "r1": { "task_type": ["Base"], "text": ["one"], diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py index 3c145e4ca..b830aa835 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for cross-request ``codes.ref`` leakage (#4370). ``make_omni_output`` used to collapse every request's reference codec @@ -11,11 +11,14 @@ vocoder chunks then decoded with another request's reference voice as context (audible onset timbre deformation at any concurrency > 1). """ +from types import SimpleNamespace + import pytest import torch from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( Qwen3TTSTalkerForConditionalGeneration, + _qwen3_tts_gpu_resident_buffer_keys, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -23,6 +26,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _NUM_CODE_GROUPS = 16 +def test_large_tts_prefill_artifacts_stay_gpu_resident_only_on_mrv2() -> None: + v1_keys = _qwen3_tts_gpu_resident_buffer_keys(False) + v2_keys = _qwen3_tts_gpu_resident_buffer_keys(True) + + assert ("embed", "prefill") not in v1_keys + assert ("codes", "ref") not in v1_keys + assert ("meta", "codec_frame_valid") not in v1_keys + assert ("embed", "prefill") in v2_keys + assert ("codes", "ref") in v2_keys + assert ("meta", "codec_frame_valid") in v2_keys + + def _make_talker() -> Qwen3TTSTalkerForConditionalGeneration: # make_omni_output reads no instance state, so skip __init__ (no real # config / checkpoint needed), same as the preprocess tests. @@ -89,3 +104,42 @@ def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_cod assert torch.equal(out.text_hidden_states, hidden) assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS) + + +def test_make_omni_output_preserves_per_request_codec_frame_validity() -> None: + talker = _make_talker() + talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)) + valid = _info(1, None) + invalid = _info(1, None) + valid["meta"]["codec_frame_valid"] = True + invalid["meta"]["codec_frame_valid"] = False + + out = talker.make_omni_output( + torch.zeros((2, 8)), + model_intermediate_buffer=[valid, invalid], + ) + + assert out.multimodal_outputs["meta"]["codec_frame_valid"].tolist() == [1, 0] + + +def test_make_omni_output_publishes_ref_once_for_kv_resumed_first_decode() -> None: + talker = _make_talker() + talker.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(async_chunk=True), + ) + ref_code = torch.ones((5, _NUM_CODE_GROUPS), dtype=torch.long) + info = _info(1, ref_code) + info["_omni_num_computed_tokens"] = 9 + info["_omni_prompt_len"] = 5 + + first = talker.make_omni_output( + torch.zeros((1, 8)), + model_intermediate_buffer=[info], + ) + second = talker.make_omni_output( + torch.zeros((1, 8)), + model_intermediate_buffer=[info], + ) + + assert first.multimodal_outputs["codes"]["ref"][0] is ref_code + assert "ref" not in second.multimodal_outputs["codes"] diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py index 3552c5356..3b58eb7d2 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048). ``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary @@ -59,6 +59,7 @@ def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None): object.__setattr__(talker, "_silence_mask", silence_mask) object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool)) object.__setattr__(talker, "_batch_req_ids", req_ids) + object.__setattr__(talker, "_mrv2_silence_ban_mask", None) object.__setattr__( talker, "_req_x_vector_only", @@ -78,7 +79,26 @@ def _banned(row: torch.Tensor) -> set[int]: return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()} -# -------------------- the step gate -------------------- +def test_mrv2_silence_ban_tracks_batch_order_scope_and_decode_window() -> None: + talker = _make_talker(ban_frames=3, logits=torch.zeros((3, _VOCAB)), modes=[True, False, True]) + with talker.mrv2_sampling_context( + req_ids=["req2", "req1", "req0"], + num_output_tokens=torch.tensor([3, 0, 1]), + ): + logits = talker.compute_logits(torch.zeros((3, 4))) + + assert _banned(logits[0]) == set() + assert _banned(logits[1]) == set() + assert _banned(logits[2]) == set(_SILENCE_IDS) + + # Prompt logprobs are computed after sampling and must remain unmasked. + assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any() + + with talker.mrv2_sampling_context( + req_ids=["req2", "req1", "req0"], + num_output_tokens=torch.tensor([4, 1, 3]), + ): + assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any() def test_ban_disabled_leaves_logits_unchanged() -> None: 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 4b542bf77..f929e41c9 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 @@ -7,6 +7,9 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.model_executor.stage_input_processors import ( + qwen3_tts as qwen3_tts_processors, +) from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( AdaptiveChunkController, compute_adaptive_emit, @@ -32,7 +35,52 @@ _FRAME = [1, 2, 3, 4] _Q = len(_FRAME) -def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=None, full_utterance_decode=None): +@pytest.mark.parametrize("frames", [1, 26]) +def test_terminal_without_new_frames_does_not_replay_last_chunk(frames): + tm = _tm(initial_chunk_frames=1) + emitted = 0 + for count in range(1, frames + 1): + payload = _call(tm, "r", n_frames=count) + if payload is not None: + emitted += payload.codes.audio.numel() // _Q + tm.put_req_chunk["r"] += 1 + terminal = _call(tm, "r", n_frames=frames, finished=True) + assert terminal is not None and terminal.meta.finished.item() + assert terminal.codes.audio.numel() == 0 + assert emitted == frames + + +def test_repeated_nonterminal_callback_does_not_republish_boundary(): + tm = _tm(initial_chunk_frames=1) + first = _call(tm, "r", n_frames=1) + assert first is not None + tm.put_req_chunk["r"] += 1 + assert _call(tm, "r", n_frames=1) is None + next_chunk = _call(tm, "r", n_frames=26) + assert next_chunk.codes.audio.numel() // _Q == 25 + + +def test_terminal_flush_contains_only_frames_not_previously_emitted(): + tm = _tm(initial_chunk_frames=1) + _call(tm, "r", n_frames=1) + tm.put_req_chunk["r"] += 1 + _call(tm, "r", n_frames=26) + tm.put_req_chunk["r"] += 1 + final = _call(tm, "r", n_frames=30, finished=True) + assert final.codes.audio.numel() // _Q == 4 + tm.put_req_chunk["r"] += 1 + assert _call(tm, "r", n_frames=30, finished=True).codes.audio.numel() == 0 + + +def _req( + rid, + *, + finished, + initial_codec_chunk_frames=None, + non_streaming_mode=None, + full_utterance_decode=None, + output_token_ids=None, +): entries = {} if initial_codec_chunk_frames is not None: entries["initial_codec_chunk_frames"] = SimpleNamespace(list_data=[initial_codec_chunk_frames]) @@ -40,11 +88,11 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=N entries["non_streaming_mode"] = SimpleNamespace(list_data=[non_streaming_mode]) if full_utterance_decode is not None: entries["full_utterance_decode"] = SimpleNamespace(list_data=[full_utterance_decode]) - ai = SimpleNamespace(entries=entries) if entries else None return SimpleNamespace( external_req_id=rid, is_finished=lambda: finished, - additional_information=ai, + additional_information=SimpleNamespace(entries=entries) if entries else None, + output_token_ids=output_token_ids or [], ) @@ -206,6 +254,112 @@ def test_voicedesign_prompt_mode_with_streaming_keeps_windowed_emit(): assert len(p.codes.audio) == _Q * 25 +def test_async_chunk_skips_stop_token_zero_frame_without_gpu_value_check(): + tm = _tm() + rid = "r-stop" + tm.code_prompt_token_ids[rid] = [_FRAME[:]] + + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}}, + request=_req(rid, finished=True, initial_codec_chunk_frames=2, output_token_ids=[2150]), + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.tolist() == _FRAME + assert len(tm.code_prompt_token_ids[rid]) == 1 + + +def test_async_chunk_uses_authoritative_frame_validity_without_token_history(): + tm = _tm() + rid = "r-native-validity" + + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={ + "codes": {"audio": torch.tensor([[1, 2, 3, 4]])}, + "meta": {"codec_frame_valid": False}, + }, + request=_req(rid, finished=False, initial_codec_chunk_frames=1, output_token_ids=[]), + is_finished=False, + ) + + assert payload is None + assert tm.code_prompt_token_ids[rid] == [] + + +def test_async_chunk_keeps_final_real_frame_when_eos_was_just_sampled(): + tm = _tm() + rid = "r-final-real-frame" + + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={ + "codes": {"audio": torch.tensor([[1, 2, 3, 4]])}, + "meta": {"codec_frame_valid": True}, + }, + request=_req(rid, finished=True, initial_codec_chunk_frames=1, output_token_ids=[4198]), + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.tolist() == [1, 2, 3, 4] + + +def test_async_chunk_batch_matches_scalar_payloads_and_state(): + batch_builder = getattr(qwen3_tts_processors, "talker2code2wav_async_chunk_batch", None) + assert callable(batch_builder), "Qwen3-TTS must provide the native MRv2 batch payload builder" + + scalar_tm = _tm(chunk_frames=4, left_context=3, initial_chunk_frames=1) + batch_tm = _tm(chunk_frames=4, left_context=3, initial_chunk_frames=1) + requests = [ + _req("r-batch-0", finished=False, output_token_ids=[7]), + _req("r-batch-1", finished=False, output_token_ids=[8]), + ] + ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) + # Request 2 has a multi-token span: the last validity flag wins the chunk. + outputs = [ + {"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}}, + { + "codes": {"audio": torch.tensor([[9, 9, 9, 9], [5, 6, 7, 8]])}, + "meta": {"codec_frame_valid": torch.tensor([False, True])}, + }, + ] + + scalar_payloads = [ + talker2code2wav_async_chunk( + transfer_manager=scalar_tm, + multimodal_output=output, + request=request, + is_finished=False, + ) + for request, output in zip(requests, outputs) + ] + batch_payloads = batch_builder( + transfer_manager=batch_tm, + pooling_outputs=outputs, + requests=requests, + is_finished=[False, False], + ) + + assert len(batch_payloads) == len(scalar_payloads) + for scalar, batched in zip(scalar_payloads, batch_payloads): + assert scalar is not None and batched is not None + torch.testing.assert_close(batched.codes.audio, scalar.codes.audio) + assert batched.meta.left_context_size == scalar.meta.left_context_size + assert batched.meta.finished.item() == scalar.meta.finished.item() + assert batched.speaker == scalar.speaker + assert batched.language == scalar.language + for request in requests: + rid = request.external_req_id + torch.testing.assert_close( + torch.stack(batch_tm.code_prompt_token_ids[rid]), + torch.stack(scalar_tm.code_prompt_token_ids[rid]), + ) + torch.testing.assert_close(batch_tm.request_payload["r-batch-0"], scalar_tm.request_payload["r-batch-0"]) + + _CASES = [ # ── IC boundary rule ────────────────────────────────────────────── # initial_codec_chunk_frames only controls the first emitted chunk. @@ -703,35 +857,27 @@ class TestRampHelpers: def test_parse_ramp_mixed_list_returns_none(self): assert parse_chunk_ramp({"codec_chunk_ramp": [4, "x"]}) is None - def test_parse_ramp_warns_on_tail_mismatch(self, caplog): - import logging + def test_parse_ramp_warns_on_tail_mismatch(self, monkeypatch): + warnings = [] + monkeypatch.setattr( + "vllm_omni.model_executor.stage_input_processors.chunk_size_utils.logger.warning", + lambda message, *args: warnings.append(message % args), + ) - target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") - target_logger.addHandler(caplog.handler) - prev_level = target_logger.level - target_logger.setLevel(logging.WARNING) - try: - result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25) - finally: - target_logger.removeHandler(caplog.handler) - target_logger.setLevel(prev_level) + result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25) assert result == [4, 4, 8] - assert "reintroduces" in caplog.text + assert any("reintroduces" in message for message in warnings) - def test_parse_ramp_no_warn_on_tail_match(self, caplog): - import logging + def test_parse_ramp_no_warn_on_tail_match(self, monkeypatch): + warnings = [] + monkeypatch.setattr( + "vllm_omni.model_executor.stage_input_processors.chunk_size_utils.logger.warning", + lambda message, *args: warnings.append(message % args), + ) - target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") - target_logger.addHandler(caplog.handler) - prev_level = target_logger.level - target_logger.setLevel(logging.WARNING) - try: - result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25) - finally: - target_logger.removeHandler(caplog.handler) - target_logger.setLevel(prev_level) + result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25) assert result == [4, 4, 8, 16, 25] - assert "reintroduces" not in caplog.text + assert warnings == [] @pytest.mark.parametrize( "index,ramp,steady,expected", @@ -1028,6 +1174,7 @@ class TestChunkRampEmission: tm.put_req_chunk[rid] = 2 tm.code_prompt_token_ids[rid] = [] + tm._qwen3_tts_emitted_frames.pop(rid, None) tm.ramp_chunk_count.pop(rid, None) p_seg2_0 = self._emit(tm, rid, 4) @@ -1051,6 +1198,7 @@ class TestChunkRampEmission: tm.put_req_chunk[rid] = 1 tm.code_prompt_token_ids[rid] = [] + tm._qwen3_tts_emitted_frames.pop(rid, None) tm.ramp_chunk_count.pop(rid, None) p_seg2_0 = self._emit(tm, rid, 4) @@ -1429,6 +1577,7 @@ class TestAdaptiveEmission: tm.put_req_chunk[rid] = 1 tm.code_prompt_token_ids[rid] = [] + tm._qwen3_tts_emitted_frames.pop(rid, None) tm.ramp_chunk_count.pop(rid, None) tm._adaptive_states.pop(rid, None) diff --git a/tests/worker/test_min_tokens_stop_ids_guard.py b/tests/worker/test_min_tokens_stop_ids_guard.py index 5d7591c99..a99e74b7d 100644 --- a/tests/worker/test_min_tokens_stop_ids_guard.py +++ b/tests/worker/test_min_tokens_stop_ids_guard.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Guard against out-of-vocabulary stop ids in min-tokens masking (#4962). vLLM's input processor folds the stage tokenizer's EOS id into @@ -19,7 +19,7 @@ from vllm.v1.sample.logits_processor import ( MinTokensLogitsProcessor, ) -from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids, sanitize_sampling_params_min_tokens_stop_ids pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -28,6 +28,12 @@ CODEC_EOS = 2150 TEXT_EOS = 151645 +@pytest.fixture(autouse=True) +def _disable_vllm_pin_memory_for_cpu(monkeypatch: pytest.MonkeyPatch): + """Keep the module-level H2D helper in CPU-only mode.""" + monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False) + + def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor: """Build a MinTokensLogitsProcessor with one active request, mirroring the production path (engine folds the tokenizer EOS via @@ -113,3 +119,14 @@ def test_guard_handles_multiple_requests(): assert out[0, CODEC_EOS] == -float("inf") assert out[1, 7] == -float("inf") assert int((out == -float("inf")).sum()) == 2 + + +def test_sampling_params_guard_filters_only_unreachable_min_tokens_ids(): + params = SamplingParams(min_tokens=2, stop_token_ids=[CODEC_EOS]) + params.update_from_generation_config({}, TEXT_EOS) + + sanitize_sampling_params_min_tokens_stop_ids(params, TALKER_VOCAB) + + assert params.all_stop_token_ids == {CODEC_EOS} + assert params.stop_token_ids == [CODEC_EOS] + assert params.eos_token_id == TEXT_EOS diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index ec2d4a704..bdf9936a4 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -1724,3 +1724,16 @@ class TestSendRetry(unittest.TestCase): if __name__ == "__main__": unittest.main() + + +@pytest.mark.parametrize( + "fallback,external,expected", [("r", None, "mapped"), (None, "ext", "ext"), (None, None, None)] +) +def test_resolve_request_id_requires_internal_or_external_id(fallback, external, expected): + host = SimpleNamespace(_request_ids_mapping={"r": "mapped"}) + request = SimpleNamespace(external_req_id=external) + if expected is None: + with pytest.raises(ValueError, match="request ID"): + OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback) + else: + assert OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback) == expected diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 4155e5ffb..ee51b0f58 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -770,8 +770,10 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh def test_full_payload_output_accumulation_hook_matrix(): assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output() assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output() - assert not _make_full_payload_accumulation_runner( - model_stage="code2wav", final_output=True + # A stage may publish a user-visible output and still feed another stage + # (Qwen3-Omni thinker publishes text while sending audio state to talker). + assert _make_full_payload_accumulation_runner( + model_stage="thinker", final_output=True )._should_accumulate_full_payload_output() assert not _make_full_payload_accumulation_runner( model_stage="token2audio", diff --git a/tests/worker/test_omni_v2_profile_dispatch.py b/tests/worker/test_omni_v2_profile_dispatch.py new file mode 100644 index 000000000..4ff6fbfd4 --- /dev/null +++ b/tests/worker/test_omni_v2_profile_dispatch.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +from vllm.compilation.cuda_graph import CUDAGraphMode + +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("dp_size", [1, 2]) +def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size): + runner = object.__new__(OmniGPUModelRunner) + runner.cudagraph_manager = None + runner.dp_size = dp_size + runner.dp_rank = 1 + + import vllm.v1.worker.gpu.dp_utils as dp_utils + + expected_tokens = object() + sync_calls = [] + + def sync_padding(*args, **kwargs): + sync_calls.append((args, kwargs)) + return args[1], expected_tokens + + monkeypatch.setattr( + dp_utils, + "sync_cudagraph_and_dp_padding", + sync_padding, + ) + + batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor( + num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=True, max_query_len=8 + ) + assert batch_desc.cg_mode == CUDAGraphMode.NONE + if dp_size == 1: + assert num_tokens_across_dp is None + assert sync_calls == [] + else: + assert num_tokens_across_dp is expected_tokens + + # Non-eager dispatch goes through the cudagraph manager. + expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1) + dispatch_calls = [] + + def dispatch(*args, **kwargs): + dispatch_calls.append((args, kwargs)) + return expected + + runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch) + runner.dp_size = 1 + batch_desc, _ = runner._dispatch_batch_descriptor( + num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8 + ) + assert batch_desc is expected + assert dispatch_calls == [((1, 8, 8), {"num_active_loras": 0, "max_query_len": 8})] diff --git a/tests/worker/test_tts_sampling_seed.py b/tests/worker/test_tts_sampling_seed.py new file mode 100644 index 000000000..f54280c36 --- /dev/null +++ b/tests/worker/test_tts_sampling_seed.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.worker import sampling_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_fixed_seed_contract_separates_stage_sampler_from_talker_mtp_rng() -> None: + assert hasattr(sampling_utils, "apply_fixed_seed_to_sampling_params") + assert hasattr(sampling_utils, "get_tts_local_seed") + + stage_only = SimpleNamespace(seed=None, extra_args=None) + sampling_utils.apply_fixed_seed_to_sampling_params( + stage_only, + 17, + seed_talker_mtp=False, + ) + assert stage_only.seed == 17 + assert sampling_utils.get_tts_local_seed(stage_only) is None + + talker = SimpleNamespace(seed=None, extra_args={"preserved": True}) + sampling_utils.apply_fixed_seed_to_sampling_params( + talker, + 23, + seed_talker_mtp=True, + ) + assert talker.seed == 23 + assert talker.extra_args == {"preserved": True, "tts_local_seed": 23} + assert sampling_utils.get_tts_local_seed(talker) == 23 diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index 18131a31c..a04cc35c2 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Characterization tests for ``vllm_omni.worker.base.OmniGPUWorkerBase``. Pins the CURRENT behaviour of ``determine_available_memory`` (device-level diff --git a/tests/worker_v2/__init__.py b/tests/worker_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/worker_v2/test_intermediate_buffer.py b/tests/worker_v2/test_intermediate_buffer.py new file mode 100644 index 000000000..a84fef711 --- /dev/null +++ b/tests/worker_v2/test_intermediate_buffer.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for OmniIntermediateBuffer lifecycle. + +Retained contracts: wire decode with malformed-payload rejection, slot +lifecycle with index migration on reuse, gather in batch order, update merge +semantics (cpu/gpu-resident keys), single owned snapshot for +update_gpu_tensor_rows, and decode-span merge until ack. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.engine import ( + AdditionalInformationEntry, + AdditionalInformationPayload, + PromptEmbedsPayload, +) +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + OmniIntermediateBuffer, + _resolve_additional_information, + _resolve_prompt_embeds, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_new_req_data(req_id="r1", **kwargs): + kwargs.setdefault("mm_features", []) + return SimpleNamespace(req_id=req_id, **kwargs) + + +def test_wire_decode_passthrough_and_malformed_rejection(): + tensor = torch.randn(3, 4) + assert _resolve_prompt_embeds(tensor) is not None + assert _resolve_prompt_embeds(None) is None + with pytest.raises(ValueError, match="Failed to decode prompt_embeds payload"): + _resolve_prompt_embeds(PromptEmbedsPayload(data=b"\x00", shape=[2, 2], dtype="float32")) + + assert _resolve_additional_information({"key": "val"}) == {"key": "val"} + assert _resolve_additional_information(None) == {} + payload = AdditionalInformationPayload( + entries={ + "embed.decode": AdditionalInformationEntry(tensor_data=b"\x00", tensor_shape=[2], tensor_dtype="float32") + } + ) + with pytest.raises(ValueError, match="Failed to decode additional_information payload"): + _resolve_additional_information(payload) + + +def test_slot_lifecycle_fields_reuse_and_remove_idempotent(): + buf = OmniIntermediateBuffer(max_num_reqs=4) + embeds = torch.randn(2, 8) + features = [torch.randn(3, 4)] + params = SimpleNamespace(extra_args={"tts_local_seed": 1234}) + buf.add_request(0, _make_new_req_data("r0", prompt_embeds=embeds, mm_features=features, sampling_params=params)) + assert torch.equal(buf.buffers[0]["prompt_embeds_cpu"], embeds) + assert buf.buffers[0]["mm_features"] is features + assert buf.buffers[0]["sampling_params"] is params + + buf.add_request(2, _make_new_req_data("r2")) + buf.add_request(3, _make_new_req_data("r3")) + batch = SimpleNamespace(idx_mapping_np=[2, 0, 3]) + assert [g["req_id"] for g in buf.gather(batch)] == ["r2", "r0", "r3"] + + # Slot reuse: req_id mapping migrates; duplicate removal is a no-op. + buf.add_request(0, _make_new_req_data("new")) + assert buf.req_id_to_index["new"] == 0 + buf.remove_request(0) + buf.remove_request(0) + assert "new" not in buf.req_id_to_index + assert buf.buffers[0] == {} + + +def test_update_merge_semantics(monkeypatch): + buf = OmniIntermediateBuffer(max_num_reqs=1) + buf.add_request(0, _make_new_req_data("r0")) + + buf.update(0, {"hidden": torch.randn(4)}) + assert buf.buffers[0]["hidden"].device == torch.device("cpu") # default migrates to cpu + buf.update(0, {}) + buf.update(0, {"a": torch.tensor([1.0])}) + buf.update(0, {"b": torch.tensor([2.0])}) + assert "a" in buf.buffers[0] and "b" in buf.buffers[0] + + # GPU-resident keys keep values (including list items) off the cpu path. + tensor = torch.randn(2) + monkeypatch.setattr(torch.Tensor, "cpu", lambda *_a, **_k: pytest.fail("GPU-resident value must not move to CPU")) + buf.update(0, {"items": [tensor], "kv": torch.randn(4)}, gpu_resident_keys={"items", "kv"}) + stored = buf.buffers[0]["items"][0] + assert stored is not tensor + assert torch.equal(stored, tensor) + assert "kv" in buf.buffers[0] diff --git a/tests/worker_v2/test_intermediate_buffer_cuda_residency.py b/tests/worker_v2/test_intermediate_buffer_cuda_residency.py new file mode 100644 index 000000000..8368d9fe2 --- /dev/null +++ b/tests/worker_v2/test_intermediate_buffer_cuda_residency.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CUDA-lane checks for GPU-resident intermediate buffer values.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize("key,dtype", [(("embed", "prefill"), torch.float32), (("codes", "ref"), torch.long)]) +def test_gpu_resident_values_stay_on_device_after_update(key, dtype): + # Declared gpu-resident keys must keep tensors on device (CLA snapshot real). + buffer = OmniIntermediateBuffer(1) + value = torch.ones(2, 3, device="cuda", dtype=dtype) + buffer.add_request(0, SimpleNamespace(req_id="r1", mm_features=[])) + + buffer.update(0, {key[0]: {key[1]: value}}, {key}) + + stored = buffer.buffers[0][key[0]][key[1]] + assert stored.device.type == "cuda" + assert stored.dtype == dtype + assert torch.equal(stored.cpu(), value.cpu()) diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py new file mode 100644 index 000000000..d1f45b5cc --- /dev/null +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for OmniARModelRunner v2: async output staging, snapshot ownership, payload slicing.""" + +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +from vllm.v1.outputs import RoutedExpertsTensors +from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors + +import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner +from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput +from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeStream: + def wait_stream(self, _stream) -> None: + pass + + +class _FakeEvent: + def record(self, _stream) -> None: + pass + + def synchronize(self) -> None: + pass + + +def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput: + rid2idx = {rid: i for i, rid in enumerate(req_ids)} + mro = omni_ar_model_runner.OmniModelRunnerOutput(list(req_ids), rid2idx, None, prompt_logprobs_dict={}) + sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, None) + kwargs = dict(model_runner_output=mro, sampler_output=sampler_output) + kwargs.update(num_sampled_tokens=torch.tensor([1] * len(req_ids)), copy_event=_FakeEvent()) + kwargs.update(main_stream=_FakeStream(), copy_stream=_FakeStream()) + return OmniAsyncOutput(**(kwargs | overrides)) + + +def test_async_output_blocking_event_and_routing_masks(monkeypatch) -> None: + event_kwargs = [] + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + + def make_event(**kwargs): + event_kwargs.append(kwargs) + return _FakeEvent() + + monkeypatch.setattr(torch.cuda, "Event", make_event) + masks = SamplingMaskTensors(torch.tensor([[5], [0]], dtype=torch.uint8), torch.tensor([2, 0]), 4) + sampler_output = SamplerOutput(torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), None, masks) + routed = RoutedExpertsTensors(torch.tensor([[[2, 3]], [[4, 5]]]), torch.tensor([7, 9])) + + output = _async_output( + req_ids=["decode", "prefill"], + sampler_output=sampler_output, + num_sampled_tokens=torch.tensor([1, 0]), + copy_event=None, # None → constructor builds the default blocking event + routed_experts=routed, + ).get_output() + assert event_kwargs == [{"blocking": True}] # blocking event by default + assert output.sampled_token_ids == [[2], []] and output.sampling_masks.cu_num_generated_tokens == [0, 1, 1] + np.testing.assert_array_equal(output.routed_experts.routing_data, [[[2, 3]], [[4, 5]]]) + np.testing.assert_array_equal(output.routed_experts.slot_mapping, [7, 9]) + np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2]) + + +@pytest.mark.parametrize("needs_history", [False, True]) +def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + input_batch = SimpleNamespace(req_ids=["req"], num_reqs=1, seq_lens=torch.tensor([3])) + input_batch.idx_mapping, input_batch.query_start_loc = torch.tensor([0]), torch.tensor([0, 1]) + state = SimpleNamespace(input_batch=input_batch, hidden_states=torch.zeros(1, 2)) + state.finished_req_ids, state.ec_connector_output, state.routed_experts = {"finished"}, None, None + runner.execute_model_state = state + runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None + runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False + runner.model_config = SimpleNamespace(async_chunk=False) + runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text")) + runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(torch.zeros(1, 2), None))) + runner.model_state.intermediate_buffer = SimpleNamespace(buffers={0: {"global_request_id": "global-req"}}) + runner.req_states = SimpleNamespace(req_id_to_index={"req": 0}) + runner.req_states.all_token_ids = SimpleNamespace(gpu=torch.tensor([[1]])) + runner.req_states.num_computed_tokens = SimpleNamespace(gpu=torch.tensor([0])) + runner.req_states.prompt_len = SimpleNamespace(np=np.array([1]), gpu=torch.tensor([1])) + runner.main_stream = runner.output_copy_stream = MagicMock() + runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = MagicMock() + sampler_out = (SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), MagicMock(), MagicMock()) + sampling_active = False + + @contextmanager + def sampling_context(*, req_ids, num_output_tokens): + nonlocal sampling_active + assert needs_history and req_ids == ["req"] + sampling_active = True + yield + sampling_active = False + + runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history) + runner.model.mrv2_sampling_context = sampling_context + runner.sample = MagicMock(return_value=sampler_out) + runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail() + logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {}) + runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=logprobs_mock) + runner.postprocess_sampled, connector_output = MagicMock(), object() + + def post_forward(finished_req_ids): + runner.postprocess_sampled.assert_called_once() # postprocess precedes kv_connector.post_forward + assert finished_req_ids == {"finished"} + return connector_output + + runner.kv_connector = SimpleNamespace(post_forward=MagicMock(side_effect=post_forward)) + mock_out = SimpleNamespace(copy_event=None) + monkeypatch.setattr(omni_ar_model_runner, "OmniAsyncOutput", MagicMock(return_value=mock_out)) + + assert runner.sample_tokens(None) is mock_out + built = omni_ar_model_runner.OmniAsyncOutput.call_args.kwargs["model_runner_output"] + assert built.kv_connector_output is connector_output + assert runner._resolve_global_request_id("req") == "global-req" # from the intermediate buffer + assert runner._resolve_global_request_id("unknown") == "unknown" # fallback to the local id + + +def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.model_config = SimpleNamespace(async_chunk=True) + 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] = [] + runner.main_stream = SimpleNamespace(wait_event=waited.append) + source = torch.tensor([[7, 8]], dtype=torch.long) + + snapshot = runner._retain_multimodal_outputs({"codes": {"audio": source}}) + source.fill_(99) + + # Snapshot owns the data (graph replay cannot overwrite it)... + snap = snapshot["codes"]["audio"] + assert snap.tolist() == [[7, 8]] and snap.data_ptr() != source.data_ptr() + assert runner._last_multimodal_snapshot_slot == 0 + # ...and slot reuse waits for the previous D2H copy event. + runner._release_multimodal_snapshot(0, copy_event := object()) + runner._retain_multimodal_outputs({"codes": {"audio": torch.zeros(1, 2)}}) + assert waited == [copy_event] + + +def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None: + slot: dict[tuple[Any, ...], torch.Tensor] = {} + omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot) + omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(4, 2), slot) + assert len(slot) == 2 # separate shapes do not share a buffer + monkeypatch.setattr(omni_ar_model_runner, "_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT", 1) + bounded: dict[tuple[Any, ...], torch.Tensor] = {} + kept = omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(2, 2), bounded) + overflow = omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(5, 2), bounded) + assert len(bounded) == 1 and overflow.data_ptr() != kept.data_ptr() # overflow clones without evicting + + source = torch.arange(12, dtype=torch.int64).view(3, 4).t() + payload = {"noncontiguous": source, "nested": [torch.tensor(7), (torch.tensor([1.5]),)], "meta": "ok"} + pack_slot: dict[tuple[Any, ...], torch.Tensor] = {} + snapshot = pack_output_snapshot(payload, pack_slot, max_buckets=4) + source.fill_(99) # ownership: snapshot keeps pre-mutation values + copies = [] + + def copy(tensor): + copies.append(tensor.numel()) + return tensor.clone() + + host = snapshot.copy_to_cpu(copy) + assert len(copies) == 2 # grouped by dtype (int64, float32), not tensor count + assert host["noncontiguous"].tolist() == torch.arange(12).view(3, 4).t().tolist() + assert host["nested"][0].item() == 7 and isinstance(host["nested"][1], tuple) and host["meta"] == "ok" + pack_output_snapshot(payload, pack_slot, max_buckets=4) + assert host["noncontiguous"][0, 0].item() == 0 # repacking must not corrupt published host data + + +@pytest.mark.parametrize("need_pooler,async_chunk", [(False, False), (False, True), (True, False), (True, True)]) +def test_guard_graph_replay_for_pooler_copy(need_pooler, async_chunk) -> None: + main_stream = MagicMock() + omni_ar_model_runner._guard_graph_replay_for_pooler_copy( + main_stream, object(), need_pooler=need_pooler, async_chunk=async_chunk + ) + # Non-async pooler copies must gate the next graph replay on the copy event. + assert main_stream.wait_event.call_count == (1 if need_pooler and not async_chunk else 0) + + +def test_build_async_chunk_outputs_slices_padded_axis_and_splits_channels() -> None: + # Graph-padded batch: padded_total_tokens > total_tokens; slice by the real token axis. + padded_codes = torch.arange(16, dtype=torch.long).reshape(8, 2) + build = OmniARModelRunner._build_async_chunk_outputs_from_mm + inter_stage, client = build({"codes": {"audio": padded_codes}}, np.array([0, 1, 2]), np.array([1, 1]), 2, 2, 8) + assert client is None + assert torch.equal(inter_stage[0]["codes.audio"], padded_codes[0:1]) + assert torch.equal(inter_stage[1]["codes.audio"], padded_codes[1:2]) + + codes, audio = torch.arange(16, dtype=torch.long).reshape(4, 4), torch.randn(4, 8) + req_codes = [torch.arange(16 * i, 16 * (i + 1), dtype=torch.long).reshape(1, 16) for i in range(2)] + inter_stage, client = build( + {"codes": {"audio": codes}, "audio": audio}, np.array([0, 2, 4]), np.array([2, 2]), 2, 4 + ) + assert "hidden" not in inter_stage[0] and torch.equal(client[1]["audio"], audio[2:]) # channel split + assert torch.equal(inter_stage[1]["codes.audio"], codes[2:]) # inter-stage channel + # Per-request code lists pass through unsliced. + inter_stage, client = build({"codes": {"audio": req_codes}}, np.array([0, 1, 2]), np.array([1, 1]), 2, 2) + assert client is None and torch.equal(inter_stage[0]["codes.audio"], req_codes[0]) + + +def test_async_chunk_output_stages_mm_on_copy_stream_before_get_output(monkeypatch) -> None: + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + calls = [] + + def copy_mm(mm_outputs, total_tokens, **ctx): + calls.append((total_tokens, ctx)) + return {"codes": {"audio": mm_outputs["codes"]["audio"].clone()}} + + monkeypatch.setattr(omni_ar_model_runner, "_async_copy_mm", copy_mm) + source_codes = torch.tensor([[7, 8]], dtype=torch.long) + input_batch = SimpleNamespace(query_start_loc_np=np.array([0, 1]), num_scheduled_tokens=[1], num_reqs=1) + output = _async_output( + multimodal_outputs={"codes": {"audio": source_codes}}, + input_batch=input_batch, + copy_stream=(copy_stream := _FakeStream()), + async_chunk=True, + ) + # Staged once on the copy stream during construction, with one resolved pin-memory context for all D2H helpers. + [(total_tokens, ctx)] = calls + assert total_tokens == 1 and ctx["copy_stream"] is copy_stream and ctx["pin_memory"] is not None + assert output._mm_snapshot["codes"]["audio"].device.type == "cpu" + source_codes.fill_(99) # a later graph replay cannot leak into the snapshot + assert torch.equal(output.get_output().inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8]])) + + +@pytest.mark.parametrize("prefill_first", [False, True]) +@pytest.mark.parametrize("padded", [False, True]) +def test_request_reference_codes_preserve_local_axis(prefill_first, padded): + lengths = np.array([255, 1] if prefill_first else [1, 255]) + offsets = np.array([0, lengths[0], 256]) + size = 512 if padded else 256 + ref = torch.arange(size * 16).reshape(size, 16) + refs = [ref, torch.empty(0)] if prefill_first else [torch.empty(0), ref] + codes = torch.arange(size * 16).reshape(size, 16) + outputs, _ = OmniARModelRunner._build_async_chunk_outputs_from_mm( + {"codes": {"audio": codes, "ref": refs}}, offsets, lengths, 2, 256, size + ) + index = 0 if prefill_first else 1 + assert torch.equal(outputs[index]["codes.ref"], ref) + assert outputs[index]["codes.ref"].data_ptr() != ref.data_ptr() + for i in range(2): + assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]]) diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py new file mode 100644 index 000000000..a1629e0bf --- /dev/null +++ b/tests/worker_v2/test_omni_data_plane.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lifecycle contracts for OmniRunnerDataPlane/NativeOutputWorker.""" + +import threading +import time +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from vllm_omni.worker_v2.delivery import DeliveryCancelledError, DeliveryState, OmniDeliveryManager +from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker +from vllm_omni.worker_v2.omni_data_plane import OmniRunnerDataPlane + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _complete(plane, payloads, token=21, req="internal"): + return plane.complete_outputs( + req_ids=[req], inter_stage_outputs=payloads, sampled_token_ids=[[token]] * len(payloads) + ) + + +def _new_request(req_id="internal", external_req_id="external"): + return SimpleNamespace( + req_id=req_id, + external_req_id=external_req_id, + prompt_token_ids=[10, 11], + num_computed_tokens=2, + resumable=True, + additional_information=SimpleNamespace(entries={}), + sampling_params=SimpleNamespace(stop_token_ids=[2150]), + ) + + +def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0): + plane = object.__new__(OmniRunnerDataPlane) + plane.__dict__.update( + _native_requests={}, + _native_output_lock=threading.RLock(), + _native_send_lock=threading.Lock(), + _native_output_error=None, + _native_output_error_lock=threading.Lock(), + _native_output_closed=False, + _native_outputs_in_flight=defaultdict(int), + _native_terminal_pending=set(), + _put_req_chunk=defaultdict(int), + _ramp_chunk_count=defaultdict(int), + _delivery_manager=OmniDeliveryManager( + delivery_timeout_s=delivery_timeout_s, shutdown_timeout_s=shutdown_timeout_s + ), + ) + return plane + + +def _make_plane(replace_send): + p = _bare_plane() + p.record = SimpleNamespace(batches=[], cleaned=[]) + if replace_send: + p.send_chunks = lambda entries, **_kw: p.record.batches.append(entries) or len(entries) + p.get_omni_connector_output, p.cleanup_finished_request = (lambda: None), p.record.cleaned.append + return p + + +@pytest.fixture +def raw_plane(): # keeps the real connector send path + p = _make_plane(replace_send=False) + yield p + if getattr(p, "_native_output_worker", None) is not None: + p._stop_output_worker() + + +@pytest.fixture +def plane(): + p = _make_plane(replace_send=True) + yield p + if getattr(p, "_native_output_worker", None) is not None: + p._stop_output_worker() + + +def _gated_connector(): + # put() blocks its first call until released (or forever without release). + conn = SimpleNamespace( + put_keys=[], put_started=threading.Event(), release=threading.Event(), close_called=threading.Event() + ) + + def put(**kwargs): + conn.put_keys.append(kwargs["put_key"]) + if len(conn.put_keys) == 1: + conn.put_started.set() + assert conn.release.wait(timeout=2) + return True, 1, None + + conn.put, conn.close = put, conn.close_called.set + return conn + + +def _start_save_thread(plane, connector): + plane.__dict__.update( + _omni_connector=connector, + _stage_id=0, + _next_stage_id=1, + _lock=threading.Lock(), + _pending_save_reqs={}, + _pending_save_counts=defaultdict(int), + _deferred_send_cleanup=set(), + _request_ids_mapping={}, + _cached_ic={}, + _send_side_request_payload={}, + _code_prompt_token_ids=defaultdict(list), + _work_available=threading.Event(), + _stop_event=threading.Event(), + _MAX_SEND_RETRIES=0, + _can_send=True, + _recv_thread=None, + _custom_process_batch_func=lambda **kwargs: kwargs["pooling_outputs"], + is_data_transfer_rank=lambda: True, + _connector_send_error_sink=plane._record_output_error, + ) + plane._save_thread = threading.Thread(target=plane._save_loop, daemon=True) + plane._save_thread.start() + + +def _stop_save_thread(plane): + plane._stop_event.set() + plane._work_available.set() + plane._save_thread.join(timeout=2) + assert not plane._save_thread.is_alive() + + +class _Output: + copy_event = None + + def __init__(self, rid="internal", gate=None, started=None, error=None, payload=None, token=1): + self.rid, self.gate, self.started = rid, gate, started + self.error, self.payload, self.token = error, payload, token + + def get_output(self): + if self.started is not None: + self.started.set() + if self.gate is not None: + assert self.gate.wait(timeout=3) + if self.error is not None: + raise self.error + return SimpleNamespace( + req_ids=[self.rid], + inter_stage_outputs=[self.payload] if self.payload is not None else None, + sampled_token_ids=[[self.token]], + ) + + +def test_emit_cohort_one_batch_and_trims_non_resumable_history(plane): + plane.send_chunk = lambda *_a, **_kw: pytest.fail("must not fall back to per-request send_chunk") + plane.register_request(_new_request("r0", "ext-0")) + plane.register_request(_new_request("r1", "ext-1")) + emitted = plane.emit_chunks( + req_ids=["r0", "r1"], + inter_stage_outputs=[{"codes.audio": "c0"}, {"codes.audio": "c1"}], + sampled_token_ids=[[21]], + terminal_req_ids={"r1"}, + ) + assert emitted == 2 and len(plane.record.batches) == 1 # the whole step cohort goes out in one batch + (req0, _), (req1, _) = plane.record.batches[0] + assert ( + req0.all_token_ids == [10, 11, 21] + and not req0.is_finished() + and req1.is_finished() + and plane.record.cleaned == ["r1"] + ) + # A non-resumable request omits token history after its first chunk. + plane.register_request(_new_request("r2", "ext-2")) + state = plane._native_requests["r2"] + state.resumable, state.output_token_ids, plane._put_req_chunk["ext-2"] = False, list(range(1024)), 1 + plane.emit_chunks( + req_ids=["r2"], inter_stage_outputs=[{"h": "d"}], sampled_token_ids=[[1024]], terminal_req_ids={"r2"} + ) + request, _ = plane.record.batches[1][0] + assert request.output_token_count == 1025 and request.all_token_ids == [] + + +def test_terminal_reserved_abort_and_cleanup_lifecycle(plane): + batches, cleaned = plane.record.batches, plane.record.cleaned + plane.register_request(_new_request()) + plane.reserve_outputs(["internal"]) + plane.reserve_outputs(["internal"]) + # The terminal defers while reservations are in flight, and is idempotent. + assert plane.request_terminal({"internal"}) == 0 + assert plane.request_terminal({"internal"}) == 0 and batches == [] + assert _complete(plane, [{"codes.audio": "c0"}]) == 1 and not batches[0][0][0].is_finished() + assert _complete(plane, [{"codes.audio": "c1"}], token=22) == 2 # data chunk, then the released terminal + chunk, terminal = batches[1][0], batches[2][0] + assert chunk[1] == {"codes": {"audio": "c1"}} and chunk[0].output_token_ids == [21, 22] + assert terminal[1] is None and terminal[0].is_finished() and cleaned == ["internal"] + # Stale terminal/complete after cleanup are no-ops; cleanup ran exactly once. + assert plane.request_terminal({"internal"}) == 0 and _complete(plane, [{"codes.audio": "stale"}]) == 0 + # Abort emits exactly one terminal and silently drops the stale deferred output. + plane.register_request(_new_request("ra")) + plane.reserve_outputs(["ra"]) + plane.request_terminal({"ra"}) + assert plane.abort_requests({"ra"}) == 1 + request, payload = batches[-1][0] + assert request.is_finished() and payload is None and cleaned == ["internal", "ra"] + assert ( + plane.abort_requests({"ra"}) == 0 and _complete(plane, [{"s": 1}], req="ra") == 0 + ) # duplicate abort, stale output + assert sum(b[0][0].request_id == "ra" for b in batches) == 1 # terminal exactly once + + +def test_abort_cannot_overtake_committed_output(plane): + in_lock, release, sent = threading.Event(), threading.Event(), threading.Event() + + def gated(entries, _orig=plane._send_chunk_entries, **kw): + terminal = entries[0][0].is_finished() + if not terminal: + in_lock.set() + assert release.wait(timeout=2) # hold the committed data send inside _native_send_lock + result = _orig(entries, **kw) + if terminal: + sent.set() + return result + + plane._send_chunk_entries = gated + plane.register_request(_new_request()) + plane.reserve_outputs(["internal"]) + threads = [ + threading.Thread(target=lambda: _complete(plane, [{"c": 0}])), + threading.Thread(target=plane.abort_requests, args=({"internal"},)), + ] + threads[0].start() + try: + assert in_lock.wait(timeout=2) + threads[1].start() + overtook = sent.wait(timeout=0.1) # the terminal must not pass the committed data send + finally: + release.set() + for t in threads: + if t.ident is not None: + t.join(timeout=2) + assert not t.is_alive() + assert not overtook + + +def test_output_worker_fifo_and_scheduler_not_blocked(plane): + batches, send_started, release_send = [], threading.Event(), threading.Event() + + def send_chunks(entries, **_kw): + if not entries[0][0].is_finished(): + send_started.set() + assert release_send.wait(timeout=2) + batches.append(entries) + return len(entries) + + plane.send_chunks = send_chunks + plane._start_output_worker(max_pending_batches=2) + plane.register_request(_new_request()) + plane.reserve_outputs(["internal"]) + plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]]) + assert send_started.wait(timeout=2) + # A stuck deferred send must not block the scheduler path. + assert plane.request_terminal({"internal"}) == 0 and batches == [] + release_send.set() + plane.drain_outputs() + assert [e[0].is_finished() for batch in batches for e in batch] == [False, True] + + +def test_worker_failure_surfaces_and_close_drains_in_order(plane): + events = [] + plane.send_chunks = lambda _e, **_kw: (_ for _ in ()).throw(RuntimeError("enqueue failed")) + plane.shutdown_omni_connectors = lambda: events.append("shutdown") + plane._start_output_worker(max_pending_batches=2) + plane.register_request(_new_request()) + plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]]) + with pytest.raises(RuntimeError, match="enqueue failed"): + plane.drain_outputs() + with pytest.raises(RuntimeError, match="enqueue failed"): # an errored worker rejects new submissions + plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 1}], sampled_token_ids=[[22]]) + # close() drains pending batches before shutting connectors down. + plane._native_output_error = None + + def send_chunks(entries, **_kw): + events.append("send") + return len(entries) + + plane.send_chunks = send_chunks + plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 1}], sampled_token_ids=[[22]]) + plane.close() + assert events == ["send", "shutdown"] + + +def test_permanent_put_failure_quarantines_and_holds_state(raw_plane): + puts = [] + + def put(**_kw): + puts.append(1) + return False, 0, None + + connector = SimpleNamespace(put=put, close=lambda: None) + _start_save_thread(raw_plane, connector) + raw_plane.register_request(_new_request()) + raw_plane.reserve_outputs(["internal"]) + raw_plane.request_terminal({"internal"}) + with pytest.raises(RuntimeError, match="connector send failed"): + raw_plane.complete_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]]) + assert len(puts) == 1 and raw_plane._native_output_error is not None # fail fast, no retry + assert raw_plane._native_outputs_in_flight["internal"] == 1 and "internal" in raw_plane._native_terminal_pending + with pytest.raises(RuntimeError, match="quarantined"): # the quarantine rejects new tickets + raw_plane._delivery_manager.create_ticket(request_id="other", put_key="other_0_0") + _stop_save_thread(raw_plane) + + +def test_delivery_timeout_drops_queued_ticket(raw_plane): + raw_plane._delivery_manager = OmniDeliveryManager(delivery_timeout_s=0.05, shutdown_timeout_s=0.5) + connector = _gated_connector() + _start_save_thread(raw_plane, connector) + req = SimpleNamespace(request_id="external") + _, first = raw_plane._enqueue_chunk_payload(req, {"c": 1}, wait_for_delivery=True) + _, second = raw_plane._enqueue_chunk_payload(req, {"c": 2}, wait_for_delivery=True) + assert connector.put_started.wait(timeout=1) + with pytest.raises(TimeoutError, match="delivery timed out"): + first.wait() + assert raw_plane._delivery_manager.is_quarantined + connector.release.set() + deadline = time.monotonic() + 1 + while time.monotonic() < deadline and "external" in raw_plane._pending_save_counts: # the save loop drops it + time.sleep(0.001) + # After quarantine the queued ticket is never put and is marked FAILED. + assert connector.put_keys == [first.put_key] and second.state is DeliveryState.FAILED + assert "external" not in raw_plane._pending_save_counts and "external" not in raw_plane._pending_save_reqs + _stop_save_thread(raw_plane) + + +def test_close_cancels_waiting_delivery_once_and_shutdown_stays_bounded(): + # A queued delivery waiter is cancelled exactly once on shutdown. + manager = OmniDeliveryManager(delivery_timeout_s=30.0, shutdown_timeout_s=0.5) + ticket = manager.create_ticket(request_id="r", put_key="k") + manager.shutdown(RuntimeError("close")) + manager.shutdown(RuntimeError("again")) # idempotent: cancelled at most once + with pytest.raises(DeliveryCancelledError, match="cancelled"): + ticket.wait() + # A connector put() that cannot be cancelled must not hang shutdown. + plane = _bare_plane(delivery_timeout_s=0.05, shutdown_timeout_s=0.2) + connector = _gated_connector() + _start_save_thread(plane, connector) + assert plane._enqueue_chunk_payload(SimpleNamespace(request_id="external"), {"c": 0}, wait_for_delivery=True)[0] + assert connector.put_started.wait(timeout=1) + start = time.monotonic() + with pytest.raises(RuntimeError, match="shutdown exceeded"): + plane.shutdown_omni_connectors() + assert time.monotonic() - start < 1.5 and connector.close_called.is_set() + connector.release.set() + _stop_save_thread(plane) + + +@pytest.mark.parametrize("abort", [False, True]) +def test_materialization_fences_terminal_and_abort(plane, abort): + batches, cleaned, gate, started = plane.record.batches, plane.record.cleaned, threading.Event(), threading.Event() + plane._start_output_worker(max_pending_batches=2) + materializer = NativeOutputWorker(2) + try: + plane.register_request(_new_request()) + plane.reserve_outputs(["internal"]) + output = materializer.submit( + _Output(gate=gate, started=started, payload={"codes.audio": "c0"}, token=21), plane + ) + assert started.wait(timeout=3) + assert plane.request_terminal({"internal"}) == 0 # fenced behind the in-flight materialization + assert not abort or plane.abort_requests({"internal"}) == 1 + gate.set() + output.get_output() + materializer.close() + plane.drain_outputs() + assert cleaned == ["internal"] and len(batches) == (1 if abort else 2) + assert abort or (batches[0][0][1] == {"codes": {"audio": "c0"}} and not batches[0][0][0].is_finished()) + assert batches[-1][0][0].is_finished() + finally: + gate.set() + materializer.close() + + +def test_native_worker_fifo_thread_affinity_and_signals(): + owner, gate = threading.get_ident(), threading.Event() + calls: list[tuple[Any, ...]] = [] + + def get_signals(): + calls.append(("signals", threading.get_ident())) + return "ready" + + plane = SimpleNamespace( + enqueue_outputs=lambda **kw: calls.append(("publish", kw["req_ids"], threading.get_ident())), + get_omni_connector_output=get_signals, + ) + worker = NativeOutputWorker(2) + try: + a = worker.submit(_Output("a", gate=gate, payload={"v": 1}), plane) + b = worker.submit(_Output("b", payload={"v": 1}), plane) + gate.set() + b.get_output() # publication of b implies a already published (FIFO, single worker thread) + publishes = [c for c in calls if c[0] == "publish"] + assert [c[1] for c in publishes] == [["a"], ["b"]] + assert all(c[2] != owner for c in publishes) # published on the worker thread + result = a.get_output() + assert a.get_output() is result and result.inter_stage_outputs is None # resolution is cached + assert result.omni_connector_output == "ready" + assert calls[-2:] == [("signals", owner)] * 2 # signals resolve lazily on the caller thread + finally: + gate.set() + worker.close() + + +def test_native_worker_capacity_error_cache_and_tp_gates(): + from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + + plane = SimpleNamespace(enqueue_outputs=lambda **kw: None, get_omni_connector_output=lambda: None) + worker, gate = NativeOutputWorker(1), threading.Event() + try: + first = worker.submit(_Output("a", gate=gate), plane) + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(worker.submit, _Output("b"), plane) + time.sleep(0.05) + assert not future.done() # capacity 1: pending materialization is bounded + gate.set() + second = future.result(3) + assert first.get_output().req_ids == ["a"] and second.get_output().req_ids == ["b"] + bad = worker.submit(_Output("x", error=ValueError("crash")), plane) + finally: + gate.set() + worker.close() + # A failed materialization is cached: repeated reads raise the same error. + for _ in range(2): + with pytest.raises(ValueError, match="crash"): + bad.get_output() + runner = SimpleNamespace( + _omni_data_plane=object(), vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2)) + ) + assert OmniGPUModelRunner._uses_native_output_materializer(runner) is False # TP>1 keeps TP consumers + runner.vllm_config.parallel_config.tensor_parallel_size = 1 + assert OmniGPUModelRunner._uses_native_output_materializer(runner) is True + + +def test_abort_before_first_chunk_cleans_receiver_state(): + plane = _bare_plane() + plane._stage_id, plane._lock, plane._work_available = 1, threading.Lock(), threading.Event() + for name in ( + "_pending_full_payload_send", + "_request_ids_mapping", + "_pending_save_counts", + "_send_side_request_payload", + "_code_prompt_token_ids", + "_cached_ic", + "_adaptive_states", + "_kv_pending_transfers", + "_get_req_chunk", + "_pending_load_reqs", + "_local_stage_payload_cache", + "_local_request_metadata", + ): + setattr(plane, name, {}) + for name in ( + "_deferred_send_cleanup", + "_kv_active_transfers", + "_kv_completed_transfers", + "_kv_triggered_requests", + "_finished_load_reqs", + "_chunk_ready_req_ids", + "_chunk_finished_req_ids", + "_chunk_stream_completed", + "_stage_recv_req_ids", + "_full_payload_pending_broadcast_req_ids", + "_async_chunk_updated_req_ids", + ): + setattr(plane, name, set()) + plane.register_receivers([SimpleNamespace(request_id="r", external_req_id="external")]) + plane._local_stage_payload_cache["r"] = {"codes": torch.ones(1)} + assert "r" in plane._pending_load_reqs and plane._request_ids_mapping["r"] == "external" + assert plane.abort_requests({"r"}) == 0 + assert plane.abort_requests({"r"}) == 0 + assert not plane._pending_load_reqs and not plane._request_ids_mapping and not plane._local_stage_payload_cache diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py new file mode 100644 index 000000000..63c480291 --- /dev/null +++ b/tests/worker_v2/test_omni_generation_model_runner.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniGenerationModelRunner contracts: empty-step lifecycle, output partition, +CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeStagedField: + def __init__(self, data): + self.np = data + self._staged = [] + + def stage_write_elem(self, idx, value): + self._staged.append((idx, value)) + + def apply_write(self): + for idx, value in self._staged: + self.np[idx] = value + self._staged.clear() + + +def _make_runner(model_output, num_reqs=1, prompt_len=10): + runner = object.__new__(OmniGenerationModelRunner) + runner.device = torch.device("cpu") + runner.model_config = MagicMock(hf_text_config=None) + del runner.model_config.eos_token_id + runner.postprocess = lambda *a, **kw: None + runner._gen_model_output = model_output + runner._gen_input_batch = SimpleNamespace( + num_reqs=num_reqs, req_ids=[f"req-{i}" for i in range(num_reqs)], idx_mapping_np=np.arange(num_reqs) + ) + runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None) + runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None)) + runner.check_ep_fault = False + runner.req_states = MagicMock( + prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)), + num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)), + ) + return runner + + +def test_control_only_step_keeps_lifecycle_and_skips_input_construction(): + runner = object.__new__(OmniGenerationModelRunner) + order = [] + for name in ( + "_prepare_native_data_plane", + "finish_requests", + "free_states", + "_apply_block_table_staged_writes_if_available", + ): + setattr(runner, name, lambda *args, name=name: order.append(name)) + for name in ("_handle_async_chunk_updates", "add_requests", "update_requests", "_sync_native_data_plane_payloads"): + setattr(runner, name, MagicMock(side_effect=AssertionError("control step built model inputs"))) + output = object() + runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output) + runner._merge_ec_connector_no_forward = lambda _s, value: value + runner._attach_native_data_plane_signals = lambda value: value + scheduler_output = SimpleNamespace( + total_num_scheduled_tokens=0, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace(req_ids=[]) + ) + assert runner.execute_model(scheduler_output) is output + assert order == [ + "_prepare_native_data_plane", + "finish_requests", + "free_states", + "_apply_block_table_staged_writes_if_available", + ] + + +def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state(): + from vllm.v1.core.sched.output import SchedulerOutput + + from vllm_omni.core.sched.output import OmniCachedRequestData + + runner = object.__new__(OmniGenerationModelRunner) + runner.model = SimpleNamespace(logits_processor=None) + runner._remove_request = MagicMock() + req_id_to_index = MagicMock() + req_id_to_index.get.return_value = None + runner.req_states = SimpleNamespace( + add_request=MagicMock(), req_id_to_index=req_id_to_index, apply_staged_writes=MagicMock() + ) + runner.pooling_runner = None + runner.encoder_cache = None + runner.model_state = SimpleNamespace( + add_request=MagicMock(), + apply_staged_writes=MagicMock(), + remove_request=MagicMock(), + intermediate_buffer=SimpleNamespace(remove_request=MagicMock()), + ) + runner.block_tables = SimpleNamespace(append_block_ids=MagicMock()) + runner.lora_state = SimpleNamespace(add_request=MagicMock()) + runner.is_last_pp_rank = False + runner.adaptive_verification = None + runner.sampler = None + cached = OmniCachedRequestData( + req_ids=["req"], + resumed_req_ids=set(), + new_token_ids=[[]], + all_token_ids={"req": [1]}, + new_block_ids=[()], + num_computed_tokens=[0], + num_output_tokens=[0], + prompt_token_ids={"req": [1]}, + additional_information={"req": None}, + ) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req": 1}, + total_num_scheduled_tokens=1, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + runner._handle_async_chunk_updates(scheduler_output) + added = runner.model_state.add_request.call_args.args[1] + assert (added.req_id, added.prompt_token_ids) == ("req", [1]) + + +def test_sample_tokens_cpu_sync_owns_outputs_and_runs_connector_last(): + waveform = torch.arange(4, dtype=torch.float32) + output = OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={"codes": {"audio": [waveform]}}) + runner = _make_runner(output, num_reqs=1) + runner._reserve_native_data_plane_outputs = MagicMock() + runner._finalize_native_data_plane_output = MagicMock(side_effect=lambda value: value) + + def post_forward(_finished): + assert runner.req_states.num_computed_tokens.np.tolist() == [10] + + runner.kv_connector.post_forward.side_effect = post_forward + result = OmniGenerationModelRunner.sample_tokens(runner) + assert isinstance(result, OmniModelRunnerOutput) + # Reserve precedes finalize; connector runs after token state is applied. + runner._reserve_native_data_plane_outputs.assert_called_once_with(["req-0"]) + runner.kv_connector.post_forward.assert_called_once_with({"finished"}) + waveform.fill_(-1) + assert torch.equal(result.multimodal_outputs[0]["codes.audio"], torch.arange(4, dtype=torch.float32)) + + +def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypatch): + from vllm_omni.worker_v2 import omni_generation_model_runner as generation_runner + + output = OmniOutput(text_hidden_states=torch.zeros(1), multimodal_outputs={"model_outputs": [torch.randn(4)]}) + runner = _make_runner(output, num_reqs=1) + runner.device = SimpleNamespace(type="cuda") + runner.main_stream = object() + runner.output_copy_stream = object() + runner.model_config.async_chunk = True + runner._release_generation_slots = MagicMock() + runner._finalize_native_data_plane_output = MagicMock() + monkeypatch.setattr(generation_runner, "_contains_cuda_tensor", lambda _: True) + captured = {} + + class _FakeAsyncOutput: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(generation_runner, "OmniGenerationAsyncOutput", _FakeAsyncOutput) + input_batch = runner._gen_input_batch + result = generation_runner.OmniGenerationModelRunner.sample_tokens(runner) + + assert isinstance(result, _FakeAsyncOutput) + assert captured["finalize_output"] is runner._finalize_native_data_plane_output + assert captured["model_runner_output"].sampled_token_ids == [[]] + runner._release_generation_slots.assert_called_once() + output_req_ids = captured["model_runner_output"].req_ids + input_batch.req_ids[0] = "reused" # snapshot must insulate the published output + assert output_req_ids == ["req-0"] diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py new file mode 100644 index 000000000..77cd6ee98 --- /dev/null +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MRV2 admission, capture, dispatch and request lifecycle contracts.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from vllm import SamplingParams +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.worker.gpu.model_runner import GPUModelRunner +from vllm.v1.worker.gpu.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.model_states import init_omni_model_state +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_runner(): + """Create an OmniGPUModelRunner without calling __init__.""" + runner = object.__new__(OmniGPUModelRunner) + runner.model = MagicMock() + runner.req_states = SimpleNamespace(req_id_to_index={"r1": 0, "r2": 1}) + runner.execute_model_state = None + return runner + + +def test_add_requests_empty_admission_and_stop_id_sanitization(): + runner = _make_runner() + runner.sampler = Sampler.__new__(Sampler) + with patch.object(GPUModelRunner, "add_requests", return_value=None) as parent: + runner.add_requests(SchedulerOutput.make_empty()) + parent.assert_not_called() + + # New requests always go upstream; narrow logits heads sanitize stop ids. + sampling_params = SamplingParams(min_tokens=2, stop_token_ids=[2150]) + sampling_params.update_from_generation_config({}, 151645) + runner.model = SimpleNamespace(logits_processor=SimpleNamespace(vocab_size=3072)) + output = SchedulerOutput.make_empty() + output.scheduled_new_reqs = [SimpleNamespace(sampling_params=sampling_params)] + runner.add_requests(output) + parent.assert_called_once_with(output) + assert sampling_params.all_stop_token_ids == {2150} + assert sampling_params.eos_token_id == 151645 + + +def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(): + runner = _make_runner() + plane = SimpleNamespace( + register_request=MagicMock(), + register_receivers=MagicMock(), + request_terminal=MagicMock(), + abort_requests=MagicMock(), + ) + runner._omni_data_plane = plane + new_req = SimpleNamespace(req_id="r1") + warmup_req = SimpleNamespace(req_id="_warmup_0_") + handle = SimpleNamespace(request_id="r2", external_req_id="ext-r2") + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[new_req, warmup_req], + pending_input_registrations=[handle], + data_plane_terminal_req_ids={"r0"}, + finished_req_ids={"r0", "aborted"}, + ) + + runner._prepare_native_data_plane(scheduler_output) + + plane.register_request.assert_called_once_with(new_req) # warmup skipped + plane.register_receivers.assert_called_once_with([handle]) + plane.request_terminal.assert_called_once_with({"r0"}) + plane.abort_requests.assert_called_once_with({"aborted"}) + + +@pytest.mark.parametrize("output_form", ["tuple", "omni"]) +def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): + runner = object.__new__(OmniGPUModelRunner) + hidden = torch.ones(1, 2) + + def original_forward(): + if output_form == "tuple": + return hidden, {"layers": {}} + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + runner.model = SimpleNamespace(forward=original_forward) + runner._model_returns_tuple = True + runner._exclude_full_graph = True + runner.use_aux_hidden_state_outputs = False + piecewise = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE) + full = SimpleNamespace(cg_mode=CUDAGraphMode.FULL) + runner.cudagraph_manager = SimpleNamespace( + _capture_descs={CUDAGraphMode.PIECEWISE: [piecewise], CUDAGraphMode.FULL: [full]}, + _candidates={(1, 0): [piecewise, full]}, + ) + runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock()) + runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc") + + def assert_unwrapped(_self): + assert torch.equal(runner.model.forward(), hidden) # forward unwrapped during capture + return 3 + + with patch.object(GPUModelRunner, "capture_model", assert_unwrapped): + assert runner.capture_model() == 3 + + assert runner.model.forward is original_forward # restored after capture + assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]} + runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor) + + +@pytest.mark.parametrize( + "parallel_config,match", + [ + (dict(pipeline_parallel_size=2, prefill_context_parallel_size=1), "pipeline parallel"), + (dict(pipeline_parallel_size=1, prefill_context_parallel_size=2), "prefill context parallelism"), + ], +) +def test_mrv2_rejects_parallel_modes_at_startup(parallel_config, match): + runner = object.__new__(OmniGPUModelRunner) + runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(**parallel_config)) + with pytest.raises(NotImplementedError, match=match): + runner._validate_parallel_support() + + +@pytest.mark.parametrize( + "flag,expect_omni", + [(None, False), ("has_preprocess", True), ("has_postprocess", True), ("have_multimodal_outputs", True)], +) +def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect_omni): + upstream = MagicMock(return_value=object()) + monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", upstream) + monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None) + cfg = SimpleNamespace(model_config=SimpleNamespace(architectures=["CustomModel"])) + model = SimpleNamespace(**({flag: True} if flag else {})) + state = init_omni_model_state(cfg, model, None, torch.device("cpu")) + + if expect_omni: + assert isinstance(state, OmniModelState) + upstream.assert_not_called() + else: + upstream.assert_called_once() + assert state is upstream.return_value + + +def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch): + runner = _make_runner() + calls = [] + runner.model = SimpleNamespace(on_requests_finished=lambda ids: calls.append(set(ids))) + runner.model_state = MagicMock() + monkeypatch.setattr(GPUModelRunner, "finish_requests", lambda *args: None) + runner.finish_requests(SimpleNamespace(finished_req_ids={"released"}, preempted_req_ids={"r1"})) + assert calls == [{"released"}] + assert sorted(c.args[0] for c in runner.model_state.remove_request.call_args_list) == [0] + + +@pytest.mark.parametrize("stage,declared", [("thinker", False), ("custom_ar", True)]) +def test_capture_contract_uses_model_declaration(stage, declared): + runner = _make_runner() + runner.model = SimpleNamespace(model_stage=stage, _returns_tuple=declared) + runner._configure_cudagraph_output_contract() + assert runner._model_returns_tuple is declared + assert runner._exclude_full_graph is declared diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py new file mode 100644 index 000000000..34d0d8010 --- /dev/null +++ b/tests/worker_v2/test_omni_model_state.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniModelState: mixed-batch reorder, state isolation/slot reuse, MTP +graph/eager paths with per-request seed independence, async snapshot ownership.""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from vllm.v1.worker.gpu.model_states.default import DefaultModelState + +from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyInputBatch: + input_ids: SimpleNamespace + query_start_loc: torch.Tensor + + def __init__(self, indices, *, num_computed_tokens_cpu=None): + self.idx_mapping_np = indices + self.num_reqs = len(indices) + self.num_scheduled_tokens = [1] * len(indices) + self.query_start_loc_np = list(range(len(indices))) + if num_computed_tokens_cpu is not None: + self.num_computed_tokens_cpu = np.array(num_computed_tokens_cpu, dtype=np.int32) + + +def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False): + state = object.__new__(OmniModelState) + model = MagicMock() + model.has_preprocess = has_preprocess + model.has_postprocess = has_postprocess + model.have_multimodal_outputs = have_multimodal_outputs + model.gpu_resident_buffer_keys = set() + model.batched_gpu_staging_keys = set() + model.mtp_accepts_per_row_generators = False + model.mtp_accepts_req_infos = False + model.mtp_output_key = ("codes", "audio") + model.mtp_validity_key = None + model.mtp_sampling_params = {} + model.get_mtp_seed = lambda params: (getattr(params, "extra_args", None) or {}).get("test_seed") + model.preprocess_batch_mrv2 = None + model.preprocess_decode_batch_mrv2 = None + model.preprocess_decode_batch = None + model.postprocess_batch_mrv2 = None + state.model = model + state.scheduler_config = SimpleNamespace(max_num_seqs=max_num_reqs) + state.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={})) + state.has_preprocess = has_preprocess + state.has_postprocess = has_postprocess + state.have_multimodal_outputs = have_multimodal_outputs + + from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer + + state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) + state._static_inputs_embeds = None + state._mtp_generators = {} + state._mtp_runner = None + for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"): + setattr(state, name, None) + return state + + +def _add(state, req_id="r1", idx=0, **req_kwargs): + with patch.object(DefaultModelState, "add_request", return_value=None): + state.add_request(idx, SimpleNamespace(req_id=req_id, mm_features=[], **req_kwargs)) + + +def _fill_buffers(state, *req_ids): + for idx, req_id in enumerate(req_ids): + state.intermediate_buffer.buffers[idx] = {"req_id": req_id} + + +def _seeded(seed): + return SimpleNamespace(extra_args={"test_seed": seed}, seed=None) + + +def _init_static(state, bsz, dim=3): + state._mtp_input_ids = torch.zeros(bsz, dtype=torch.long) + state._mtp_input_embeds = torch.zeros((bsz, dim)) + state._mtp_hidden = torch.zeros((bsz, dim)) + state._mtp_text_step = torch.zeros((bsz, dim)) + state._mtp_offsets = torch.zeros(bsz, dtype=torch.long) + + +def _mtp_batches(dim=3): + return [(0, 0, (torch.ones(dim), torch.ones(dim) * 2)), (1, 1, (torch.ones(dim) * 3, torch.ones(dim) * 4))] + + +@pytest.fixture(autouse=True) +def _fwd_ctx(): + with patch("vllm.forward_context.set_forward_context", return_value=nullcontext()): + yield + + +@pytest.mark.parametrize("req_id,expect_validity", [("_warmup_0_", True), ("r1", False)]) +def test_add_request_declared_validity_only_for_warmup(req_id, expect_validity): + # Real requests must not be masked with a fabricated validity key. + state = _make_state() + state.model.mtp_validity_key = ("meta", "codec_frame_valid") + _add(state, req_id, 0) + buf = state.intermediate_buffer.buffers[0] + assert ("meta" in buf) is expect_validity + if expect_validity: + assert buf["meta"]["codec_frame_valid"].item() is False + + +@pytest.mark.parametrize("mode", ["index", "req_id", "unknown"]) +def test_remove_request_state_isolation(mode): + state = _make_state() + _add(state, "r1", 0) + state._mtp_generators["r1"] = torch.Generator(device="cpu") + if mode == "unknown": + state.remove_request("missing") + assert state.intermediate_buffer.buffers[0]["req_id"] == "r1" + assert "r1" in state._mtp_generators + return + state.remove_request(0 if mode == "index" else "r1") + assert state.intermediate_buffer.buffers[0] == {} + assert "r1" not in state._mtp_generators # recycled slot inherits no seed stream + + +def test_output_spans_follow_reordered_mixed_batch(): + # batch=[2, 0]: prefill (3 tokens) reordered ahead of decode (1 token). + state = _make_state(have_multimodal_outputs=True) + batch = _DummyInputBatch([2, 0]) + batch.num_scheduled_tokens = [3, 1] + batch.query_start_loc_np = [0, 3] + state.intermediate_buffer.buffers[2] = {"req_id": "prefill"} + state.intermediate_buffer.buffers[0] = {"req_id": "decode"} + seen = {} + state.model.make_omni_output = lambda hidden, **kwargs: ( + seen.update(kwargs) or OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + ) + state.postprocess_model_output(torch.zeros(4, 2), batch, SimpleNamespace()) + assert seen["request_token_spans"] == [(0, 3), (3, 4)] + assert [info["req_id"] for info in seen["model_intermediate_buffer"]] == ["prefill", "decode"] + + +@pytest.mark.parametrize("hook_name", ["preprocess_decode_batch_mrv2", "preprocess_decode_batch", "legacy"]) +def test_decode_cohort_uses_batch_hook(hook_name): + state = _make_state(max_num_reqs=2, has_preprocess=True) + state.intermediate_buffer.buffers[0] = {"req_id": "r1", "meta": {"step": 1}} + state.intermediate_buffer.buffers[1] = {"req_id": "r2", "meta": {"step": 2}} + batch_calls = [] + + def preprocess_decode_batch(*, input_ids, input_embeds, req_infos): + batch_calls.append(list(req_infos)) + return ( + input_ids + 10, + input_embeds + 20, + torch.tensor([[31.0, 32.0], [41.0, 42.0]]), + torch.tensor([[51.0, 52.0], [61.0, 62.0]]), + [{"meta": {"step": 11}}, {"meta": {"step": 12}}], + ) + + if hook_name == "legacy": + # Legacy V1 hook without input_embeds is wrapped by the adapter. + state.model.preprocess_decode_batch = lambda *, input_ids, req_infos: preprocess_decode_batch( + input_ids=input_ids, input_embeds=model_inputs["inputs_embeds"], req_infos=req_infos + ) + else: + setattr(state.model, hook_name, preprocess_decode_batch) + if hook_name == "preprocess_decode_batch_mrv2": + state.model.preprocess_decode_batch = MagicMock( + side_effect=AssertionError("MRV2 hook must take precedence") + ) + state.model.preprocess = MagicMock(side_effect=AssertionError("decode requests must use the batch hook")) + seen_prepacked = [] + + def capture(batches, *_a, prepacked_mtp_inputs=None): + seen_prepacked.append(prepacked_mtp_inputs) + + state._run_batched_mtp = capture + model_inputs = { + "input_ids": torch.tensor([101, 202], dtype=torch.long), + "inputs_embeds": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + } + state.run_preprocess( + _DummyInputBatch([0, 1], num_computed_tokens_cpu=[1, 1]), + model_inputs, + SimpleNamespace(prompt_len=np.array([1, 1], dtype=np.int32)), + ) + state.model.preprocess.assert_not_called() + assert [info["req_id"] for info in batch_calls[0]] == ["r1", "r2"] + assert torch.equal(model_inputs["input_ids"], torch.tensor([111, 212])) + assert state.intermediate_buffer.buffers[0]["meta"]["step"] == 11 + assert torch.equal(seen_prepacked[0][0], torch.tensor([[31.0, 32.0], [41.0, 42.0]])) + + +def test_static_decode_embeddings_refresh_from_input_ids(): + # FULL-graph replay reads the static buffer: it must hold fresh embeddings. + state = _make_state(has_preprocess=True) + static_embeds = torch.zeros(1, 4) + state._static_inputs_embeds = static_embeds + expected = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + state.model.embed_input_ids = lambda input_ids: expected.to(input_ids.device) + state.model.preprocess = lambda input_ids, input_embeds, **_info: (input_ids, input_embeds, {}) + state.intermediate_buffer.buffers[0] = {"req_id": "r1"} + state.run_preprocess( + _DummyInputBatch([0]), + {"input_ids": torch.tensor([1049], dtype=torch.long), "inputs_embeds": static_embeds[:1]}, + ) + assert torch.equal(static_embeds, expected) + original = torch.tensor([[1.0, 2.0]]) + assert OmniModelState._preprocess_result_needs_writeback(original, original) is False + assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True + + +@pytest.mark.parametrize("owned", [False, True]) +def test_batched_postprocess_gpu_snapshot_writeback(owned): + # batch=[1, 0]: last-token indices follow the reordered batch, the scalar + # hook is not called, and both rows share one snapshot storage so a later + # graph replay cannot overwrite published rows. Owned producer batches + # skip the extra snapshot while borrowed graph output still gets one. + state = _make_state(max_num_reqs=2, has_postprocess=True) + state.model.gpu_resident_buffer_keys = {("hidden_states", "last")} + _fill_buffers(state, "r0", "r1") + state.model.postprocess = MagicMock(side_effect=AssertionError("batch hook must replace scalar postprocess")) + calls = [] + produced = [] + + def postprocess_batch(*, hidden_states, last_token_indices): + calls.append(last_token_indices.clone()) + gathered = hidden_states.index_select(0, last_token_indices) + produced.append(gathered) + if owned: + return (("hidden_states", "last"), OwnedBatchTensor(gathered)) + return (("hidden_states", "last"), gathered) + + state.model.postprocess_batch_mrv2 = postprocess_batch + batch = _DummyInputBatch([1, 0]) + batch.num_scheduled_tokens = [2, 2] + batch.query_start_loc_np = [0, 2] + batch.query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]) + state.run_postprocess(hidden, batch) + + state.model.postprocess.assert_not_called() + assert torch.equal(calls[0], torch.tensor([1, 3], dtype=torch.int32)) + first = state.intermediate_buffer.buffers[0]["hidden_states"]["last"] + second = state.intermediate_buffer.buffers[1]["hidden_states"]["last"] + assert torch.equal(first, hidden[3]) + assert torch.equal(second, hidden[1]) + assert first.untyped_storage().data_ptr() == second.untyped_storage().data_ptr() + if owned: + # Owned output: rows view into the producer batch, no extra snapshot. + assert first.untyped_storage().data_ptr() == produced[0].untyped_storage().data_ptr() + else: + # Borrowed/graph output: the snapshot owns its storage, not the producer. + assert first.untyped_storage().data_ptr() != produced[0].untyped_storage().data_ptr() + + +def test_seed_independence_resolve_once_and_sampling_kwargs(): + state = _make_state(max_num_reqs=2) + # 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 + # Same model-local seed reproduces identical uniforms regardless of batch makeup. + state._mtp_sample_uniforms = torch.empty((2, 2, 4)) + assert torch.equal( + state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1).clone(), + state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1), + ) + # Per-request generators are resolved once per step and ride sampling kwargs. + state.model.mtp_sampling_params = {"do_sample": True, "temperature": 0.7} + state.model.mtp_accepts_per_row_generators = True + _init_static(state, 2) + state.intermediate_buffer.buffers[0] = {"req_id": "r0", "sampling_params": _seeded(11)} + state.intermediate_buffer.buffers[1] = {"req_id": "r1", "sampling_params": _seeded(22)} + resolver = MagicMock(wraps=state._get_mtp_generator) + state._get_mtp_generator = resolver + + def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs): + assert (kwargs["do_sample"], kwargs["temperature"]) == (True, 0.7) + assert all(isinstance(g, torch.Generator) for g in kwargs["generators"]) + return input_embeds + 10, torch.tensor([[1, 2, 3], [4, 5, 6]]) + + state.model.mtp = mtp + state._run_batched_mtp( + _mtp_batches(), torch.tensor([101, 202]), torch.zeros(2, 3), _DummyInputBatch([0, 1]), {("codes", "audio")} + ) + assert resolver.call_count == 2 + assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[1, 2, 3]])) + assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[4, 5, 6]])) + + +def test_seeded_mtp_bypasses_outer_graph_runner(): + state = _make_state(max_num_reqs=2) + raw_calls = [] + + def mtp(*args, **kwargs): + raw_calls.append(kwargs) + return args[1], args[0].reshape(-1, 1) + + state.model.mtp = mtp + state._mtp_runner = MagicMock(side_effect=AssertionError("seeded sampling must not replay the graph")) + generators = [torch.Generator().manual_seed(11), torch.Generator().manual_seed(22)] + state._call_mtp_runner( + torch.tensor([101, 202]), torch.zeros(2, 3), torch.zeros(2, 3), torch.zeros(2, 3), generators=generators + ) + assert raw_calls == [{"generators": generators}] + state._mtp_runner.assert_not_called() + + +def test_run_batched_mtp_uses_dispatched_graph_descriptor(): + state = _make_state(max_num_reqs=4) + + class _FakeGraphRunner: + def __call__(self, input_ids, input_embeds, last_hidden, text_step, **kwargs): + assert input_ids.shape[0] == 4 # graph runs the padded descriptor size + return input_embeds + 2, torch.arange(12, dtype=torch.long).reshape(4, 3) + + state._mtp_runner = _FakeGraphRunner() + _init_static(state, 4) + _fill_buffers(state, "r0", "r1") + graph_desc = SimpleNamespace(cg_mode="FULL", num_tokens=4) + set_ctx = MagicMock(return_value=nullcontext()) + with ( + patch( + "vllm_omni.worker_v2.model_states.omni_model_state.current_omni_platform.get_graph_wrapper_cls", + return_value=_FakeGraphRunner, + ), + patch("vllm.forward_context.set_forward_context", set_ctx), + ): + state._run_batched_mtp( + _mtp_batches(), + torch.tensor([101, 202]), + torch.zeros(2, 3), + _DummyInputBatch([0, 1]), + {("codes", "audio")}, + lambda bsz: graph_desc, + ) + _, kwargs = set_ctx.call_args + assert kwargs["batch_descriptor"] is graph_desc + assert kwargs["num_tokens"] == 4 + assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[0, 1, 2]])) + assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[3, 4, 5]])) + + +def test_rope_shim_propagates_type_error(): + # vLLM rope API drift (TypeError) must not be swallowed by the shim. + def broken_get_rope(*_args, **_kwargs): + raise TypeError("vLLM rope API drift") + + with pytest.raises(TypeError, match="rope API drift"): + _make_safe_get_rope(broken_get_rope)(SimpleNamespace(uses_mrope=False), object()) + + +@pytest.mark.parametrize("output_key", [None, ("only_one",)]) +def test_mtp_requires_model_declared_output_key(output_key): + state = _make_state() + model = SimpleNamespace(mtp=lambda: None, mtp_output_key=output_key) + with pytest.raises(TypeError, match="must declare mtp_output_key"): + state._init_mtp_runner(model) diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 7be569f47..a56e6148c 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from dataclasses import MISSING, field from typing import Any @@ -124,6 +127,8 @@ class OmniModelConfig(ModelConfig): async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + 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. active_stream_window: int = 0 duplex_max_sessions: int = 1 @@ -134,6 +139,7 @@ class OmniModelConfig(ModelConfig): # Optional dotted path of a per-stage pooling-output decoder applied # worker-side before IPC. Read by the AR scheduler. pooling_output_decoder: str | None = None + final_output: bool = False hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None stage_connector_config: dict[str, Any] = field( diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index ef48603cf..d64372b84 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -170,7 +170,10 @@ class _ModelEngineOverrides(TypedDict, total=False): limit_mm_per_prompt: dict[str, Any] interleave_mm_strings: bool media_io_kwargs: dict[str, Any] + final_output: bool active_stream_window: int + use_v2_model_runner: bool + supports_native_mrv2_data_plane: bool enable_sleep_mode: bool subtalker_sampling_params: dict[str, Any] silence_ban_frames: int @@ -493,9 +496,12 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): # MiniCPM interleaved AV packing and media decode knobs (Daily-Omni). interleave_mm_strings: bool | None = None media_io_kwargs: dict[str, Any] | None = None + final_output: bool = False active_stream_window: int = Field(default=0, ge=0) session_mode: str = "turn" duplex_max_sessions: int = Field(default=1, ge=1) + use_v2_model_runner: bool = False + supports_native_mrv2_data_plane: bool = False enable_sleep_mode: bool = False default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None @@ -1965,8 +1971,14 @@ def _build_model_config( kwargs["dtype"] = _copy_value(deploy.dtype) if "active_stream_window" not in kwargs: kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window) + kwargs["final_output"] = topology.final_output 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) + kwargs.setdefault("use_v2_model_runner", deploy.model_runner == "v2") + kwargs.setdefault( + "supports_native_mrv2_data_plane", + topology.supports_native_mrv2_data_plane, + ) if "has_sampling_extra_args" not in kwargs: kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args")) if "model_subdir" not in kwargs and topology.model_subdir is not None: diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9f45d357d..aab4752e0 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping from dataclasses import asdict, dataclass, field, fields from enum import Enum from pathlib import Path -from typing import Any, NamedTuple +from typing import Any, Literal, NamedTuple from transformers import PretrainedConfig from vllm.logger import init_logger @@ -250,6 +250,7 @@ class StagePipelineConfig: # Alternates picked by ``merge_pipeline_deploy`` based on ``deploy.async_chunk``. async_chunk_process_next_stage_input_func: str | None = None sync_process_input_func: str | None = None + supports_native_mrv2_data_plane: bool = False # Rewrites the Stage-0 view of a raw prompt before vLLM input processing. # The callable receives ``(prompt, sampling_params_list)``; downstream # stages continue to receive the original prompt. @@ -566,6 +567,7 @@ class DeployConfig: async_chunk: bool = True session_mode: str = "turn" + model_runner: Literal["v1", "v2"] = "v1" # Stage-1 active stream slots; 0 preserves legacy all-stream cycling. active_stream_window: int = 0 duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig) @@ -736,6 +738,16 @@ def _merge_platforms( return merged +def _merge_connectors( + base: dict[str, Any] | None, + overlay: dict[str, Any] | None, +) -> dict[str, Any] | None: + """Deep-merge named connector definitions from a deploy overlay.""" + if not base and not overlay: + return None + return _get_recursively_merged_dict(base or {}, overlay or {}) + + def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: """Load a deploy YAML with optional ``base_config`` inheritance.""" raw_dict = to_dict(load_yaml_config(path)) @@ -748,12 +760,15 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: base_path = Path(path).parent / base_path base_dict = resolve_deploy_yaml(base_path) - # Merge top-level scalars: overlay wins. ``stages:`` and ``platforms:`` - # are deep-merged below so an overlay can layer on top of the base. + # Merge top-level scalars: overlay wins. Structured sections are merged + # below so a thin overlay does not discard inherited runtime contracts. merged = { **base_dict, - **{k: v for k, v in raw_dict.items() if k not in ("stages", "platforms")}, + **{k: v for k, v in raw_dict.items() if k not in ("connectors", "stages", "platforms")}, } + merged_connectors = _merge_connectors(base_dict.get("connectors"), raw_dict.get("connectors")) + if merged_connectors is not None: + merged["connectors"] = merged_connectors merged["stages"] = _merge_stage_lists(base_dict.get("stages"), raw_dict.get("stages")) merged_platforms = _merge_platforms(base_dict.get("platforms"), raw_dict.get("platforms")) if merged_platforms is not None: @@ -773,9 +788,14 @@ def load_deploy_config(path: str | Path) -> DeployConfig: stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] + model_runner = raw_dict.get("model_runner", "v1") + if model_runner not in ("v1", "v2"): + raise ValueError(f"model_runner must be one of ('v1', 'v2'), got {model_runner!r}") + kwargs: dict[str, Any] = { "async_chunk": raw_dict.get("async_chunk", True), "session_mode": raw_dict.get("session_mode", "turn"), + "model_runner": model_runner, "active_stream_window": int(raw_dict.get("active_stream_window", 0) or 0), "duplex_session": DuplexSessionRuntimeConfig(**(raw_dict.get("duplex_session") or {})), "connectors": raw_dict.get("connectors", None), @@ -828,15 +848,25 @@ def _apply_platform_overrides( deploy: DeployConfig, platform: str | None = None, ) -> DeployConfig: - """Merge platform-specific stage overrides into deploy config.""" + """Merge platform-specific runner and stage overrides into deploy config.""" if platform is None: from vllm_omni.platforms import current_omni_platform device_name = current_omni_platform.device_name platform = device_name.lower() if device_name is not None else None + platform_section = (deploy.platforms or {}).get(platform) if platform is not None else None + if platform_section is not None and "model_runner" in platform_section: + model_runner = platform_section["model_runner"] + if model_runner not in ("v1", "v2"): + raise ValueError(f"platform model_runner must be one of ('v1', 'v2'), got {model_runner!r}") + deploy.model_runner = model_runner + if deploy.model_runner == "v2" and platform in {"npu", "xpu"}: + raise NotImplementedError( + f"Model Runner V2 is not supported on {platform.upper()}: " + "the platform worker still uses the legacy chunk-transfer data plane." + ) if platform is None or deploy.platforms is None: return deploy - platform_section = deploy.platforms.get(platform) if platform_section is None: return deploy @@ -998,6 +1028,28 @@ def _build_engine_args( # model-owned streaming state. Propagate it to every stage instead of # making individual models duplicate the value in connector extras. engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions + # The runner selection is a deploy-topology decision owned by the + # ``model_runner`` field; do not let an opaque ``engine_extras`` entry + # silently veto or force it per stage. + if ds is not None: + for reserved in ("use_v2_model_runner", "supports_native_mrv2_data_plane"): + if reserved in ds.engine_extras: + raise ValueError( + f"stage {ds.stage_id}: {reserved!r} must not be set via engine_extras; " + "it is derived from the deploy-level `model_runner` field and the " + "pipeline's `supports_native_mrv2_data_plane` declaration." + ) + engine_args["use_v2_model_runner"] = deploy.model_runner == "v2" + engine_args["supports_native_mrv2_data_plane"] = bool(ps.supports_native_mrv2_data_plane) + if deploy.model_runner == "v2" and not ps.supports_native_mrv2_data_plane: + logger.warning( + "Stage %s (%s) selects model_runner=v2 without declaring " + "supports_native_mrv2_data_plane. It will use the legacy transport path; " + "MRV2 support for this pipeline has not been validated. Use model_runner=v1 " + "unless you are validating a new MRV2 integration.", + ps.stage_id, + ps.model_arch or pipeline.model_arch or pipeline.model_type, + ) 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 @@ -1196,6 +1248,9 @@ class StageConfig: else: engine_args[key] = value + # Terminal-stage ownership comes from topology, not engine overrides. + engine_args["final_output"] = self.final_output + # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) runtime.setdefault("process", True) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 385e180af..4d55c55f8 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -33,6 +33,21 @@ from vllm_omni.engine.serialization import deserialize_additional_information logger = init_logger(__name__) +def _should_emit_engine_output( + model_config: Any, + *, + stopped: bool, + has_control: bool, +) -> bool: + if stopped or has_control: + return True + return not ( + bool(getattr(model_config, "use_v2_model_runner", False)) + and bool(getattr(model_config, "async_chunk", False)) + and not bool(getattr(model_config, "final_output", False)) + ) + + class SampledLogprobContractError(RuntimeError): """The model runner returned unusable sampled-token logprobs.""" @@ -84,6 +99,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): core scheduling logic. """ + max_num_running_reqs: int + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Track requests that need KV cache transfer when finished @@ -200,9 +217,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finished_status: RequestStatus, ) -> list[Request]: """Finish requests and discard any incomplete KV-wait timing.""" + cleanup_ids: Iterable[str] if isinstance(request_ids, str): cleanup_ids = (request_ids,) - finish_request_ids: str | tuple[str, ...] | None = request_ids + finish_request_ids: str | Iterable[str] | None = request_ids elif request_ids is None: cleanup_ids = () finish_request_ids = None @@ -329,15 +347,22 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._process_pending_omni_inputs(model_mode="ar") self._drop_aborted_queued_requests() self._resync_streaming_input_counter() - original_waiting = None if self._should_defer_waiting_admission(): original_waiting = waiting self.waiting = create_request_queue(self.policy) + original_max_num_running_reqs = self.max_num_running_reqs + async_chunk_transport = self._async_chunk_transport_enabled() + reserved_running_slots = ( + self._get_async_chunk_reserved_running_slots() if async_chunk_transport and self.use_v2_model_runner else 0 + ) + if reserved_running_slots: + self.max_num_running_reqs = max(0, original_max_num_running_reqs - reserved_running_slots) try: scheduler_output = super().schedule(throttle_prefills) finally: + self.max_num_running_reqs = original_max_num_running_reqs if original_waiting is not None: deferred_waiting = list(self.waiting) if deferred_waiting: @@ -441,7 +466,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): assert num_tokens_scheduled > 0 request = self.requests.get(req_id) if request is not None: - # vLLM 0.26: settle the in-flight tokens counted in schedule(). + # Settle the in-flight tokens counted in schedule(). # Must happen before the skips below — failed-KV-load and # already-finished requests were incremented too, and the two # readers (allocate_slots, _connector_finished) clamp with @@ -609,6 +634,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): confirmed_num_computed_tokens = None boundary_generation = None + # Capture before resumable stop handling can clear token history. + output_token_ids: Any = getattr(request, "output_token_ids", None) + if output_token_ids is None: + output_token_ids = getattr(request, "_output_token_ids", ()) + num_generation_tokens = len(output_token_ids) if stopped: if self.chunk_transfer_adapter is not None: confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request) @@ -627,6 +657,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): is_segment_finished = not finished if finished: request.resumable = False + if self._native_data_plane: + self._pending_data_plane_terminal_req_ids.add(req_id) if not finished: # for streaming input request only if self.chunk_transfer_adapter: @@ -683,7 +715,18 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) - if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped: + has_stage_output = ( + bool(new_token_ids) + or mm_output is not None + or pooler_output is not None + or kv_transfer_params + or stopped + ) + if has_stage_output and _should_emit_engine_output( + self.vllm_config.model_config, + stopped=stopped, + has_control=kv_transfer_params is not None, + ): OmniSchedulerMixin._append_request_output( self, outputs, @@ -702,6 +745,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): num_nans_in_logits=request.num_nans_in_logits, is_segment_finished=is_segment_finished, new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id), + num_generation_tokens=num_generation_tokens, ) else: # Invariant: EngineCore returns no partial prefill outputs. @@ -877,7 +921,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # This streaming update has already been dequeued. Report the # permanent contract failure so the next scheduling pass # finishes only this request instead of crashing EngineCore. - self.chunk_transfer_adapter.record_receive_failure(req_id, str(exc)) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.record_receive_failure(req_id, str(exc)) return if replaced is not None: if replaced: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index db0eedcf7..9ec55ddf7 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -5,7 +5,7 @@ from __future__ import annotations import time from collections import defaultdict -from typing import Any +from typing import TYPE_CHECKING, Any from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.logger import init_logger @@ -15,11 +15,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all -from vllm.v1.engine import ( - EngineCoreEventType, - EngineCoreOutput, - EngineCoreOutputs, -) +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs from vllm.v1.metrics.perf import PerfStats from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats @@ -27,12 +23,33 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData from vllm_omni.core.sched.utils import omni_routed_experts_for_request +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.outputs import OmniModelRunnerOutput +if TYPE_CHECKING: + from vllm.v1.core.sched.request_queue import RequestQueue + + logger = init_logger(__name__) +def _has_async_chunk_payload_to_run(request: Request) -> bool: + additional_information = getattr(request, "additional_information", None) + if not isinstance(additional_information, dict): + additional_information = deserialize_additional_information(additional_information) + + codes = additional_information.get("codes") + audio = codes.get("audio") if isinstance(codes, dict) else additional_information.get("codes.audio") + if audio is None: + return False + if hasattr(audio, "numel"): + return bool(audio.numel()) + return bool(audio) + + class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): + waiting: RequestQueue + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) model_config = self.vllm_config.model_config @@ -40,6 +57,25 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] + def _is_done_receiving_chunks(self, request_id: str) -> bool: + if self.chunk_transfer_adapter is not None: + return self.chunk_transfer_adapter.is_done_receiving_chunks(request_id) + coordinator = self.input_coordinator + return bool(self._native_data_plane and coordinator is not None and request_id in coordinator.finished_requests) + + def _input_execution_is_terminal( + self, + request_id: str, + scheduler_output: SchedulerOutput, + ) -> bool: + if self._native_data_plane: + return request_id in getattr( + scheduler_output, + "input_terminal_req_ids", + set(), + ) + return self._is_done_receiving_chunks(request_id) + @staticmethod def _record_prefill_stats(request: Request) -> None: """Mirror upstream first-prefill prompt accounting for generation stages. @@ -73,6 +109,29 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): return False return super()._handle_stopped_request(request) + def _requeue_completed_native_chunks(self) -> None: + """Return completed native generation chunks to the admission queue. + + A generation batch slot belongs to one chunk, not the entire stream. + Keep requests with outstanding output in ``running`` so their state + cannot be reused before completion. Other live requests join the tail + of ``waiting`` without freeing their decoder state or KV blocks. + Preserve WAITING_FOR_CHUNK until the coordinator observes fresh input. + Stateful codecs retain lifetime admission until their decoder state is + released, even when no chunk is currently executing. + """ + if not getattr(self, "_native_data_plane", False) or self._retains_state_across_chunks: + return + in_flight: list[Request] = [] + for request in self.running: + if request.num_in_flight_tokens > 0 or request.is_finished(): + in_flight.append(request) + continue + if request.status == RequestStatus.RUNNING: + request.status = RequestStatus.WAITING + self.waiting.add_request(request) + self.running = in_flight + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: """One-shot generation fast path: - Feed all input tokens of the request at once @@ -82,6 +141,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): """ token_budget = self.max_num_scheduled_tokens + execution_batch_size = self.max_num_running_reqs if self._pause_state == PauseState.PAUSED_ALL: token_budget = 0 scheduled_timestamp = time.monotonic() @@ -102,14 +162,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 self._drop_aborted_queued_requests() + self._requeue_completed_native_chunks() self._process_pending_omni_inputs(model_mode="generation") self._drop_aborted_queued_requests() self._resync_streaming_input_counter() + async_chunk_transport = self._async_chunk_transport_enabled() + native_chunks = bool(getattr(self, "_native_data_plane", False)) + # Parking releases an execution slot, but stateful codecs retain their + # request-owned state until completion or abort, in either runner. + reserved_running_slots = ( + self._get_async_chunk_reserved_running_slots() if self._retains_state_across_chunks else 0 + ) # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled already_finished_reqs: set[Request] = set() - while req_index < len(self.running) and token_budget > 0: + while req_index < len(self.running) and token_budget > 0 and len(num_scheduled_tokens) < execution_batch_size: request = self.running[req_index] # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( @@ -119,6 +187,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): req_index += 1 continue + if native_chunks and request.num_in_flight_tokens > 0: + # Overlap different streams, never two stateful chunks from + # the same stream. Readiness remains pending for the next step. + req_index += 1 + continue + num_computed_tokens = request.num_computed_tokens required_tokens = len(request.prompt_token_ids) - num_computed_tokens if not self.scheduler_config.enable_chunked_prefill and required_tokens > token_budget: @@ -127,12 +201,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): break # async_chunk: don't schedule placeholder tokens when no new chunk is available. if required_tokens <= 0: - if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks( - request.request_id - ): - self._pending_finish_reqs.append(request) - req_index += 1 - continue + if async_chunk_transport and self._is_done_receiving_chunks(request.request_id): + if len(request.prompt_token_ids) == 0 and _has_async_chunk_payload_to_run(request): + required_tokens = 1 + else: + self._pending_finish_reqs.append(request) + req_index += 1 + continue + else: + req_index += 1 + continue num_new_tokens = min(required_tokens, token_budget) new_blocks = self.kv_cache_manager.allocate_slots( request, @@ -161,20 +239,27 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if already_finished_reqs: self.running = remove_all(self.running, already_finished_reqs) - # Fast path selection and scheduling for one-shot generation requests, - # independent of pooling_params. - while self.waiting and token_budget > 0 and self._pause_state == PauseState.UNPAUSED: - # Requests waiting for their next chunk are temporarily absent - # from `running`, but stateful models still retain their model - # runner slot. Mirror vLLM's treatment of - # `num_waiting_for_streaming_input` when enforcing max_num_seqs. - num_running = len(self.running) - if self._retains_state_across_chunks and self.chunk_transfer_adapter is not None: - num_running += self.chunk_transfer_adapter.num_running_waiting_for_chunk - if num_running >= self.max_num_running_reqs: - break - + # Fast path selection and scheduling (treat all as diffusion requests, + # independent of pooling_params) + while ( + self.waiting + and token_budget > 0 + and len(num_scheduled_tokens) < execution_batch_size + and ( + len(num_scheduled_tokens) + if native_chunks and not self._retains_state_across_chunks + else len(self.running) + reserved_running_slots + ) + < self.max_num_running_reqs + and self._pause_state == PauseState.UNPAUSED + ): request = self.waiting.peek_request() + if native_chunks and request.num_in_flight_tokens > 0: + # A restored waiting entry can still own an outstanding batch. + # Do not let it block other ready streams or execute it twice. + self.waiting.pop_request() + skipped_waiting_requests.add_request(request) + continue # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED @@ -184,11 +269,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): continue # async_chunk: wait for the first upstream chunk (don't start with placeholders). - if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0: - if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id): - self.waiting.pop_request() - self._pending_finish_reqs.append(request) - continue + if async_chunk_transport and len(request.prompt_token_ids) == 0: + if self._is_done_receiving_chunks(request.request_id): + if not _has_async_chunk_payload_to_run(request): + self.waiting.pop_request() + self._pending_finish_reqs.append(request) + continue else: self.waiting.pop_request() skipped_waiting_requests.prepend_request(request) @@ -197,6 +283,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Allocate all input tokens for the request in one shot # (allocate 1 placeholder if zero) required_tokens = max(len(request.prompt_token_ids), 1) + if native_chunks: + required_tokens = len(request.prompt_token_ids) - request.num_computed_tokens + if required_tokens <= 0: + if ( + self._is_done_receiving_chunks(request.request_id) + and not request.prompt_token_ids + and _has_async_chunk_payload_to_run(request) + ): + required_tokens = 1 + else: + self.waiting.pop_request() + if self._is_done_receiving_chunks(request.request_id): + self._pending_finish_reqs.append(request) + else: + skipped_waiting_requests.add_request(request) + continue num_new_tokens = min(required_tokens, token_budget) new_blocks = self.kv_cache_manager.allocate_slots( request, @@ -229,7 +331,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # If fast path scheduled none, fall back to the original scheduling if not num_scheduled_tokens: - if self.chunk_transfer_adapter: + if async_chunk_transport: # Don't fall back: base scheduler doesn't handle async_chunk # requests with empty prompt_token_ids. self._restore_omni_wait_queues() @@ -245,7 +347,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): any_request = self.running[0] num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id) - # Assemble SchedulerOutput (align with v0.14.0) + # Assemble the current SchedulerOutput contract. if self.use_v2_model_runner: # No resumed reqs in fast path; pass prefill_token_ids for new reqs. new_reqs_data = [ @@ -284,7 +386,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) - # Record the request ids scheduled in this step (v0.14.0 behavior). + # Record request ids scheduled in this step for the next update. self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) @@ -404,7 +506,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): assert num_tokens_scheduled > 0 request = self.requests.get(req_id) if request is not None: - # vLLM 0.26: settle the in-flight tokens counted in schedule(). + # Settle the in-flight tokens counted in schedule(). # Must happen before the skips below — failed-KV-load and # already-finished requests were incremented too, and the two # readers (allocate_slots, _connector_finished) clamp with @@ -496,10 +598,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if ( request.status == RequestStatus.FINISHED_ERROR or request.status == RequestStatus.FINISHED_STOPPED - or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens) or ( - self.chunk_transfer_adapter is not None - and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id) + not self._async_chunk_transport_enabled() + and request.num_computed_tokens >= request.num_prompt_tokens + ) + or ( + self._async_chunk_transport_enabled() + and self._input_execution_is_terminal( + request.request_id, + scheduler_output, + ) and request.num_computed_tokens >= len(request.prompt_token_ids) ) ): @@ -526,6 +634,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason = request.get_finished_reason() finished = self._handle_stopped_request(request) is_segment_finished = not finished + if finished and self._native_data_plane: + self._pending_data_plane_terminal_req_ids.add(req_id) if not finished: # for streaming input request only if self.chunk_transfer_adapter: diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index f33359835..d714d2665 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -5,17 +5,20 @@ from __future__ import annotations import math import os +import queue import time from collections.abc import Iterable, Iterator from typing import Any import torch from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.config import VllmConfig from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.request_queue import RequestQueue from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import ( EngineCoreEventType, @@ -31,6 +34,7 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, uses_full_payload_input_coordinator, + uses_native_mrv2_data_plane, ) from vllm_omni.core.sched.output import ( OmniChunkRecvHandle, @@ -41,6 +45,8 @@ from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapt OmniChunkTransferAdapter, ) from vllm_omni.engine import OmniEngineCoreOutput +from vllm_omni.engine.serialization import serialize_additional_information +from vllm_omni.outputs import OmniConnectorOutput logger = init_logger(__name__) @@ -101,6 +107,19 @@ elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0: class OmniSchedulerMixin: """Shared scheduler helpers for omni-specific request handling.""" + # Provided by the concrete vLLM scheduler. + vllm_config: VllmConfig + requests: dict[str, Request] + waiting: RequestQueue + running: list[Request] + + def _init_omni_connector_output_inbox(self) -> None: + self._omni_connector_output_inbox: queue.SimpleQueue[OmniConnectorOutput] = queue.SimpleQueue() + + def enqueue_omni_connector_output(self, output: OmniConnectorOutput) -> None: + """Accept a lightweight readiness event from the local runner thread.""" + self._omni_connector_output_inbox.put(output) + # ------------------------------------------------------------------ # # Shared scheduler/output helpers (lift the AR / generation duplicates) # ------------------------------------------------------------------ # @@ -119,15 +138,30 @@ class OmniSchedulerMixin: self.use_v2_model_runner = False model_config = self.vllm_config.model_config - self.chunk_transfer_adapter = ( - OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None + self.use_v2_model_runner = bool(getattr(model_config, "use_v2_model_runner", False)) + self._native_data_plane = uses_native_mrv2_data_plane( + model_config, + use_v2_model_runner=self.use_v2_model_runner, ) - self.input_coordinator = ( - OmniSchedulingCoordinator(stage_id=getattr(model_config, "stage_id", 0)) - if uses_full_payload_input_coordinator(model_config) + self.chunk_transfer_adapter = ( + OmniChunkTransferAdapter(self.vllm_config) + if getattr(model_config, "async_chunk", False) and not self._native_data_plane else None ) - self._latest_omni_connector_output = None + self.input_coordinator: OmniSchedulingCoordinator | None = None + if self._native_data_plane: + self.input_coordinator = OmniSchedulingCoordinator( + scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs, + stage_id=getattr(model_config, "stage_id", 0), + async_chunk=True, + ) + elif uses_full_payload_input_coordinator(model_config): + self.input_coordinator = OmniSchedulingCoordinator( + stage_id=getattr(model_config, "stage_id", 0), + ) + self._latest_omni_connector_output: OmniConnectorOutput | None = None + self._init_omni_connector_output_inbox() + self._pending_data_plane_terminal_req_ids: set[str] = set() # Optional per-stage pooling-output decoder hook (dotted path in # model_config); applied worker-side before IPC. self._pooling_output_decoder = None @@ -214,6 +248,35 @@ class OmniSchedulerMixin: session.num_prompt_tokens = len(new_prompt) self._finish_streaming_session_update(session, update) + def _async_chunk_transport_enabled(self) -> bool: + return getattr(self, "chunk_transfer_adapter", None) is not None or bool( + getattr(self, "_native_data_plane", False) + ) + + def _get_async_chunk_reserved_running_slots(self) -> int: + adapter = getattr(self, "chunk_transfer_adapter", None) + live_requests = self.requests + reserved_ids: set[str] = set() + if adapter is not None: + parked = getattr(adapter, "waiting_for_chunk_running_requests", ()) + held = getattr(adapter, "_held_non_active", ()) + else: + coordinator = getattr(self, "input_coordinator", None) + parked = getattr(coordinator, "_waiting_for_chunk_running", ()) + held = () + for request in parked: + req_id = getattr(request, "request_id", None) + if req_id is not None and req_id in live_requests: + reserved_ids.add(req_id) + for request in held: + req_id = getattr(request, "request_id", None) + if req_id is not None and req_id in live_requests: + reserved_ids.add(req_id) + return len(reserved_ids) + + # ------------------------------------------------------------------ # + # Shared scheduler/output helpers (lift the AR / generation duplicates) + # ------------------------------------------------------------------ # def _release_replaced_streaming_prompt_cache(self, session: Request) -> None: """Discard cache state that belongs to a replaced prompt.""" # A prompt replacement is not a normal streaming extension: none of @@ -259,19 +322,55 @@ class OmniSchedulerMixin: AR and generation schedulers except for the ``model_mode`` argument forwarded to ``update_request_metadata``. """ + connector_outputs: list[OmniConnectorOutput] = [] + inbox = getattr(self, "_omni_connector_output_inbox", None) + if inbox is not None: + while True: + try: + connector_outputs.append(inbox.get_nowait()) + except queue.Empty: + break connector_output = getattr(self, "_latest_omni_connector_output", None) self._latest_omni_connector_output = None + if connector_output is not None: + connector_outputs.append(connector_output) input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is None: return - if connector_output and connector_output.request_metadata: + request_metadata: dict[str, dict[str, Any]] = {} + chunk_ready_req_ids: set[str] = set() + chunk_finished_req_ids: set[str] = set() + stage_recv_req_ids: set[str] = set() + for output in connector_outputs: + request_metadata.update(output.request_metadata) + chunk_ready_req_ids.update(output.chunk_ready_req_ids) + chunk_finished_req_ids.update(output.chunk_finished_req_ids) + stage_recv_req_ids.update(output.stage_recv_req_ids) + live_request_ids = self.requests.keys() + request_metadata = { + req_id: metadata for req_id, metadata in request_metadata.items() if req_id in live_request_ids + } + chunk_ready_req_ids.intersection_update(live_request_ids) + chunk_finished_req_ids.intersection_update(live_request_ids) + stage_recv_req_ids.intersection_update(live_request_ids) + if request_metadata: input_coordinator.update_request_metadata( - self.requests, connector_output.request_metadata, model_mode=model_mode + self.requests, + request_metadata, + model_mode=model_mode, + ) + if input_coordinator._async_chunk: + input_coordinator.process_pending_chunks( + self.waiting, + self.running, + chunk_ready_req_ids, + chunk_finished_req_ids, + ) + else: + input_coordinator.process_pending_full_payload_inputs( + self.waiting, + stage_recv_req_ids, ) - input_coordinator.process_pending_full_payload_inputs( - self.waiting, - connector_output.stage_recv_req_ids if connector_output else set(), - ) def _process_pending_omni_inputs(self, model_mode: str) -> None: """Apply pending connector inputs, timeouts, and async chunks.""" @@ -316,7 +415,7 @@ class OmniSchedulerMixin: scheduler_requests=self.requests, ) if self.input_coordinator: - self.input_coordinator.restore_queues(self.waiting) + self.input_coordinator.restore_queues(self.waiting, self.running) def _process_pending_input_timeouts(self) -> None: """Force-fail requests waiting on the full-payload coordinator too long. @@ -458,6 +557,8 @@ class OmniSchedulerMixin: *, finished_requests_needing_kv_transfer: dict | None = None, pending_input_registrations: list[OmniChunkRecvHandle] | None = None, + data_plane_terminal_req_ids: set[str] | None = None, + input_terminal_req_ids: set[str] | None = None, ) -> OmniSchedulerOutput: """Wrap a base ``SchedulerOutput`` in ``OmniSchedulerOutput``. @@ -468,11 +569,31 @@ class OmniSchedulerMixin: base_data = {name: getattr(base, name) for name in SchedulerOutput.__dataclass_fields__} input_coordinator = getattr(self, "input_coordinator", None) if pending_input_registrations is None: - pending_input_registrations = input_coordinator.pending_input_registrations if input_coordinator else [] + if input_coordinator is None: + pending_input_registrations = [] + elif input_coordinator._async_chunk: + pending_input_registrations = input_coordinator.pending_chunk_registrations + else: + pending_input_registrations = input_coordinator.pending_input_registrations + if data_plane_terminal_req_ids is None: + pending_terminal: set[str] = getattr(self, "_pending_data_plane_terminal_req_ids", set()) + data_plane_terminal_req_ids = set(pending_terminal) + pending_terminal.clear() + if input_coordinator is not None: + scheduled_terminal_req_ids = input_coordinator.get_scheduled_input_terminal_req_ids(base) + if input_terminal_req_ids is None: + input_terminal_req_ids = scheduled_terminal_req_ids + else: + input_terminal_req_ids = set(input_terminal_req_ids).union( + scheduled_terminal_req_ids, + ) + input_coordinator.postprocess_scheduler_output(base) return OmniSchedulerOutput( **base_data, finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {}, pending_input_registrations=pending_input_registrations, + data_plane_terminal_req_ids=data_plane_terminal_req_ids, + input_terminal_req_ids=input_terminal_req_ids or set(), ) def _rewrap_scheduled_new_reqs(self, scheduler_output: SchedulerOutput) -> None: @@ -520,8 +641,14 @@ class OmniSchedulerMixin: num_nans_in_logits: int = 0, is_segment_finished: bool | None = False, new_prompt_len_snapshot: int | None = None, + num_generation_tokens: int | None = None, ) -> OmniEngineCoreOutput: """Build the common request-output envelope used by LLM schedulers.""" + pooling_output_payload = None + if isinstance(pooling_output, dict): + pooling_output_payload = serialize_additional_information(pooling_output) + pooling_output = None + return OmniEngineCoreOutput( request_id=request.request_id, new_token_ids=new_token_ids, @@ -529,6 +656,7 @@ class OmniSchedulerMixin: new_logprobs=new_logprobs, new_prompt_logprobs_tensors=new_prompt_logprobs_tensors, pooling_output=pooling_output, + pooling_output_payload=pooling_output_payload, multimodal_output=multimodal_output, stop_reason=stop_reason, events=request.take_events(), @@ -540,6 +668,7 @@ class OmniSchedulerMixin: num_nans_in_logits=num_nans_in_logits, is_segment_finished=is_segment_finished, new_prompt_len_snapshot=new_prompt_len_snapshot, + num_generation_tokens=num_generation_tokens, ) def _append_request_output( @@ -953,10 +1082,10 @@ class OmniSchedulerMixin: self.running[:] = [req for req in self.running if keep_running(req)] def _drop_aborted_queued_requests(self) -> None: - for queue in (self.waiting, self.skipped_waiting): - aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED] + for request_queue in (self.waiting, self.skipped_waiting): + aborted = [req for req in request_queue if req.status == RequestStatus.FINISHED_ABORTED] if aborted: - queue.remove_requests(aborted) + request_queue.remove_requests(aborted) self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED] def _resync_streaming_input_counter(self) -> None: diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 3d516f5e3..70b682ac0 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Scheduling-side coordination for full_payload input waiting. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Scheduling-side coordination for chunk and full-payload input waiting. -Manages WAITING_FOR_INPUT state transitions based on readiness signals -from OmniConnectorOutput, without ever calling connector.put()/get(). - -Chunk waiting (WAITING_FOR_CHUNK) lives on OmniChunkTransferAdapter. +Manages WAITING_FOR_CHUNK and WAITING_FOR_INPUT state transitions based on +readiness signals from OmniConnectorOutput, without calling connector I/O. """ from __future__ import annotations @@ -22,6 +20,18 @@ from vllm_omni.core.sched.output import OmniChunkRecvHandle logger = init_logger(__name__) +def uses_native_mrv2_data_plane( + model_config: Any, + *, + use_v2_model_runner: bool, +) -> bool: + return bool( + use_v2_model_runner + and getattr(model_config, "async_chunk", False) + and getattr(model_config, "supports_native_mrv2_data_plane", False) + ) + + def uses_full_payload_input_coordinator(model_config: Any) -> bool: """Returns True if this stage parks pending requests in WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector. @@ -38,7 +48,7 @@ def uses_full_payload_input_coordinator(model_config: Any) -> bool: class OmniSchedulingCoordinator: - """Pure-scheduling coordinator for full_payload input waiting. + """Pure-scheduling coordinator for chunk and full-payload input waiting. The Scheduler owns an instance of this class. It consumes readiness signals produced by the Model Runner's ``OmniConnectorModelRunnerMixin`` @@ -46,12 +56,28 @@ class OmniSchedulingCoordinator: transitions accordingly. """ - def __init__(self, stage_id: int = 0): + def __init__( + self, + scheduler_max_num_seqs: int = 0, + stage_id: int = 0, + async_chunk: bool = False, + ) -> None: self._stage_id = stage_id + self._scheduler_max_num_seqs = scheduler_max_num_seqs + self._async_chunk = async_chunk self.finished_requests: set[str] = set() + self.requests_with_ready_chunks: set[str] = set() + self.input_terminal_req_ids: set[str] = set() self._full_payload_input_received: set[str] = set() + self._waiting_for_chunk_waiting: deque[Any] = deque() + self._waiting_for_chunk_running: deque[Any] = deque() + + # Request IDs that were newly registered for chunk recv this cycle. + # The engine/Model Runner should call register_chunk_recv() for these + # so the bg thread starts polling. + self.pending_chunk_registrations: list[OmniChunkRecvHandle] = [] # Requests waiting for full_payload stage input (WAITING_FOR_INPUT). self._waiting_for_input: deque[Any] = deque() # Per-cycle list of minimal handles to ship to the model runner so it @@ -69,6 +95,47 @@ class OmniSchedulingCoordinator: # Core scheduling methods # ------------------------------------------------------------------ # + def process_pending_chunks( + self, + waiting_queue: Any, + running_queue: list[Request], + chunk_ready_req_ids: set[str], + chunk_finished_req_ids: set[str], + ) -> None: + """Transition requests whose chunks have arrived. + + Args: + waiting_queue: Scheduler's waiting request queue. + running_queue: Scheduler's running request list. + chunk_ready_req_ids: IDs with a newly arrived chunk this cycle. + chunk_finished_req_ids: IDs whose final chunk has arrived. + """ + if self._stage_id == 0 or not self._async_chunk: + return + + self.finished_requests.update(chunk_finished_req_ids) + self.pending_chunk_registrations = [] + + self._process_chunk_queue( + waiting_queue, + self._waiting_for_chunk_waiting, + RequestStatus.WAITING, + chunk_ready_req_ids, + ) + self._process_chunk_queue( + running_queue, + self._waiting_for_chunk_running, + RequestStatus.RUNNING, + chunk_ready_req_ids, + ) + while len(running_queue) > self._scheduler_max_num_seqs: + request = running_queue.pop() + # Must reset status to WAITING so the scheduler treats it as + # schedulable work. KV blocks are NOT freed here (unlike a + # real preemption), so PREEMPTED would be incorrect. + request.status = RequestStatus.WAITING + waiting_queue.prepend_requests([request]) + def process_pending_full_payload_inputs( self, waiting_queue: Any, @@ -144,10 +211,29 @@ class OmniSchedulingCoordinator: waiting_queue.remove_requests(to_remove) def free_finished_request(self, request_id: str) -> None: - """Prune internal tracking sets for a freed request to prevent unbounded growth.""" + """Prune all coordinator state owned by a freed request.""" self._full_payload_input_received.discard(request_id) self.finished_requests.discard(request_id) + self.requests_with_ready_chunks.discard(request_id) + self.input_terminal_req_ids.discard(request_id) self._waiting_since.pop(request_id, None) + for queue_attr in ( + "_waiting_for_chunk_waiting", + "_waiting_for_chunk_running", + "_waiting_for_input", + ): + queue = getattr(self, queue_attr) + setattr( + self, + queue_attr, + deque(request for request in queue if request.request_id != request_id), + ) + self.pending_chunk_registrations = [ + handle for handle in self.pending_chunk_registrations if handle.request_id != request_id + ] + self.pending_input_registrations = [ + handle for handle in self.pending_input_registrations if handle.request_id != request_id + ] def collect_timed_out_request_ids( self, @@ -178,11 +264,14 @@ class OmniSchedulingCoordinator: # Defensively remove from coordinator internal queues (may already # be empty if restore_queues() has run). - remaining: deque[Any] = deque() - for request in self._waiting_for_input: - if request.request_id not in timed_out_ids: - remaining.append(request) - self._waiting_for_input = remaining + for queue_attr in ( + "_waiting_for_chunk_waiting", + "_waiting_for_chunk_running", + "_waiting_for_input", + ): + queue = getattr(self, queue_attr) + remaining = deque(request for request in queue if request.request_id not in timed_out_ids) + setattr(self, queue_attr, remaining) for req_id in timed_out_ids: self._waiting_since.pop(req_id, None) @@ -198,8 +287,17 @@ class OmniSchedulingCoordinator: def restore_queues( self, waiting_queue: Any, + running_queue: list[Request] | None = None, ) -> None: - """Return waiting-for-input requests to the waiting queue.""" + """Return waiting-for-chunk/input requests to scheduling queues.""" + for request in self._waiting_for_chunk_waiting: + waiting_queue.add_request(request) + self._waiting_for_chunk_waiting = deque() + + if running_queue is not None and self._waiting_for_chunk_running: + running_queue.extend(self._waiting_for_chunk_running) + self._waiting_for_chunk_running = deque() + for request in self._waiting_for_input: waiting_queue.add_request(request) self._waiting_for_input = deque() @@ -247,6 +345,9 @@ class OmniSchedulingCoordinator: if request is None: continue + if metadata.get("input_terminal") is True: + self.input_terminal_req_ids.add(req_id) + # Handle next_stage_prompt_len if present (for models like Qwen3-Omni). # Only apply when the request has not started decoding yet # (no output tokens). Resetting a mid-decode request would @@ -255,8 +356,7 @@ class OmniSchedulingCoordinator: next_len = metadata["next_stage_prompt_len"] if isinstance(next_len, int) and next_len > 0: output_token_ids = getattr(request, "_output_token_ids", None) - has_decode_output = output_token_ids is not None and len(output_token_ids) > 0 - if has_decode_output: + if output_token_ids is not None and len(output_token_ids) > 0: logger.debug( "[Coordinator stage-%s] Skipping prompt resize for req %s: " "request already has %s output tokens", @@ -284,6 +384,12 @@ class OmniSchedulingCoordinator: if model_mode != "ar": new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes")) + runtime_seed = None + if "left_context_size" in metadata: + runtime_seed = { + "meta": {"left_context_size": metadata["left_context_size"]}, + } + request._omni_initial_model_buffer = runtime_seed if new_ids: request.prompt_token_ids = new_ids request.num_prompt_tokens = len(new_ids) @@ -291,3 +397,103 @@ class OmniSchedulingCoordinator: request._all_token_ids.extend(new_ids) request._output_token_ids.clear() request.num_computed_tokens = 0 + + def postprocess_scheduler_output( + self, + scheduler_output: Any, + requests: dict[str, Request] | None = None, + ) -> None: + """Clear per-cycle ready state after scheduler output is materialized.""" + self._clear_chunk_ready(scheduler_output) + self.input_terminal_req_ids.difference_update( + self._scheduled_request_ids(scheduler_output), + ) + + def get_scheduled_input_terminal_req_ids( + self, + scheduler_output: Any, + ) -> set[str]: + """Return terminal inputs executed by this immutable scheduler step.""" + return self.input_terminal_req_ids.intersection( + self._scheduled_request_ids(scheduler_output), + ) + + @staticmethod + def _scheduled_request_ids(scheduler_output: Any) -> set[str]: + num_scheduled_tokens = getattr( + scheduler_output, + "num_scheduled_tokens", + None, + ) + if num_scheduled_tokens is not None: + return set(num_scheduled_tokens) + + request_ids = {req.req_id for req in getattr(scheduler_output, "scheduled_new_reqs", ())} + cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) + request_ids.update(getattr(cached_reqs, "req_ids", ())) + return request_ids + + # ------------------------------------------------------------------ # + # Internal helpers + # ------------------------------------------------------------------ # + + def _process_chunk_queue( + self, + queue: Any, + waiting_for_chunk_list: deque[Any], + target_status: RequestStatus, + chunk_ready_req_ids: set[str], + ) -> None: + queue_snapshot = list(queue) + for request in queue_snapshot: + request_can_run = self._request_can_run(request, chunk_ready_req_ids) + if request.status != RequestStatus.WAITING_FOR_CHUNK: + if request_can_run: + if request.request_id in chunk_ready_req_ids: + self.requests_with_ready_chunks.add(request.request_id) + continue + if request.status == RequestStatus.WAITING_FOR_INPUT: + continue + self.pending_chunk_registrations.append( + OmniChunkRecvHandle( + request_id=request.request_id, + external_req_id=getattr(request, "external_req_id", None), + ) + ) + request.status = RequestStatus.WAITING_FOR_CHUNK + self._waiting_since.setdefault(request.request_id, time.monotonic()) + else: + if request_can_run: + request.status = target_status + if request.request_id in chunk_ready_req_ids: + self.requests_with_ready_chunks.add(request.request_id) + self._waiting_since.pop(request.request_id, None) + continue + if isinstance(queue, list): + queue.remove(request) + else: + queue.remove_request(request) + waiting_for_chunk_list.append(request) + + def _request_can_run( + self, + request: Request, + chunk_ready_req_ids: set[str], + ) -> bool: + req_id = request.request_id + if req_id in self.finished_requests: + # Admit the final payload or terminal marker for cleanup. + return True + + return req_id in chunk_ready_req_ids or req_id in self.requests_with_ready_chunks + + def _clear_chunk_ready(self, scheduler_output: Any) -> None: + if scheduler_output.scheduled_new_reqs: + for req_data in scheduler_output.scheduled_new_reqs: + req_id = getattr(req_data, "req_id", None) + if req_id is not None: + self.requests_with_ready_chunks.discard(req_id) + + if scheduler_output.scheduled_cached_reqs: + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + self.requests_with_ready_chunks.discard(req_id) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index deaf44df8..58aafea01 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from dataclasses import dataclass, field, fields from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput @@ -18,14 +21,14 @@ class OmniNewRequestData(NewRequestData): Args: external_req_id: Optional external request ID for tracking - additional_information: Optional serialized additional information + additional_information: Optional serialized or materialized additional information dictionary containing tensors or lists model_intermediate_buffer: Optional runner-owned payload for GPUModelRunner.model_intermediate_buffer """ external_req_id: str | None = None - additional_information: AdditionalInformationPayload | None = None + additional_information: AdditionalInformationPayload | dict[str, object] | None = None model_intermediate_buffer: dict[str, object] | None = None @classmethod @@ -114,3 +117,5 @@ class OmniSchedulerOutput(SchedulerOutput): finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict) pending_input_registrations: list[OmniChunkRecvHandle] = field(default_factory=list) + data_plane_terminal_req_ids: set[str] = field(default_factory=set) + input_terminal_req_ids: set[str] = field(default_factory=set) diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 3e92eb37a..52c27f20e 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -34,9 +34,14 @@ class HiddenStates(TypedDict, total=False): class Embeddings(TypedDict, total=False): + prepared_prefill: torch.Tensor prefill: torch.Tensor decode: torch.Tensor + decode_token_start: int + decode_token_end: int cached_decode: torch.Tensor + cached_decode_token_start: int + cached_decode_token_end: int tts_bos: torch.Tensor tts_eos: torch.Tensor tts_pad: torch.Tensor @@ -54,6 +59,7 @@ class Codes(TypedDict, total=False): class Ids(TypedDict, total=False): + prepared_prefill: torch.Tensor all: list[int] prompt: list[int] output: list[int] @@ -81,6 +87,7 @@ class OmniPayloadMeta(TypedDict, total=False): next_stage_prompt_len: int next_stage_generation_tokens: int replace_streaming_prompt: bool + next_stage_prompt_ids: list[int] streaming_prompt_recompute: bool streaming_condition_seq: int replace_runtime_additional_information: bool @@ -94,6 +101,7 @@ class OmniPayloadMeta(TypedDict, total=False): width: int decode_flag: bool codec_streaming: bool + codec_frame_valid: bool | torch.Tensor ref_code_len: int ref_context_size: int ref_context_request_id: str @@ -148,6 +156,8 @@ class EmbeddingsStruct(_StructBase): decode_token_start: int | None = None decode_token_end: int | None = None cached_decode: torch.Tensor | None = None + cached_decode_token_start: int | None = None + cached_decode_token_end: int | None = None tts_bos: torch.Tensor | None = None tts_eos: torch.Tensor | None = None tts_pad: torch.Tensor | None = None @@ -190,6 +200,7 @@ class MetaStruct(_StructBase): next_stage_prompt_len: int | None = None next_stage_generation_tokens: int | None = None replace_streaming_prompt: bool | None = None + next_stage_prompt_ids: list[int] | None = None streaming_prompt_recompute: bool | None = None streaming_condition_seq: int | None = None replace_runtime_additional_information: bool | None = None @@ -203,6 +214,7 @@ class MetaStruct(_StructBase): width: int | None = None decode_flag: bool | None = None codec_streaming: bool | None = None + codec_frame_valid: torch.Tensor | None = None ref_code_len: int | None = None # Expected FINAL length of a growing async-chunk sequence, when the # producer knows it up front (e.g. a frame-locked AR stage whose diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml new file mode 100644 index 000000000..b6ebcf67f --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml @@ -0,0 +1,28 @@ +# Opt-in CUDA Model Runner V2 profile. The base profile retains V1. +base_config: qwen3_tts_high_concurrency.yaml +model_runner: v2 + +platforms: + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 + +connectors: + connector_of_shared_memory: + extra: + # Keep unqualified prefix-graph optimizations disabled in this profile. + code_predictor_prefix_graphs: false + ref_audio_artifact_cache_max_entries: 1024 + # Opt-in B2 codec graphs. Validate quality for the deployment workload. + decode_batch_max_size: 2 + decode_cudagraph_batch_sizes: [1, 2] + +stages: + - stage_id: 1 + # Enable the decoder's inner graphs; the generation runner stays eager. + enforce_eager: false diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml new file mode 100644 index 000000000..4551e35d9 --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml @@ -0,0 +1,8 @@ +# Experimental B4 profile. Early first audio can require startup buffering. +# MPS is a separate operator setting and is not enabled by this file. +base_config: qwen3_tts_high_concurrency_mrv2.yaml +connectors: + connector_of_shared_memory: + extra: + decode_batch_max_size: 4 + decode_cudagraph_batch_sizes: [1, 2, 3, 4] diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml new file mode 100644 index 000000000..da6a90cf4 --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml @@ -0,0 +1,23 @@ +# Opt-in single-GPU throughput profile, measured on H200 at concurrency 128 +# with Qwen3-TTS-12Hz-1.7B-Base. Warm the full pipeline before benchmarking. +# B8 uses more graph memory than B4; this is not a low-latency profile. +# MPS is optional and is not enabled by this file. +# +# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \ +# --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml +base_config: qwen3_tts_high_concurrency_mrv2.yaml + +connectors: + connector_of_shared_memory: + extra: + decode_batch_max_size: 8 + decode_cudagraph_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 128 + - stage_id: 1 + devices: "0" + # Admission candidates are grouped by codec phase before B8 execution. + max_num_seqs: 32 diff --git a/vllm_omni/deploy/qwen3_tts_mrv2.yaml b/vllm_omni/deploy/qwen3_tts_mrv2.yaml new file mode 100644 index 000000000..228076c8e --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_mrv2.yaml @@ -0,0 +1,18 @@ +# Opt-in CUDA Model Runner V2 profile. The base profile retains V1. +base_config: qwen3_tts.yaml +model_runner: v2 + +platforms: + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 + +stages: + - stage_id: 0 + # Bound per-step prefill work, matching the high-concurrency TTS profile. + max_num_batched_tokens: 512 diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 3a51795c4..f2c01da9c 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import fcntl import os diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 826ccec72..23c218c71 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -7,6 +7,7 @@ from __future__ import annotations import importlib import inspect from collections import deque +from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast import torch @@ -65,7 +66,6 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): extracted: dict[str, Any] = {} meta = payload.get("meta") if isinstance(payload, dict) else None meta = meta if isinstance(meta, dict) else {} - if "next_stage_prompt_len" in meta: extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] else: @@ -85,6 +85,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): elif "left_context_size" in payload: logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'") + if cls._payload_finished(payload): + extracted["input_terminal"] = True + return extracted _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = { @@ -331,67 +334,11 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): """ if not hasattr(self, "_lock"): return OmniConnectorOutput() - - tp_group = self._get_local_tp_group() - if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: - if self.is_data_transfer_rank(): - with self._lock: - fanout_packet = self._collect_async_chunk_fanout_packet_locked() - else: - fanout_packet = None - fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) - if fanout_packet is None: - newly_finished = set() - chunk_finished = set() - request_metadata = {} - else: - if not self.is_data_transfer_rank(): - self._apply_async_chunk_fanout_packet(fanout_packet) - newly_finished = set(fanout_packet["newly_finished"]) - chunk_finished = set(fanout_packet["chunk_finished"]) - request_metadata = dict(fanout_packet["request_metadata"]) - else: - with self._lock: - newly_finished = set(self._finished_load_reqs) - self._finished_load_reqs.clear() - chunk_finished = set(self._chunk_finished_req_ids) - self._chunk_finished_req_ids.clear() - request_metadata = dict(self._local_request_metadata) - self._local_request_metadata.clear() - # _send_side_request_payload is the async accumulation buffer for - # future recv chunks. Clearing it on every consumable wake-up drops - # intermediate - # thinker decode spans before the model side can consume them. - # Only terminal chunk_finished requests may release that buffer. - for req_id in chunk_finished: - if req_id not in self._local_stage_payload_cache: - continue - ext_req_id = self._request_ids_mapping.get(req_id, req_id) - self._send_side_request_payload.pop(ext_req_id, None) - if ext_req_id != req_id: - self._send_side_request_payload.pop(req_id, None) - self._chunk_ready_req_ids.update(newly_finished) - - output = OmniConnectorOutput( - chunk_ready_req_ids=set(self._chunk_ready_req_ids), - chunk_finished_req_ids=chunk_finished, - request_metadata=request_metadata, - kv_sent_req_ids=list(self._kv_sent_req_ids), - stage_recv_req_ids=set(self._stage_recv_req_ids), - has_pending_kv_work=self.has_pending_kv_work(), - ) - if output.stage_recv_req_ids or chunk_finished or newly_finished: - logger.debug( - "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", - self._stage_id, - output.stage_recv_req_ids, - chunk_finished, - output.chunk_ready_req_ids, - ) - self._chunk_ready_req_ids.clear() - self._kv_sent_req_ids.clear() - self._stage_recv_req_ids.clear() - return output + drain_lock = getattr(self, "_omni_connector_output_drain_lock", None) + if drain_lock is None: + return self._drain_omni_connector_output() + with drain_lock: + return self._drain_omni_connector_output() @staticmethod def _connector_output_has_signals(output: OmniConnectorOutput) -> bool: @@ -854,6 +801,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self, request: Any, pooling_output: Any | None = None, + *, + propagate_errors: bool = False, + wait_for_delivery: bool = False, ) -> bool: """Derive and enqueue one chunk for async sending. @@ -882,10 +832,11 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): request_id=request_id, request=request, pooling_output=pooling_output, + propagate_errors=propagate_errors, ) if payload_data is None: if chunk_id == 0: - logger.warning( + logger.debug( "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)", self._stage_id, request_id, @@ -894,31 +845,14 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): ) return False - self._put_req_chunk[request_id] += 1 - self._ramp_chunk_count[request_id] += 1 - next_stage_id = self._next_stage_id - connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" - - if chunk_id == 0: - logger.debug( - "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", - self._stage_id, - request_id, - connector_put_key, - ) - - task = { - "stage_id": self._stage_id, - "next_stage_id": next_stage_id, - "put_key": connector_put_key, - "data": payload_data, - "request_id": request_id, - } - with self._lock: - self._pending_save_reqs.setdefault(request_id, deque()).append(task) - self._pending_save_counts[request_id] += 1 - self._work_available.set() - return True + enqueued, completion = self._enqueue_chunk_payload( + request, + payload_data, + wait_for_delivery=wait_for_delivery, + ) + if enqueued and completion is not None: + completion.wait() + return enqueued # ------------------------------------------------------------------ # # Background I/O threads @@ -946,14 +880,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): _recv_poll_count, ) - made_progress = False - for req_id in pending_ids: - if self._stop_event.is_set(): - break - try: - made_progress = self._poll_single_request(req_id) or made_progress - except Exception: - logger.warning("Error receiving data for %s", req_id, exc_info=True) + made_progress = self._poll_pending_requests_once(pending_ids) if not made_progress and not self._stop_event.is_set(): self._work_available.wait(timeout=0.005) @@ -977,22 +904,29 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if task is not None: success = False + send_error: BaseException | None = None try: success = self._send_single_request(task) - except Exception: + except Exception as error: + send_error = error logger.error( "Error saving data for %s", task.get("request_id"), exc_info=True, ) if not success: - self._requeue_or_drop_failed_send(task) + self._requeue_or_drop_failed_send(task, error=send_error) continue self._work_available.wait(timeout=0.01) self._work_available.clear() - def _requeue_or_drop_failed_send(self, task: dict) -> None: + def _requeue_or_drop_failed_send( + self, + task: dict, + *, + error: BaseException | None = None, + ) -> None: """Re-enqueue a failed send task or drop it after max retries.""" retry_count = task.get("_retry_count", 0) + 1 req_id = task.get("request_id") @@ -1002,6 +936,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): getattr(self, "_stage_id", "?"), ) return + if error is not None: + task["_last_send_error"] = error if retry_count <= self._MAX_SEND_RETRIES: task["_retry_count"] = retry_count logger.warning( @@ -1015,19 +951,36 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): dq = self._pending_save_reqs.setdefault(req_id, deque()) dq.appendleft(task) else: + cause = error or task.get("_last_send_error") + failure = RuntimeError( + "connector send failed permanently for " + f"request {req_id!r}, key {task.get('put_key')!r} " + f"after {self._MAX_SEND_RETRIES + 1} attempts" + ) + if cause is not None: + failure.__cause__ = cause logger.error( "[Stage-%s] Giving up on send for %s after %d retries", getattr(self, "_stage_id", "?"), req_id, self._MAX_SEND_RETRIES, ) + sink = getattr(self, "_connector_send_error_sink", None) + if sink is not None: + try: + sink(failure) + except Exception: + logger.exception("Connector send error sink failed for %s", req_id) + completion = task.get("completion") + if completion is not None: + completion.set_error(failure) self._decrement_pending_save_count(req_id) # ------------------------------------------------------------------ # # Chunk-level poll / send (ported from OmniChunkTransferAdapter) # ------------------------------------------------------------------ # - def _poll_single_request(self, req_id: str) -> bool: + def _poll_single_request(self, req_id: str, *, publish_ready: bool = True) -> bool: """Poll connector for one chunk of a request (non-blocking).""" connector = self._omni_connector if connector is None: @@ -1099,16 +1052,20 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): is_finished = self._payload_finished(payload_data) incoming_payload_consumable = self._payload_is_consumable(payload_data) - if self._model_mode == "ar": - payload_data = self._accumulate_payload(external_req_id, payload_data) - payload_consumable = incoming_payload_consumable - else: - new_ids = self._payload_audio_codes(payload_data) or [] - if not new_ids and not is_finished: + if self._model_mode != "ar": + new_ids = self._payload_audio_codes(payload_data) + if not self._payload_value_has_content(new_ids) and not is_finished: return False payload_consumable = self._payload_is_consumable(payload_data) with self._lock: + if self._model_mode == "ar": + # Accumulation, staging, and model-side consume/ack share + # this lock. Keeping the transition atomic prevents the + # model thread from clearing a just-merged decode span + # before the recv thread publishes it to the local cache. + payload_data = self._accumulate_payload(external_req_id, payload_data) + payload_consumable = incoming_payload_consumable if is_finished: self._chunk_finished_req_ids.add(req_id) self._chunk_stream_completed.add(req_id) @@ -1167,6 +1124,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): type(engine_inputs).__name__, ) + if publish_ready and self._async_chunk and (payload_consumable or is_finished): + self.publish_omni_connector_output_to_sink() logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key) return True @@ -1175,14 +1134,27 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): request_id: str | None, request: Any | None, pooling_output: Any | None, + *, + propagate_errors: bool = False, ) -> Any | None: """Run the custom process hook with a best-effort finished kwarg.""" if self._custom_process_func is None: return None + payload_kwarg = getattr(self, "_custom_process_payload_kwarg", None) + if payload_kwarg is None: + payload_kwarg = self._connector_payload_kwarg(self._custom_process_func) + self._custom_process_payload_kwarg = payload_kwarg + if payload_kwarg is None: + message = f"custom process hook {self._custom_process_func} has no connector payload argument" + logger.error(message) + if propagate_errors: + raise RuntimeError(message) + return None + kwargs = { "transfer_manager": self, - "pooling_output": pooling_output, + payload_kwarg: pooling_output, "request": request, } supports_is_finished = getattr( @@ -1197,21 +1169,29 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): kwargs["is_finished"] = bool(is_finished_fn()) except Exception: logger.debug("request.is_finished() failed for %s", request_id, exc_info=True) + if propagate_errors: + raise try: return self._custom_process_func(**kwargs) except TypeError as exc: if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc): logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + if propagate_errors: + raise return None kwargs.pop("is_finished", None) try: return self._custom_process_func(**kwargs) except Exception: logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + if propagate_errors: + raise return None except Exception: logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + if propagate_errors: + raise return None def _custom_process_supports_is_finished_kwarg(self) -> bool | None: @@ -1247,15 +1227,17 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): def _send_single_request(self, task: dict) -> bool: """Send one queued task via connector.put(). - Returns True on success. On failure (put() raises or returns - ``success=False``), returns False **without** decrementing - ``_pending_save_counts`` so the caller can retry or clean up. + Returns True when the task is consumed, either by a successful send or + because its delivery completion was already terminal. On transport + failure (put() raises or returns ``success=False``), returns False + **without** decrementing ``_pending_save_counts`` so the caller can + retry or clean up. """ connector = self._omni_connector if connector is None: return True - request_id = task.get("request_id") + request_id = task["request_id"] payload_data = task.get("data") if payload_data is None and task.get("request") is not None: payload_data = self._build_custom_process_payload( @@ -1264,6 +1246,15 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): pooling_output=task.get("pooling_output"), ) put_key = task.get("put_key") + completion = task.get("completion") + if completion is not None and not completion.mark_in_flight(): + logger.debug( + "[Stage-%s] Skipping terminal connector delivery: put_key=%s", + task["stage_id"], + put_key, + ) + self._decrement_pending_save_count(request_id) + return True success, _size, _metadata = connector.put( from_stage=str(task["stage_id"]), @@ -1282,8 +1273,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if not success: return False - if isinstance(request_id, str): - self._decrement_pending_save_count(request_id) + self._decrement_pending_save_count(request_id) + if completion is not None: + completion.set_result() return True def _decrement_pending_save_count(self, request_id: str) -> None: @@ -1302,6 +1294,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._put_req_chunk.pop(cleanup_req_id, None) self._send_side_request_payload.pop(cleanup_req_id, None) self._code_prompt_token_ids.pop(cleanup_req_id, None) + getattr(self, "_qwen3_tts_emitted_frames", {}).pop(cleanup_req_id, None) self._cached_ic.pop(cleanup_req_id, None) self._ramp_chunk_count.pop(cleanup_req_id, None) self._adaptive_states.pop(cleanup_req_id, None) @@ -1354,3 +1347,238 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._send_side_request_payload[req_id] = merged return cast(OmniPayload, dict(merged)) + + def _drain_omni_connector_output(self) -> OmniConnectorOutput: + tp_group = self._get_local_tp_group() + if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: + if self.is_data_transfer_rank(): + with self._lock: + fanout_packet = self._collect_async_chunk_fanout_packet_locked() + else: + fanout_packet = None + fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) + if fanout_packet is None: + newly_finished = set() + chunk_finished = set() + request_metadata = {} + else: + if not self.is_data_transfer_rank(): + self._apply_async_chunk_fanout_packet(fanout_packet) + newly_finished = set(fanout_packet["newly_finished"]) + chunk_finished = set(fanout_packet["chunk_finished"]) + request_metadata = dict(fanout_packet["request_metadata"]) + else: + with self._lock: + newly_finished = set(self._finished_load_reqs) + self._finished_load_reqs.clear() + chunk_finished = set(self._chunk_finished_req_ids) + self._chunk_finished_req_ids.clear() + request_metadata = dict(self._local_request_metadata) + self._local_request_metadata.clear() + # _send_side_request_payload is the async accumulation buffer for + # future recv chunks. Clearing it on every consumable wake-up drops + # intermediate + # thinker decode spans before the model side can consume them. + # Only terminal chunk_finished requests may release that buffer. + for req_id in chunk_finished: + if req_id not in self._local_stage_payload_cache: + continue + ext_req_id = self._request_ids_mapping.get(req_id, req_id) + self._send_side_request_payload.pop(ext_req_id, None) + if ext_req_id != req_id: + self._send_side_request_payload.pop(req_id, None) + self._chunk_ready_req_ids.update(newly_finished) + + output = OmniConnectorOutput( + chunk_ready_req_ids=set(self._chunk_ready_req_ids), + chunk_finished_req_ids=chunk_finished, + request_metadata=request_metadata, + kv_sent_req_ids=list(self._kv_sent_req_ids), + stage_recv_req_ids=set(self._stage_recv_req_ids), + has_pending_kv_work=self.has_pending_kv_work(), + ) + if output.stage_recv_req_ids or chunk_finished or newly_finished: + logger.debug( + "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", + self._stage_id, + output.stage_recv_req_ids, + chunk_finished, + output.chunk_ready_req_ids, + ) + self._chunk_ready_req_ids.clear() + self._kv_sent_req_ids.clear() + self._stage_recv_req_ids.clear() + return output + + def _enqueue_chunk_payload( + self, + request: Any, + payload_data: Any, + *, + wait_for_delivery: bool = False, + ) -> tuple[bool, Any | None]: + """Enqueue an already-built payload under its per-request chunk key.""" + raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) + request_id = self._resolve_external_req_id(request, raw_req_id) + if raw_req_id and raw_req_id != request_id: + self._request_ids_mapping.setdefault(raw_req_id, request_id) + chunk_id = self._put_req_chunk[request_id] + self._put_req_chunk[request_id] += 1 + self._ramp_chunk_count[request_id] += 1 + next_stage_id = self._next_stage_id + connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" + completion = ( + self._create_send_completion( + request_id=request_id, + put_key=connector_put_key, + ) + if wait_for_delivery + else None + ) + + if chunk_id == 0: + logger.debug( + "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", + self._stage_id, + request_id, + connector_put_key, + ) + + task = { + "stage_id": self._stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload_data, + "request_id": request_id, + "completion": completion, + } + with self._lock: + self._pending_save_reqs.setdefault(request_id, deque()).append(task) + self._pending_save_counts[request_id] += 1 + self._work_available.set() + return True, completion + + def _poll_pending_requests_once(self, pending_ids: list[str]) -> bool: + """Poll one receiver pass and publish its ready requests as a cohort.""" + made_progress = False + for req_id in pending_ids: + if self._stop_event.is_set(): + break + try: + made_progress = self._poll_single_request(req_id, publish_ready=False) or made_progress + except Exception: + logger.warning("Error receiving data for %s", req_id, exc_info=True) + if made_progress and self._async_chunk: + self.publish_omni_connector_output_to_sink() + return made_progress + + def _publish_chunk_cohort(self, entries: list[tuple[Any, Any]], *, wait_for_delivery: bool) -> int: + entries = [(request, payload) for request, payload in entries if payload is not None] + if not entries: + return 0 + emitted = 0 + completions: list[Any] = [] + for request, payload in entries: + enqueued, completion = self._enqueue_chunk_payload(request, payload, wait_for_delivery=wait_for_delivery) + if enqueued: + emitted += 1 + if completion is not None: + completions.append(completion) + for completion in completions: + completion.wait() + return emitted + + def publish_omni_connector_output_to_sink(self) -> bool: + """Drain pending readiness into the same-process scheduler inbox.""" + sink = getattr(self, "_omni_connector_output_sink", None) + if sink is None: + return False + output = self.get_omni_connector_output() + if not self._connector_output_has_signals(output): + return False + sink(output) + return True + + def send_chunks( + self, + entries: list[tuple[Any, Any | None]], + *, + propagate_errors: bool = False, + wait_for_delivery: bool = False, + ) -> int: + """Build and enqueue all request payloads from one model step. + + A model-specific batch builder may coalesce GPU post-processing and + D2H copies across requests. Models without one retain the scalar + builder contract. + """ + if not entries or self._omni_connector is None: + return 0 + if not self.is_data_transfer_rank(): + return len(entries) + + batch_func = getattr(self, "_custom_process_batch_func", None) + if batch_func is None: + payloads = [] + for request, pooling_output in entries: + raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) + request_id = self._resolve_external_req_id(request, raw_req_id) + payloads.append( + self._build_custom_process_payload( + request_id=request_id, + request=request, + pooling_output=pooling_output, + propagate_errors=propagate_errors, + ) + ) + return self._publish_chunk_cohort( + [(request, payload) for (request, _), payload in zip(entries, payloads, strict=True)], + wait_for_delivery=wait_for_delivery, + ) + + requests = [request for request, _ in entries] + pooling_outputs = [pooling_output for _, pooling_output in entries] + is_finished: list[bool] = [] + for request in requests: + finished_fn = getattr(request, "is_finished", None) + is_finished.append(bool(finished_fn()) if callable(finished_fn) else False) + + try: + payloads = batch_func( + transfer_manager=self, + pooling_outputs=pooling_outputs, + requests=requests, + is_finished=is_finished, + ) + except Exception: + logger.exception( + "batch custom_process_stage_input_func failed for %s requests", + len(entries), + ) + if propagate_errors: + raise + return 0 + if not isinstance(payloads, list) or len(payloads) != len(entries): + message = ( + "batch custom process hook " + f"{batch_func} returned " + f"{len(payloads) if isinstance(payloads, list) else type(payloads).__name__} " + f"payloads for {len(entries)} requests" + ) + logger.error( + "batch custom process hook %s returned %s payloads for %s requests", + batch_func, + len(payloads) if isinstance(payloads, list) else type(payloads).__name__, + len(entries), + ) + if propagate_errors: + raise RuntimeError(message) + return 0 + + return self._publish_chunk_cohort(list(zip(requests, payloads)), wait_for_delivery=wait_for_delivery) + + def set_omni_connector_output_sink( + self, + sink: Callable[[OmniConnectorOutput], None] | None, + ) -> None: + self._omni_connector_output_sink = sink diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 95a6cac22..2ed040703 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -9,6 +9,7 @@ import inspect import os import threading from collections import defaultdict, deque +from collections.abc import Callable from typing import TYPE_CHECKING, Any import torch @@ -21,6 +22,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ( ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id +from vllm_omni.outputs import OmniConnectorOutput logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") @@ -69,14 +71,35 @@ def should_accumulate_full_payload_output(model_config, custom_process_func) -> return False if getattr(model_config, "async_chunk", False): return False - if getattr(model_config, "final_output", False): - return False next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) if not isinstance(next_stage_func, str) or not next_stage_func: return False return getattr(model_config, "model_stage", None) is not None +class _SendCompletion: + """One connector task's delivery acknowledgement.""" + + def __init__(self) -> None: + self._event = threading.Event() + self._error: BaseException | None = None + + def mark_in_flight(self) -> bool: + return True + + def set_result(self) -> None: + self._event.set() + + def set_error(self, error: BaseException) -> None: + self._error = error + self._event.set() + + def wait(self) -> None: + self._event.wait() + if self._error is not None: + raise self._error + + class _OmniConnectorRuntimeMixin: """Own connector lifecycle, shared state, and KV transfer delegation.""" @@ -166,6 +189,8 @@ class _OmniConnectorRuntimeMixin: self._stage_id: int = stage_id if isinstance(stage_id, int) else 0 self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) + self._custom_process_batch_func = self._load_custom_batch_func(self._custom_process_func) + self._custom_process_payload_kwarg = self._connector_payload_kwarg(self._custom_process_func) self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() logger.debug( "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", @@ -229,6 +254,7 @@ class _OmniConnectorRuntimeMixin: self._pending_save_reqs: dict[str, deque] = {} self._pending_save_counts: dict[str, int] = defaultdict(int) self._deferred_send_cleanup: set[str] = set() + self._connector_send_error_sink: Callable[[BaseException], None] | None = None # -- per-cycle output accumulator -- self._chunk_ready_req_ids: set[str] = set() self._chunk_finished_req_ids: set[str] = set() @@ -247,6 +273,9 @@ class _OmniConnectorRuntimeMixin: self._local_stage_payload_cache: dict[str, dict[str, Any]] = {} # Lightweight scheduling metadata pending delivery to the Scheduler. self._local_request_metadata: dict[str, dict[str, Any]] = {} + # Optional same-process control-plane fast path. Payload tensors remain + # runner-owned; only OmniConnectorOutput readiness is published. + self._omni_connector_output_sink: Callable[[OmniConnectorOutput], None] | None = None # -- persistent set of request IDs whose chunk stream is complete -- # Prevents re-registration after the finish sentinel has been received. @@ -270,6 +299,7 @@ class _OmniConnectorRuntimeMixin: self._kv_triggered_requests: set[str] = set() self._lock = threading.Lock() + self._omni_connector_output_drain_lock = threading.Lock() self._stop_event = threading.Event() self._work_available = threading.Event() @@ -364,6 +394,7 @@ class _OmniConnectorRuntimeMixin: self._put_req_chunk.pop(k, None) self._send_side_request_payload.pop(k, None) self._code_prompt_token_ids.pop(k, None) + getattr(self, "_qwen3_tts_emitted_frames", {}).pop(k, None) self._cached_ic.pop(k, None) self._ramp_chunk_count.pop(k, None) self._adaptive_states.pop(k, None) @@ -384,6 +415,10 @@ class _OmniConnectorRuntimeMixin: self._cleanup_recv_delivery_state(req_id) def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None: + watermark = getattr(self, "_qwen3_tts_emitted_frames", {}) + watermark.pop(req_id, None) + if ext_id is not None: + watermark.pop(ext_id, None) if ext_id is not None: self._send_side_request_payload.pop(ext_id, None) self._cached_ic.pop(ext_id, None) @@ -594,31 +629,12 @@ class _OmniConnectorRuntimeMixin: return None, None - @staticmethod - def _is_connector_payload_builder(func: Any) -> bool: + @classmethod + def _is_connector_payload_builder(cls, func: Any) -> bool: """Whether *func* matches the mixin payload-builder contract.""" - try: - signature = inspect.signature(func) - except (TypeError, ValueError): - return False + return cls._connector_payload_kwarg(func) is not None - params = signature.parameters - if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): - return True - - required = {"transfer_manager", "pooling_output", "request"} - supported = { - name - for name, param in params.items() - if param.kind - in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - } - return required.issubset(supported) - - def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str: + def _resolve_external_req_id(self, request: Any, fallback_req_id: str | None) -> str: """Resolve the external request ID consistently. Checks ``_request_ids_mapping`` first (populated by @@ -626,7 +642,7 @@ class _OmniConnectorRuntimeMixin: ``external_req_id`` attribute, and finally to the given ``fallback_req_id``. """ - mapped = self._request_ids_mapping.get(fallback_req_id) + mapped = self._request_ids_mapping.get(fallback_req_id) if fallback_req_id is not None else None if mapped is not None: return mapped if request is not None: @@ -634,6 +650,8 @@ class _OmniConnectorRuntimeMixin: ext = getattr(request, "external_req_id", None) if ext is not None: return ext + if fallback_req_id is None: + raise ValueError("Connector request has neither an external nor an internal request ID") return fallback_req_id def _resolve_next_stage_id(self, model_config: Any) -> int: @@ -1073,3 +1091,67 @@ class _OmniConnectorRuntimeMixin: ) -> str: """Build connector key that includes rank info for KV transfers.""" return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" + + @staticmethod + def _connector_payload_kwarg(func: Any) -> str | None: + """Resolve the model output argument used by a connector builder. + + Full-payload builders historically name it ``pooling_output`` while + async-chunk builders name it ``multimodal_output``. Both represent the + runner-owned per-step inter-stage payload and share the same transport + contract. + """ + if func is None: + return None + try: + signature = inspect.signature(func) + except (TypeError, ValueError): + return None + + params = signature.parameters + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): + return "pooling_output" + + supported = { + name + for name, param in params.items() + if param.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + if not {"transfer_manager", "request"}.issubset(supported): + return None + if "pooling_output" in supported: + return "pooling_output" + if "multimodal_output" in supported: + return "multimodal_output" + return None + + @staticmethod + def _load_custom_batch_func(custom_process_func: Any | None) -> Any | None: + """Resolve an optional ``_batch`` peer.""" + if custom_process_func is None: + return None + module_name = getattr(custom_process_func, "__module__", None) + func_name = getattr(custom_process_func, "__name__", None) + if not module_name or not func_name: + return None + try: + module = importlib.import_module(module_name) + batch_func = getattr(module, f"{func_name}_batch", None) + except ImportError: + logger.debug( + "Could not import batch connector payload builder for %s.%s", + module_name, + func_name, + exc_info=True, + ) + return None + return batch_func if callable(batch_func) else None + + def _create_send_completion(self, *, request_id: str, put_key: str) -> Any: + """Create a delivery completion; MRv2 overrides this with a ticket.""" + del request_id, put_key + return _SendCompletion() diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index b801622f9..13d3dca3e 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -273,6 +273,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): frozen_ids = token_ids.copy() setattr(snapshot, private_name, frozen_ids) setattr(snapshot, public_name, ConstantList(frozen_ids)) + if public_name == "output_token_ids": + snapshot.output_token_count = len(frozen_ids) return snapshot @staticmethod @@ -878,6 +880,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: self.code_prompt_token_ids.pop(external_req_id, None) + getattr(self, "_qwen3_tts_emitted_frames", {}).pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) self._adaptive_states.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) @@ -992,6 +995,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.put_req_chunk.pop(external_req_id, None) self.request_payload.pop(external_req_id, None) self.code_prompt_token_ids.pop(external_req_id, None) + getattr(self, "_qwen3_tts_emitted_frames", {}).pop(external_req_id, None) self.requests_num_chunks_sent.pop(external_req_id, None) self._segment_generation.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index a96c8027f..c36d6d986 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Engine components for vLLM-Omni. """ @@ -134,10 +137,19 @@ class OmniEngineCoreOutput(EngineCoreOutput): # pooling_output is inherited from EngineCoreOutput as torch.Tensor | None # and retains its original vLLM semantics for pooling/embedding tasks. multimodal_output: dict[str, torch.Tensor] | None = None + # Bytes-serialized carrier for a dict-shaped pooling_output (MR V2 cross-stage + # hidden/codec handoff). vLLM strictly decodes ``pooling_output`` as a + # ``torch.Tensor``, so a per-request dict payload cannot ride that field. The + # producer (OmniARScheduler) sets ``pooling_output=None`` and serializes the + # dict here; the consumer (StagePool._poll_stage_raw) rehydrates it back into + # ``pooling_output`` after decode. Tensor pooling outputs are unaffected. + pooling_output_payload: AdditionalInformationPayload | None = None # Finished flag for streaming input segment is_segment_finished: bool | None = False # Streaming update prompt length new_prompt_len_snapshot: int | None = None + # Authoritative segment count when the native plane suppresses token IPC. + num_generation_tokens: int | None = None class OmniEngineCoreOutputs(EngineCoreOutputs): diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b6aba463b..939b19d50 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -190,6 +190,7 @@ class OmniEngineArgs(EngineArgs): model_stage: str = "thinker" model_arch: str | None = None engine_output_type: str | None = None + final_output: bool = False hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None requires_full_payload_input: bool = False @@ -199,6 +200,8 @@ class OmniEngineArgs(EngineArgs): async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + use_v2_model_runner: bool = False + supports_native_mrv2_data_plane: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. # Must be declared here so engine_args dict propagation does not silently # drop the value when constructing OmniEngineArgs from kwargs. @@ -423,6 +426,8 @@ class OmniEngineArgs(EngineArgs): async_chunk=self.async_chunk, session_mode=self.session_mode, retains_state_across_chunks=self.retains_state_across_chunks, + 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, duplex_max_sessions=self.duplex_max_sessions, model_stage=self.model_stage, @@ -430,6 +435,7 @@ class OmniEngineArgs(EngineArgs): worker_type=self.worker_type, pooling_output_decoder=self.pooling_output_decoder, engine_output_type=self.engine_output_type, + final_output=self.final_output, hf_config_name=self.hf_config_name, custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, requires_full_payload_input=self.requires_full_payload_input, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 277a48119..e69ebc7ae 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -192,6 +192,8 @@ class OrchestratorRequestState: final_stage_id: int = -1 final_output_stage_ids: set[int] = field(default_factory=set) finished_final_output_stage_ids: set[int] = field(default_factory=set) + finished_stage_ids: set[int] = field(default_factory=set) + pending_final_output: OutputMessage | None = None # Wall-clock timestamp when the client-facing engine request was accepted. request_timestamp: float = 0.0 @@ -763,20 +765,21 @@ class OrchestratorBase: float(kv_wait_s), ) req_state = self.request_states.get(getattr(eco, "request_id", None)) - if req_state is None or not req_state.streaming.enabled: + if req_state is None: continue - segment_finished = bool(getattr(eco, "is_segment_finished", False)) - raw_mm = self._completion_multimodal_output(eco, None) - req_state.streaming.segments[stage_id] = StreamingSegmentState( - finished=segment_finished, - token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []), - output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}), - ) - req_state.streaming.new_prompt_len_snapshot = getattr( - eco, - "new_prompt_len_snapshot", - None, - ) + if req_state.streaming.enabled: + segment_finished = bool(getattr(eco, "is_segment_finished", False)) + raw_mm = self._completion_multimodal_output(eco, None) + req_state.streaming.segments[stage_id] = StreamingSegmentState( + finished=segment_finished, + token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []), + output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}), + ) + req_state.streaming.new_prompt_len_snapshot = getattr( + eco, + "new_prompt_len_snapshot", + None, + ) if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): raw_terminal_request_ids.add(req_state.request_id) await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state) @@ -1506,18 +1509,19 @@ class OrchestratorBase: ``is_segment_finished=False``, but vLLM's output processor may remove the request state before that EngineCoreOutput is processed. - Only update ``finished_final_output_stage_ids`` here. Request cleanup stays - in ``_route_output`` so downstream async-chunk stages can still deliver - outputs after stage-0 session end. + Record raw stage completion here. Client completion is + resolved after this raw batch passes through the output processor, so a + real processed output wins over the swallowed-terminal fallback. """ if getattr(eco, "finish_reason", None) is None: return False if getattr(eco, "is_segment_finished", False): return False + req_state.finished_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if stage_id not in final_output_stage_ids: - return False + return True req_state.finished_final_output_stage_ids.add(stage_id) return True @@ -1570,16 +1574,23 @@ class OrchestratorBase: replica_id: int, request_ids: set[str], ) -> None: - """Finish streaming requests whose raw terminal had no processed output.""" + """Finish requests whose raw terminal had no processed output.""" pool = self.stage_pools[stage_id] - if not pool.final_output: - return - for request_id in request_ids: req_state = self.request_states.get(request_id) if req_state is None or req_state.session_owned: continue - + pending = req_state.pending_final_output + if pending is not None: + if set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids): + req_state.pending_final_output = None + await self.output_async_queue.put(pending) + await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)]) + # A real terminal output is pending; do not replace it with + # the swallowed-output fallback before upstream completion. + continue + if not pool.final_output: + continue final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if not final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids): continue @@ -1653,27 +1664,41 @@ class OrchestratorBase: await self._cleanup_request_ids([req_id]) return + if finished and not segment_finished: + req_state.finished_stage_ids.add(stage_id) + request_finished = False if finished and self.stage_pools[stage_id].final_output and not segment_finished: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) if req_state.session_owned: - # Session-owned outputs are delivered to their session through - # ``_intercept_stage_output`` below, never to a request queue. + # Session-owned outputs are delivered through the session interceptor. pass elif self.stage_pools[stage_id].final_output: - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=stage_id, - replica_id=replica_id, - engine_outputs=output, - metrics=stage_metrics, - finished=request_finished, - stage_submit_ts=submit_ts, - ) + message = OutputMessage( + request_id=req_id, + stage_id=stage_id, + replica_id=replica_id, + engine_outputs=output, + metrics=stage_metrics, + finished=request_finished, + stage_submit_ts=submit_ts, ) + if ( + request_finished + and self.async_chunk + and not req_state.streaming.enabled + and not set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids) + ): + # Direct chunk delivery can finish the decoder before the + # upstream engine's terminal output/metrics reaches this loop. + # Keep request state until those messages are routed, otherwise + # the frontend publishes incomplete usage and loses stop reasons. + req_state.pending_final_output = message + request_finished = False + else: + await self.output_async_queue.put(message) elif stage_metrics is not None: await self.output_async_queue.put( StageMetricsMessage( @@ -1741,6 +1766,12 @@ class OrchestratorBase: is_final_update=True, ) + pending = req_state.pending_final_output + if pending is not None and set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids): + req_state.pending_final_output = None + await self.output_async_queue.put(pending) + request_finished = True + if request_finished and not req_state.session_owned: await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) @@ -2483,8 +2514,6 @@ class OrchestratorBase: operation="async-chunk prewarm", ) else: - import copy - from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length try: @@ -2507,7 +2536,9 @@ class OrchestratorBase: original_prompt = req_state.prompt if isinstance(original_prompt, dict): - base_input = copy.deepcopy(original_prompt) + from vllm_omni.engine.request_snapshot import copy_request_snapshot + + base_input = copy_request_snapshot(original_prompt) else: base_input = {} diff --git a/vllm_omni/engine/request_snapshot.py b/vllm_omni/engine/request_snapshot.py new file mode 100644 index 000000000..4db02f4c1 --- /dev/null +++ b/vllm_omni/engine/request_snapshot.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Full request snapshots with a fast path for immutable scalar leaves.""" + +import copy +from typing import Any + +_ATOMIC_TYPES = frozenset((type(None), bool, int, float, complex, str, bytes)) + + +def copy_request_snapshot(value: Any, memo: dict[int, Any] | None = None) -> Any: + """Deep-copy all fields, preserving aliases and isolating mutable containers. + + Plain numeric/text trees do not need generic object dispatch or an identity + lookup for each immutable scalar. Subclasses and other objects retain their + regular deepcopy contracts. No model fields are filtered or shared here. + """ + kind = type(value) + if kind in _ATOMIC_TYPES: + return value + if memo is None: + memo = {} + key = id(value) + if key in memo: + return memo[key] + if kind is list: + atomic = all(type(item) in _ATOMIC_TYPES for item in value) + result = value.copy() if atomic else [] + memo[key] = result + if not atomic: + result.extend(copy_request_snapshot(item, memo) for item in value) + elif kind is dict: + result = {} + memo[key] = result + for k, item in value.items(): + result[copy_request_snapshot(k, memo)] = copy_request_snapshot(item, memo) + else: + return copy.deepcopy(value, memo) + # Keep sources alive for the duration of a copy, including custom deepcopy + # implementations that mutate a source container while being visited. + memo.setdefault(id(memo), []).append(value) + return result diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 5fdd74912..e87d19bb4 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -19,7 +19,7 @@ import vllm.v1.engine as _vllm_engine_module import vllm.v1.engine.core_client as _vllm_core_client_module from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest -from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient +from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient, MPClient from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.distributed.omni_connectors.utils.config import ( @@ -232,6 +232,9 @@ class StageEngineCoreClientBase(StageClientBase): if self.resources.engine_dead: raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead") + def _apply_ready_response(self, payload: bytes) -> None: + MPClient._apply_ready_response(self, payload) + # ==================== Overrides ==================== async def add_request_async(self, request: EngineCoreRequest) -> None: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index a52f58925..839e52495 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Stage Core Process for vLLM-Omni V1 architecture. @@ -27,6 +30,7 @@ from vllm.v1.engine.utils import ( EngineZmqAddresses, SignalCallback, ) +from vllm.v1.executor.uniproc_executor import UniProcExecutor from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest @@ -75,6 +79,24 @@ def _signal_exit_code(signum: int) -> int: return _SIGNAL_EXIT_BASE + signum +def _bind_native_data_plane_ready_sink(model_executor: Any, scheduler: Any) -> bool: + """Bind the TP1 in-process runner control plane directly to its scheduler.""" + if not isinstance(model_executor, UniProcExecutor): + return False + parallel_config = model_executor.vllm_config.parallel_config + if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1: + return False + driver_worker = getattr(model_executor, "driver_worker", None) + worker = getattr(driver_worker, "worker", None) + model_runner = getattr(worker, "model_runner", None) + data_plane = getattr(model_runner, "_omni_data_plane", None) + sink = getattr(scheduler, "enqueue_omni_connector_output", None) + if data_plane is None or not callable(sink): + return False + data_plane.set_omni_connector_output_sink(sink) + return True + + class StageEngineCoreProc(EngineCoreProc): """Stage-specific engine core process for vLLM-Omni. @@ -83,6 +105,11 @@ class StageEngineCoreProc(EngineCoreProc): ``EngineCoreProc.run_engine_core()``. """ + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler): + logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.") + def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" scheduler_request, current_wave = super().preprocess_add_request(request) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 0ab6c1da7..46a34dfdb 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -15,6 +15,7 @@ from vllm.logger import init_logger from vllm.v1.engine import EngineCoreOutputs from vllm.v1.metrics.stats import IterationStats +from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_coordinator import ( LoadBalancer, OmniCoordClientForHub, @@ -22,6 +23,7 @@ from vllm_omni.distributed.omni_coordinator import ( ReplicaStatus, ) from vllm_omni.distributed.omni_coordinator.load_balancer import Task +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_client import ( StagePoolClient, StagePoolDiffusionClient, @@ -1136,8 +1138,30 @@ class StagePool: # gauges for that interval. if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests: return None + self._rehydrate_pooling_output_payloads(outputs) return outputs + @staticmethod + def _rehydrate_pooling_output_payloads(outputs: EngineCoreOutputs) -> None: + """Restore dict-shaped pooling_output from its bytes carrier (MR V2). + + vLLM decodes EngineCoreOutput.pooling_output as a torch.Tensor, so MR + V2 runners ship the per-request dict handoff via ``pooling_output_payload`` + with ``pooling_output=None``. Decode it back here so downstream stage-input + processors see the same ``pooling_output`` shape as the legacy runner. + """ + for eco in outputs.outputs: + payload = getattr(eco, "pooling_output_payload", None) + if payload is None: + continue + if getattr(eco, "pooling_output", None) is None: + # deserialize_* returns the nested OmniPayload form; the producer + # serialized a flat (dotted-key) pooler dict, and downstream + # consumers (e.g. talker2code2wav_full_payload) read flat keys like + # "codes.audio". Re-flatten to restore the exact on-wire shape. + eco.pooling_output = flatten_payload(deserialize_additional_information(payload)) + eco.pooling_output_payload = None + async def process_llm_raw_outputs( self, replica_id: int, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index a348fdd09..d3c23ca39 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -110,8 +110,8 @@ _REF_AUDIO_MIN_DURATION = 1.0 # seconds _REF_AUDIO_MAX_DURATION = 30.0 # seconds _REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3 _REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"}) -_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256 -_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024 +_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 1024 +_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 512 * 1024 * 1024 _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _DEFAULT_VOICE_NAME = "default" @@ -1359,10 +1359,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_resolve_cache.move_to_end(cache_key) return cached[3] - def _put_resolved_ref_audio(self, cache_key: str, waveform: np.ndarray, sr: int, artifact_key: str) -> None: + def _put_resolved_ref_audio( + self, cache_key: str, waveform: np.ndarray | list[float], sr: int, artifact_key: str + ) -> None: if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0: return - size = waveform.nbytes + # Keep the finalized ndarray's ownership and identity for the array + # resolver. Legacy list callers are materialized into a compact array; + # their Python list remains independent from cached storage. + waveform = np.asarray(waveform, dtype=np.float32) + if not waveform.flags.c_contiguous: + waveform = np.ascontiguousarray(waveform) + size = int(waveform.nbytes) if size > self._ref_audio_resolve_cache_max_bytes: return previous = self._ref_audio_resolve_cache.pop(cache_key, None) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index eff3305d5..306301c94 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -682,14 +682,20 @@ class CodePredictorWrapper(nn.Module): self._lm_heads_list: list[nn.Module] | None = None self._codec_embeds_list: list[nn.Module] | None = None self._device_graphs: dict[int | tuple[int, int], tuple] = {} # (graph, static_output) per bucket + # Outer-runner execution buckets when a model host declares them (MRv2); + # ``None`` keeps the legacy power-of-two bucket derivation. + self._execution_batch_buckets: list[int] | None = None prefix_graph_cfg = self._stage_connector_extra_config(vllm_config) prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs")) is_npu = current_omni_platform.is_npu() - self._prefix_graphs_enabled = prefix_graphs_requested and wrapper_config.use_cuda_graphs and not is_npu - if prefix_graphs_requested and not self._prefix_graphs_enabled: + # MRv2 captures the whole Talker MTP call. It can therefore capture + # shorter compiled re-prefill forwards without nesting predictor-owned + # CUDA graphs inside the outer graph. + self._prefix_reprefill_enabled = prefix_graphs_requested and not is_npu + self._prefix_graphs_enabled = self._prefix_reprefill_enabled and wrapper_config.use_cuda_graphs + if prefix_graphs_requested and not self._prefix_reprefill_enabled: logger.info_once( - "code_predictor: prefix CUDA graphs requested but disabled because use_cuda_graphs=%s is_npu=%s", - wrapper_config.use_cuda_graphs, + "code_predictor: prefix re-prefill requested but disabled because is_npu=%s", is_npu, ) self._prefix_graph_buckets = self._parse_positive_int_set( @@ -698,6 +704,7 @@ class CodePredictorWrapper(nn.Module): self._prefix_graph_seq_lens = self._parse_positive_int_set( prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens") ) + self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1)) def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() @@ -761,7 +768,8 @@ class CodePredictorWrapper(nn.Module): dynamic=False, options={"epilogue_fusion": False}, ) - self._warmup_buckets() + with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()): + self._warmup_buckets() if self._wrapper_config.use_cuda_graphs: self._capture_cuda_graphs() @@ -770,7 +778,7 @@ class CodePredictorWrapper(nn.Module): logger.info("code_predictor: torch.compile (dynamic=False, no epilogue fusion)") def _padded_bsz(self, bsz: int) -> int: - """Round batch size up to nearest power-of-2 bucket.""" + """Round batch size up to the nearest warmed bucket.""" for bucket in self._bucket_sizes: if bsz <= bucket: return bucket @@ -832,15 +840,28 @@ class CodePredictorWrapper(nn.Module): return row_generators @classmethod - def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor: - """Sample ``logits`` via Gumbel-max with optional per-row generators.""" - row_generators = cls._normalize_generators(generator, int(logits.shape[0])) - u = torch.empty_like(logits, dtype=torch.float32) - if isinstance(row_generators, list): - for row, row_generator in enumerate(row_generators): - u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator) + def _sample_codes_gumbel( + cls, + logits: torch.Tensor, + generator: _GeneratorLike = None, + uniforms: torch.Tensor | None = None, + ) -> torch.Tensor: + """Sample ``logits`` via Gumbel-max with optional precomputed noise.""" + if uniforms is None: + row_generators = cls._normalize_generators(generator, int(logits.shape[0])) + u = torch.empty_like(logits, dtype=torch.float32) + if isinstance(row_generators, list): + for row, row_generator in enumerate(row_generators): + u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator) + else: + u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators) else: - u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators) + if uniforms.shape != logits.shape: + raise ValueError( + f"precomputed sampling uniforms must match logits: uniforms={tuple(uniforms.shape)} " + f"logits={tuple(logits.shape)}" + ) + u = uniforms return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) def _prefix_seq_lens(self, max_seq: int) -> list[int]: @@ -850,13 +871,64 @@ class CodePredictorWrapper(nn.Module): allowed = set(all_seq_lens) return sorted(seq_len for seq_len in self._prefix_graph_seq_lens if seq_len in allowed) - def _warmup_buckets(self) -> None: - """Warmup power-of-2 batch-size buckets to front-load Inductor compilation.""" + @staticmethod + def _synchronize_warmup(device: torch.device) -> None: + if device.type != "cpu": + current_omni_platform.synchronize() + + def configure_mtp_execution_buckets(self, sizes: Iterable[int]) -> None: + """Declare the batch buckets the outer runner can actually reach. + + Must be called before the first warmup; once buckets are captured the + existing set is only acceptable when it already covers the declaration + (never clear active graphs to rebuild them). + """ + max_bsz = int(self._vllm_config.scheduler_config.max_num_seqs) + cleaned = sorted({int(size) for size in sizes if 0 < int(size) <= max_bsz} | {max_bsz}) + if self._bucket_sizes: + missing = [size for size in cleaned if size not in self._bucket_sizes] + if missing: + raise RuntimeError( + "mtp execution buckets must be configured before the first warmup; " + f"missing={missing} existing={self._bucket_sizes}" + ) + return + self._execution_batch_buckets = cleaned + + def _batch_bucket_sizes(self) -> list[int]: max_bsz = self._vllm_config.scheduler_config.max_num_seqs + if self._execution_batch_buckets is not None: + # Outer-reachable buckets plus prefix-graph capture buckets and the + # full-size eager fallback row already merged at configure time. + bucket_sizes = list(self._execution_batch_buckets) + bucket_sizes.extend(bucket for bucket in self._prefix_graph_buckets if bucket <= max_bsz) + return sorted(set(bucket_sizes)) bucket_sizes = [1 << i for i in range(max_bsz.bit_length()) if (1 << i) <= max_bsz] - if max_bsz not in bucket_sizes: - bucket_sizes.append(max_bsz) - self._bucket_sizes = sorted(bucket_sizes) + bucket_sizes.extend(bucket for bucket in self._prefix_graph_buckets if bucket <= max_bsz) + bucket_sizes.append(max_bsz) + return sorted(set(bucket_sizes)) + + def _compile_cache_size_limit(self) -> int: + bucket_sizes = self._batch_bucket_sizes() + if not self._prefix_reprefill_enabled: + required_entries = len(bucket_sizes) + else: + prefix_buckets = ( + self._prefix_graph_buckets.intersection(bucket_sizes) + if self._prefix_graph_buckets + else set(bucket_sizes) + ) + prefix_seq_lens = self._prefix_reprefill_seq_lens + needs_full_graph = set(prefix_seq_lens) != set(range(2, self._num_groups + 1)) + full_graph_entries = len(bucket_sizes) - len(prefix_buckets) + if needs_full_graph: + full_graph_entries += len(prefix_buckets) + required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens) + return max(torch._dynamo.config.cache_size_limit, required_entries) + + def _warmup_buckets(self) -> None: + """Warm up batch-size buckets to front-load Inductor compilation.""" + self._bucket_sizes = self._batch_bucket_sizes() max_seq = self._num_groups + 1 device = next(self.model.parameters()).device @@ -866,8 +938,8 @@ class CodePredictorWrapper(nn.Module): self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes)) proj_buf = self._proj_buf - if self._prefix_graphs_enabled: - prefix_seq_lens = self._prefix_seq_lens(max_seq) + if self._prefix_reprefill_enabled: + prefix_seq_lens = self._prefix_reprefill_seq_lens needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq)) for bsz in self._bucket_sizes: capture_prefixes = not self._prefix_graph_buckets or bsz in self._prefix_graph_buckets @@ -905,6 +977,10 @@ class CodePredictorWrapper(nn.Module): self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids) logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes) + # Compiled attention warmup is asynchronous. Complete it before graph + # capture and serving reuse the same static buffers. + self._synchronize_warmup(device) + def _capture_cuda_graphs(self) -> None: """Capture a CUDA graph per bucket using vLLM's global graph pool.""" from vllm.platforms import current_platform @@ -914,7 +990,7 @@ class CodePredictorWrapper(nn.Module): proj_buf = self._proj_buf if self._prefix_graphs_enabled: - prefix_seq_lens = self._prefix_seq_lens(max_seq) + prefix_seq_lens = self._prefix_reprefill_seq_lens needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq)) for bsz in self._bucket_sizes: capture_prefixes = not self._prefix_graph_buckets or bsz in self._prefix_graph_buckets @@ -1013,13 +1089,20 @@ class CodePredictorWrapper(nn.Module): top_p: float = 1.0, generator: torch.Generator | None = None, generators: Sequence[torch.Generator | None] | None = None, + sample_uniforms: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """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)}" + ) sample_generator: _GeneratorLike = generators if generators is not None else generator - num_groups = self._num_groups device = layer0_code.device # _setup_compile caches _model_dtype on first call; use it for buffers @@ -1072,11 +1155,16 @@ class CodePredictorWrapper(nn.Module): for step in range(1, num_groups): graph_key: int | tuple[int, int] = padded_bsz seq_len = max_seq - if self._prefix_graphs_enabled: - prefix_key = (padded_bsz, step + 1) - if prefix_key in self._device_graphs: - graph_key = prefix_key - seq_len = step + 1 + if self._prefix_reprefill_enabled: + actual_seq_len = step + 1 + for prefix_seq_len in self._prefix_reprefill_seq_lens: + if prefix_seq_len < actual_seq_len: + continue + prefix_key = (padded_bsz, prefix_seq_len) + if prefix_key in self._bucket_pos_ids: + graph_key = prefix_key + seq_len = prefix_seq_len + break pos_ids = self._bucket_pos_ids.get(graph_key) if pos_ids is None: pos_ids = ( @@ -1123,7 +1211,8 @@ class CodePredictorWrapper(nn.Module): remove_mask = (cumulative_probs - sorted_probs) >= s_top_p sorted_logits[remove_mask] = float("-inf") logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) - code = self._sample_codes_gumbel(logits, generator=sample_generator) + step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None + code = self._sample_codes_gumbel(logits, generator=sample_generator, uniforms=step_uniforms) else: # "per_call" mode: temperature-scaled + top-k -> Gumbel-max if use_sampling: @@ -1131,7 +1220,8 @@ class CodePredictorWrapper(nn.Module): if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - code = self._sample_codes_gumbel(scaled, generator=sample_generator) + step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None + code = self._sample_codes_gumbel(scaled, generator=sample_generator, uniforms=step_uniforms) else: code = logits.argmax(dim=-1, keepdim=True) diff --git a/vllm_omni/model_executor/models/output_templates.py b/vllm_omni/model_executor/models/output_templates.py index 9d6a84ac4..f8848cf51 100644 --- a/vllm_omni/model_executor/models/output_templates.py +++ b/vllm_omni/model_executor/models/output_templates.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from typing import NamedTuple import torch @@ -6,6 +9,18 @@ from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload +class OwnedBatchTensor(NamedTuple): + """A batch tensor whose producer transfers ownership to the buffer. + + The producer must have freshly allocated the tensor (e.g. an index_select + result) and must not write it afterwards. The buffer keeps row views + instead of taking a second snapshot. Tensors without this marker are + treated as borrowed (graph output / shared scratch) and still get a clone. + """ + + tensor: torch.Tensor + + class OmniOutput(NamedTuple): """Output from the merged Omni model containing both text and audio.""" diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 30fc504b3..f7542d3e3 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -179,6 +179,7 @@ class Qwen3OmniMoeForConditionalGeneration( # vLLM run with no talker stage downstream, so no one consumes captured # thinker layers and the forward must return what stock vLLM expects. self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None + self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run if self.model_stage == "thinker": self.use_async_omni_output = True diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index 36bf663dd..e8776ed40 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -533,6 +536,10 @@ class Qwen3TTSConfig(PretrainedConfig): if any("Code2Wav" in str(a) for a in archs): if hasattr(config, "rope_parameters"): delattr(config, "rope_parameters") + # Code2Wav doesn't use position embeddings, but vLLM validates + # max_model_len <= max_position_embeddings. Set a large value + # to allow the pipeline.yaml max_model_len (65536) to pass validation. + config.max_position_embeddings = 131072 return config diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index abeaf243d..6aa56ab09 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py @@ -1,5 +1,6 @@ # Copyright 2026 The Alibaba Qwen team. # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder. @@ -8,10 +9,8 @@ reducing kernel launch overhead during inference. """ import bisect -import os import time -from collections import Counter -from collections.abc import Callable, Sequence +from collections.abc import Callable, Collection, Sequence import torch from torch.cuda import CUDAGraph @@ -187,22 +186,6 @@ class CUDAGraphDecoderWrapper: self._warmed_up = False self._device = None - self._stats_enabled = os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", "").lower() in ( - "1", - "true", - "yes", - "on", - ) - self._stats_log_every = int(os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY", "0") or 0) - self._stats_total = 0 - self._stats_hits = 0 - self._stats_compiled_hits = 0 - self._stats_fallbacks = 0 - self._stats_stream_capture_fallbacks = 0 - self._stats_requests: Counter[tuple[int, int]] = Counter() - self._stats_hit_shapes: Counter[tuple[int, int, int]] = Counter() - self._stats_compiled_shapes: Counter[tuple[int, int]] = Counter() - self._stats_fallback_shapes: Counter[tuple[int, int, int]] = Counter() @staticmethod def compute_capture_sizes( @@ -244,57 +227,19 @@ class CUDAGraphDecoderWrapper: shapes.update(self.extra_capture_shapes) return sorted(shapes) - def _record_decode_stats( - self, - *, - hit: bool, - batch_size: int, + @staticmethod + def _select_replay_shape( + actual_batch_size: int, actual_size: int, - padded_size: int | None, - stream_capture: bool = False, - compiled: bool = False, - ) -> None: - if not self._stats_enabled: - return - - padded_key = int(padded_size) if padded_size is not None else -1 - self._stats_total += 1 - self._stats_requests[(batch_size, actual_size)] += 1 - if hit: - self._stats_hits += 1 - if compiled: - self._stats_compiled_hits += 1 - self._stats_compiled_shapes[(batch_size, actual_size)] += 1 - else: - self._stats_hit_shapes[(batch_size, actual_size, padded_key)] += 1 - else: - self._stats_fallbacks += 1 - self._stats_fallback_shapes[(batch_size, actual_size, padded_key)] += 1 - if stream_capture: - self._stats_stream_capture_fallbacks += 1 - - if self._stats_log_every > 0 and self._stats_total % self._stats_log_every == 0: - self.log_decode_stats() - - def log_decode_stats(self) -> None: - if not self._stats_enabled or self._stats_total == 0: - return - hit_rate = 100.0 * self._stats_hits / self._stats_total - logger.info( - "Code2Wav CUDA Graph stats: total=%d hits=%d fallbacks=%d " - "compiled_hits=%d stream_capture_fallbacks=%d hit_rate=%.2f%% " - "top_requests=%s top_compiled=%s top_hits=%s top_fallbacks=%s", - self._stats_total, - self._stats_hits, - self._stats_fallbacks, - self._stats_compiled_hits, - self._stats_stream_capture_fallbacks, - hit_rate, - self._stats_requests.most_common(12), - self._stats_compiled_shapes.most_common(12), - self._stats_hit_shapes.most_common(12), - self._stats_fallback_shapes.most_common(12), + available_shapes: Collection[tuple[int, int]], + ) -> tuple[int, int] | None: + """Select the lowest-cost captured shape that contains the request.""" + candidates = ( + (batch_size, size) + for batch_size, size in available_shapes + if batch_size >= actual_batch_size and size >= actual_size ) + return min(candidates, key=lambda shape: (shape[0] * shape[1], shape[0], shape[1]), default=None) def warmup( self, @@ -486,68 +431,42 @@ class CUDAGraphDecoderWrapper: # outside the startup capture window, so normal inference still hits # the graph fast path. if torch.cuda.is_current_stream_capturing(): - self._record_decode_stats( - hit=False, - batch_size=int(codes.shape[0]), - actual_size=int(codes.shape[-1]), - padded_size=None, - stream_capture=True, - ) return self.decoder(codes) batch_size = int(codes.shape[0]) actual_size = int(codes.shape[-1]) - padded_size = self._get_padded_size(actual_size) - compile_key = (batch_size, actual_size) - if compile_key not in self._compiled_shapes and padded_size is not None: - compile_key = (batch_size, padded_size) - if compile_key in self._compiled_shapes: + compile_key = self._select_replay_shape(batch_size, actual_size, self._compiled_shapes) + if compile_key is not None: compiled_size = compile_key[1] - self._record_decode_stats( - hit=True, - batch_size=batch_size, - actual_size=actual_size, - padded_size=compiled_size, - compiled=True, - ) static_input = self._compiled_static_inputs[compile_key] - if actual_size == compiled_size: + if actual_size == compiled_size and batch_size == compile_key[0]: static_input.copy_(codes) else: static_input.zero_() - static_input[:, :, :actual_size] = codes + static_input[:batch_size, :, :actual_size] = codes self._compiled_graphs[compile_key].replay() output = self._trim_replay_output(self._compiled_static_outputs[compile_key], actual_size, compiled_size) + output = output[:batch_size] if clone_graph_output: return output.clone() return output - graph_key = (batch_size, padded_size) if padded_size is not None else None + graph_key = self._select_replay_shape(batch_size, actual_size, self.graphs) - if graph_key is None or graph_key not in self.graphs: - self._record_decode_stats( - hit=False, - batch_size=batch_size, - actual_size=actual_size, - padded_size=padded_size, - ) + if graph_key is None: return self.decoder(codes) - self._record_decode_stats( - hit=True, - batch_size=batch_size, - actual_size=actual_size, - padded_size=padded_size, - ) static_input = self.static_inputs[graph_key] - if actual_size == padded_size: + graph_size = graph_key[1] + if actual_size == graph_size and batch_size == graph_key[0]: static_input.copy_(codes) else: static_input.zero_() - static_input[:, :, :actual_size] = codes + static_input[:batch_size, :, :actual_size] = codes self.graphs[graph_key].replay() - output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, padded_size) + output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, graph_size) + output = output[:batch_size] if clone_graph_output: return output.clone() return output diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index dd33aeca8..dd753861d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen3-TTS pipeline: Talker (text → RVQ codec) → Code2Wav (codec → audio). Chunked vs end-to-end mode is dispatched from ``deploy.async_chunk``. @@ -41,6 +41,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig( owns_tokenizer=True, engine_output_type="latent", async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), + supports_native_mrv2_data_plane=True, custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": False, @@ -56,6 +57,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="Qwen3TTSCode2Wav", + supports_native_mrv2_data_plane=True, # ``sync_process_input_func`` is the only input-proc override for # this stage in sync (non-async-chunk) mode: a length-only # ``_token_only`` placeholder. The bulk codec payload itself 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 609117ce3..45d887ee2 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 @@ -198,10 +198,6 @@ def mel_spectrogram( hann_window: torch.Tensor | None = None, ) -> torch.Tensor: """Mel spectrogram via torch STFT and a (cached) mel filterbank.""" - if torch.min(y) < -1.0: - logger.warning("Min value of input waveform signal is %s", torch.min(y)) - if torch.max(y) > 1.0: - logger.warning("Max value of input waveform signal is %s", torch.max(y)) device = y.device if mel_basis is None: mel_basis = _cached_mel_filter_bank(sampling_rate, n_fft, num_mels, fmin, fmax).to(device) @@ -350,7 +346,7 @@ class Qwen3TTSPromptEmbedsBuilder: tts_pad_embed: torch.Tensor, encode_ref_audio_batch: Callable[..., list[torch.Tensor]], speaker_cache: Any | None = None, - ref_audio_artifact_cache_max_entries: int = 256, + ref_audio_artifact_cache_max_entries: int = 1024, ): self._config = config self._talker_config = talker_config @@ -366,9 +362,14 @@ class Qwen3TTSPromptEmbedsBuilder: self._embedding_dtype = torch.bfloat16 self._text_tokenizer: Any | None = None + # Projected special-token embeddings are immutable model constants. + # 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] = {} 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() + self._long_tensor_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {} # Bounded LRU; caller-supplied orig_sr can otherwise grow this without limit. self._resampler_cache: OrderedDict[tuple[int, int], AudioResampler] = OrderedDict() @@ -384,6 +385,29 @@ class Qwen3TTSPromptEmbedsBuilder: """Return the pad-token embedding on the requested device/dtype, shape ``[1, 1, H]``.""" return self._tts_pad_embed_buffer.to(device=device, dtype=dtype).reshape(1, 1, -1) + def _long_tensor(self, values: Sequence[int], device: torch.device) -> torch.Tensor: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.accelerator.current_device_index()) + key = (str(device), tuple(int(v) for v in values)) + cache = getattr(self, "_long_tensor_cache", None) + if cache is None: + cache = self._long_tensor_cache = {} + cached = cache.get(key) + if cached is None or cached.device != device: + cached = torch.tensor([list(key[1])], device=device, dtype=torch.long) + cache[key] = cached + return cached + + def _projected_tokens(self, values: Sequence[int], device: torch.device) -> torch.Tensor: + device = torch.device(device) + key = (str(device), tuple(int(v) for v in values)) + cached = self._projected_token_cache.get(key) + if cached is None: + cached = self._text_projection(self._text_embedding(self._long_tensor(values, device))) + self._projected_token_cache[key] = cached + return cached + def _get_resampler(self, orig_sr: int, target_sr: int) -> AudioResampler: key = (int(orig_sr), int(target_sr)) cache = self._resampler_cache @@ -631,9 +655,9 @@ class Qwen3TTSPromptEmbedsBuilder: if entry is None: entry = {} if isinstance(ref_code, torch.Tensor): - entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous() + entry["ref_code"] = ref_code.detach().to(dtype=torch.long).contiguous() if isinstance(ref_spk_embedding, torch.Tensor): - entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=self._embedding_dtype).reshape(-1) + entry["ref_spk_embedding"] = ref_spk_embedding.detach().to(dtype=self._embedding_dtype).reshape(-1) if not entry: return self._ref_audio_artifact_cache[cache_key] = entry @@ -659,7 +683,6 @@ class Qwen3TTSPromptEmbedsBuilder: self._speaker_encoder.to(device=dev, dtype=dtype) except StopIteration: pass - # Resample to 24kHz for speaker encoder. target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000)) if sr != target_sr: @@ -761,6 +784,23 @@ class Qwen3TTSPromptEmbedsBuilder: req_ids: list[str], model_intermediate_buffer: dict[str, dict[str, Any]], device: torch.device, + ) -> None: + """MRv1 adapter for the request-info batch implementation.""" + self.preprocess_infos_batch( + req_infos=[ + info_dict + for req_id in req_ids + if isinstance((info_dict := model_intermediate_buffer.get(req_id)), dict) + ], + device=device, + ) + + @torch.inference_mode() + def preprocess_infos_batch( + self, + *, + req_infos: list[dict[str, Any]], + device: torch.device, ) -> None: """Batch Base voice-clone ref-audio codec extraction for current prefill requests. @@ -775,11 +815,7 @@ class Qwen3TTSPromptEmbedsBuilder: pending_text: list[tuple[dict[str, Any], str]] = [] pending_ref_text: list[tuple[dict[str, Any], str]] = [] groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {} - for req_id in req_ids: - info_dict = model_intermediate_buffer.get(req_id) - if not isinstance(info_dict, dict): - continue - + for info_dict in req_infos: if ( self._needs_initial_prompt_preprocess(info_dict) and first_value(info_dict.get("task_type"), "CustomVoice") == "Base" @@ -837,17 +873,31 @@ class Qwen3TTSPromptEmbedsBuilder: except Exception as exc: logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc) continue - input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None + input_ids = tokenized.get("input_ids") if isinstance(tokenized, Mapping) else None if not isinstance(input_ids, list) or len(input_ids) != len(items): continue + valid_items: list[tuple[dict[str, Any], int]] = [] + flat_ids: list[int] = [] for (info_dict, _), ids in zip(items, input_ids, strict=True): - if isinstance(ids, list) and ids: - info_dict[key] = torch.tensor(ids, dtype=torch.long) + if not ( + isinstance(ids, list) + and ids + and all(isinstance(token_id, (int, np.integer)) for token_id in ids) + ): + continue + valid_items.append((info_dict, len(ids))) + flat_ids.extend(int(token_id) for token_id in ids) + if not valid_items: + continue + use_async_h2d = device.type == "cuda" + packed_ids_cpu = torch.tensor(flat_ids, dtype=torch.long, pin_memory=use_async_h2d) + packed_ids = packed_ids_cpu.to(device=device, non_blocking=use_async_h2d) + offset = 0 + for info_dict, num_ids in valid_items: + info_dict[key] = packed_ids.narrow(0, offset, num_ids) + offset += num_ids groups = {sr: items for sr, items in groups.items() if len(items) >= 2} - if not groups: - return - for sr, items in groups.items(): wavs = [wav for _, wav, _ in items] try: @@ -892,9 +942,7 @@ class Qwen3TTSPromptEmbedsBuilder: codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0) # [1,T,H] codec_embed_sum = torch.cat( [ - self._codec_embed( - torch.tensor([[self._talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long) - ), + self._codec_embed(self._long_tensor([self._talker_config.codec_bos_id], codec_embed_sum.device)), codec_embed_sum, ], dim=1, @@ -990,12 +1038,8 @@ class Qwen3TTSPromptEmbedsBuilder: # tts special token embeds (projected into talker hidden). # ``tts_pad_embed`` is precomputed (request-independent), so we only # need bos/eos here. - tts_tokens = torch.tensor( - [[config.tts_bos_token_id, config.tts_eos_token_id]], - device=input_ids.device, - dtype=input_ids.dtype, - ) - tts_bos_embed, tts_eos_embed = text_projection(text_embedding(tts_tokens)).chunk(2, dim=1) + projected_tts = self._projected_tokens([config.tts_bos_token_id, config.tts_eos_token_id], input_ids.device) + tts_bos_embed, tts_eos_embed = projected_tts.chunk(2, dim=1) tts_pad_embed = self._pad_embed(input_ids.device, tts_bos_embed.dtype) # Codec prefill tags. @@ -1032,9 +1076,9 @@ class Qwen3TTSPromptEmbedsBuilder: ] ] - codec_input_0 = codec_embed(torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long)) + codec_input_0 = codec_embed(self._long_tensor(codec_prefill_list[0], input_ids.device)) codec_input_1 = codec_embed( - torch.tensor([[talker_config.codec_pad_id, talker_config.codec_bos_id]], device=input_ids.device) + self._long_tensor([talker_config.codec_pad_id, talker_config.codec_bos_id], input_ids.device) ) # Speaker embedding/token (task-dependent) @@ -1339,7 +1383,7 @@ class Qwen3TTSPromptEmbedsBuilder: talker_prompt, text_all + codec_embed(pad_ids), tts_pad_embed - + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)), ], dim=1, ) @@ -1392,8 +1436,7 @@ class Qwen3TTSPromptEmbedsBuilder: [ talker_prompt, text_all + codec_embed(pad_ids), - tts_pad_embed - + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)), ], dim=1, ) @@ -1431,8 +1474,7 @@ class Qwen3TTSPromptEmbedsBuilder: [ talker_prompt, text_all + codec_embed(pad_ids), - tts_pad_embed - + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)), + tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)), ], dim=1, ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 8ab5c5e34..ce9f0b04a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -46,9 +46,9 @@ def _codec_ids_from_payload_or_input( if isinstance(codes, dict): audio = codes.get("audio") if isinstance(audio, torch.Tensor) and audio.numel() > 0: - return audio.reshape(-1).to(device=input_ids.device, dtype=torch.long) + return audio.reshape(-1).to(dtype=torch.long) if isinstance(audio, (list, tuple)) and audio: - return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) + return torch.as_tensor(audio, dtype=torch.long).reshape(-1) return input_ids.reshape(-1).to(dtype=torch.long) @@ -63,6 +63,10 @@ class Qwen3TTSCode2Wav(nn.Module): # decoder caches must use the same IDs delivered by on_requests_finished; # payload metadata carries an external ID which may differ. requires_request_ids = True + # A nonempty native codes.audio payload is authoritative; token IDs may + # serve only as per-request control slots for the generation scheduler. + supports_native_payload_input = True + batched_gpu_staging_keys = {("codes", "audio")} def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -427,12 +431,9 @@ class Qwen3TTSCode2Wav(nn.Module): try: _, c = valid_codes_qf[0] logger.info( - "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d", + "Code2Wav codec: frames=%d q=%d batch=%d", c.shape[1], q, - int(torch.unique(c).numel()), - int(c.min().item()), - int(c.max().item()), len(valid_codes_qf), ) except Exception: @@ -454,9 +455,22 @@ class Qwen3TTSCode2Wav(nn.Module): request_lengths = [int(codes_qf.shape[-1]) for _, codes_qf in valid_codes_qf] max_request_length = max(request_lengths) - request_codes = valid_codes_qf[0][1].new_zeros((len(valid_codes_qf), q, max_request_length)) - for row, (_, codes_qf) in enumerate(valid_codes_qf): - request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf) + request_codes_shape = (len(valid_codes_qf), q, max_request_length) + target_device = ids.device + if target_device.type == "cuda" and all(codes_qf.device.type == "cpu" for _, codes_qf in valid_codes_qf): + staged_codes = torch.zeros( + request_codes_shape, + dtype=torch.long, + device="cpu", + pin_memory=True, + ) + for row, (_, codes_qf) in enumerate(valid_codes_qf): + staged_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf) + request_codes = staged_codes.to(device=target_device, non_blocking=True) + else: + request_codes = torch.zeros(request_codes_shape, dtype=torch.long, device=target_device) + for row, (_, codes_qf) in enumerate(valid_codes_qf): + request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf) self._record_decode_batch_stats( group_size=len(valid_codes_qf), 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 dfb24516a..1a1358c17 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 @@ -4,7 +4,8 @@ from __future__ import annotations import copy -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence +from contextlib import contextmanager from typing import Any import numpy as np @@ -25,13 +26,14 @@ from vllm.multimodal.audio import AudioResampler from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload -from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, load_validated_profile_tensors, validate_qwen3_tts_profile, ) +from vllm_omni.worker.sampling_utils import get_tts_local_seed from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only @@ -42,6 +44,80 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2E logger = init_logger(__name__) _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64 +_ASYNC_REF_CODES_PUBLISHED_KEY = "_qwen3_tts_ref_codes_published" + + +def _qwen3_tts_gpu_resident_buffer_keys( + use_v2_model_runner: bool, +) -> set[tuple[str, str]]: + keys = { + ("codes", "audio"), + ("hidden_states", "last"), + ("hidden_states", "trailing_text"), + } + if use_v2_model_runner: + keys.update( + { + ("codes", "ref"), + ("embed", "prefill"), + ("meta", "codec_frame_valid"), + } + ) + return keys + + +def _materialize_span_scalars( + values: Sequence[int], + span_lengths: Sequence[int], + *, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + """Expand a uniform scalar cohort with one GPU allocation.""" + if len(values) != len(span_lengths): + raise ValueError( + f"values and span_lengths must have the same length, got {len(values)} and {len(span_lengths)}" + ) + total = sum(int(length) for length in span_lengths) + if not values: + return torch.empty((0,), dtype=dtype, device=device) + first = int(values[0]) + if not all(int(value) == first for value in values[1:]): + return None + return torch.full((total,), first, dtype=dtype, device=device) + + +def _should_publish_async_ref_codes(info: dict[str, Any]) -> bool: + """Publish ref codes exactly once, including after a KV-resumed prefill.""" + meta = info.get("meta") + return not (isinstance(meta, dict) and bool(meta.get(_ASYNC_REF_CODES_PUBLISHED_KEY))) + + +def _mark_async_ref_codes_published(info: dict[str, Any]) -> None: + meta = info.setdefault("meta", {}) + if not isinstance(meta, dict): + raise TypeError("Qwen3-TTS payload meta must be a dict") + meta[_ASYNC_REF_CODES_PUBLISHED_KEY] = True + + +def _ref_audio_artifact_cache_capacity(vllm_config: VllmConfig) -> int: + model_cfg = getattr(vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + connector_extra = ( + connector_cfg.get("extra", connector_cfg) + if isinstance(connector_cfg, dict) + else getattr(connector_cfg, "extra", None) + ) + if not isinstance(connector_extra, dict): + connector_extra = {} + raw_capacity = connector_extra.get("ref_audio_artifact_cache_max_entries", 1024) + try: + capacity = int(raw_capacity) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid ref_audio_artifact_cache_max_entries={raw_capacity!r}") from exc + if capacity < 0: + raise ValueError(f"ref_audio_artifact_cache_max_entries must be non-negative, got {capacity}") + return capacity def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool: @@ -267,8 +343,7 @@ class Qwen3TTSSpeakerEncoder(torch.nn.Module): class Qwen3TTSTalkerForConditionalGeneration(nn.Module): - """vLLM-AR talker: step-wise layer-0 codec decoding. - Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`.""" + """vLLM-AR talker with step-wise layer-0 codec decoding.""" hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -329,29 +404,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # per-request additional_information. self.talker_mtp_output_key = ("codes", "audio") self.talker_mtp_graph_safe = True - # talker_mtp samples with per-row generators, so explicitly-seeded - # requests stay batched instead of one scalar forward per row (#4883). - # This is only valid while talker_mtp receives the *unpadded* active - # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp - # (padded batches + a single captured RNG stream) whenever full - # cudagraphs are enabled, so disable per-row generators in that case. - # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is - # not reproducible (matching Qwen3-Omni, which has no per-row seeding); - # eager mode keeps the per-row generator fast-path. - # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an - # upstream (runner) concern, and the runner already re-derives the same - # wrap decision in ``_init_talker_mtp``. Once all TTS models declare - # ``talker_mtp_graph_safe`` we should move this gating into the model - # runner (models only expose the static capability flag, and users can - # then configure ``cudagraph_mode`` freely). This also means mtp seeding - # is silently dropped when the decode batch > 1 under full cudagraphs; - # the follow-up should refactor the per-row generator in the runner to - # be cudagraph-safe so seeded batches stay reproducible. - cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None) - talker_mtp_graph_wrapped = ( - self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() - ) - self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped + # 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. + self.talker_mtp_accepts_per_row_generators = True + self.mtp_sample_uniforms = True + self.mtp_sample_steps = max(0, int(self.talker_config.num_code_groups) - 1) + self.mtp_sample_vocab_size = self._codebook_vocab_size self.use_async_omni_output = True self.eager_omni_postprocess_before_async_output = True self.omni_pooler_payload_include_hidden = False @@ -445,14 +504,24 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # the same req_id ordering that indexes sampling_metadata. self._req_x_vector_only: dict[str, bool | None] = {} self._batch_req_ids: list[str] = [] + self._mrv2_silence_ban_mask: torch.Tensor | None = None # Keys that should stay on GPU in model_intermediate_buffer to avoid # CPU-to-GPU round-trips on every decode step. - self.gpu_resident_buffer_keys: set[tuple[str, str]] = { - ("codes", "audio"), - ("hidden_states", "last"), - ("hidden_states", "trailing_text"), - } + self._use_v2_model_runner = bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + self.gpu_resident_buffer_keys = _qwen3_tts_gpu_resident_buffer_keys(self._use_v2_model_runner) + + if self._use_v2_model_runner: + # Align predictor warmup/capture buckets with the outer MRv2 MTP + # graph buckets so the compile cache covers every reachable size. + outer_max_bsz = int(vllm_config.scheduler_config.max_num_seqs) + outer_buckets = { + int(size) + for size in vllm_config.compilation_config.cudagraph_capture_sizes + if 0 < int(size) <= outer_max_bsz + } + outer_buckets.update({1, outer_max_bsz}) # eager fallback rows + self.code_predictor.configure_mtp_execution_buckets(sorted(outer_buckets)) # ``text_proj(text_emb(tts_pad_token_id))`` is request-independent — # it depends only on frozen ``text_embedding`` / ``text_projection`` @@ -492,6 +561,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self._subtalker_sampling_params: dict[str, Any] = ( dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {} ) + self.mtp_sampling_params = self._subtalker_sampling_params self._stacked_codec_embed: torch.Tensor | None = None @@ -514,6 +584,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): tts_pad_embed=self._tts_pad_embed, encode_ref_audio_batch=self._encode_ref_audio_batch, speaker_cache=self._speaker_cache, + ref_audio_artifact_cache_max_entries=_ref_audio_artifact_cache_capacity(vllm_config), ) self._load_custom_voice_profiles() @@ -592,6 +663,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Suppress silence-region tokens for the first N decode frames (#4966). The # per-request decode step is the length of its generated-token history. ban_n = self._silence_ban_frames + if ban_n > 0 and sampling_metadata is None and self._mrv2_silence_ban_mask is not None: + if self._mrv2_silence_ban_mask.shape[0] != logits.shape[0]: + raise ValueError("MRV2 silence-ban batch must match the sampled logits rows") + return logits.masked_fill(self._mrv2_silence_ban_mask & self._silence_mask, float("-inf")) if ban_n > 0 and sampling_metadata is not None: output_token_ids = getattr(sampling_metadata, "output_token_ids", None) if output_token_ids is None or len(output_token_ids) != logits.shape[0]: @@ -645,11 +720,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): info_dicts = kwargs.get("runtime_additional_information") or [] if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) + async_chunk = bool(getattr(model_config, "async_chunk", False)) audio_codes_list: list[torch.Tensor] = [] - ref_code_len_list: list[torch.Tensor] = [] + ref_code_len_segments: list[tuple[int | torch.Tensor, int, torch.device]] = [] ref_code_list: list[torch.Tensor] = [] has_ref_code = False - codec_streaming_list: list[torch.Tensor] = [] + codec_streaming_values: list[int] = [] + codec_streaming_span_lengths: list[int] = [] + codec_streaming_device: torch.device | None = None + codec_frame_valid_device: torch.device | None = None + codec_frame_valid_parts: list[torch.Tensor] = [] for info in info_dicts: if not isinstance(info, dict): ref_code_list.append(torch.empty(0, dtype=torch.long)) @@ -659,17 +740,46 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): ac = codes.get("audio") if isinstance(ac, torch.Tensor): audio_codes_list.append(ac) - cs = meta.get("codec_streaming") - if isinstance(cs, bool): - codec_streaming_list.append( - torch.full((int(ac.shape[0]),), int(cs), dtype=torch.int8, device=ac.device) + frame_valid = meta.get("codec_frame_valid") + if isinstance(frame_valid, bool): + codec_frame_valid_device = ac.device + codec_frame_valid_parts.append( + torch.full( + (int(ac.shape[0]),), + int(frame_valid), + dtype=torch.int8, + device=ac.device, + ) ) + elif isinstance(frame_valid, torch.Tensor): + if frame_valid.numel() != 1: + raise ValueError("codec_frame_valid must be scalar per request") + codec_frame_valid_device = ac.device + codec_frame_valid_parts.append( + frame_valid.to(device=ac.device, dtype=torch.int8, non_blocking=True) + .reshape(1) + .expand(int(ac.shape[0])) + ) + cs = meta.get("codec_streaming") + if not async_chunk and isinstance(cs, bool): + codec_streaming_values.append(int(cs)) + codec_streaming_span_lengths.append(int(ac.shape[0])) + codec_streaming_device = ac.device ref_code = codes.get("ref") - if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: + publish_ref_code = not async_chunk or _should_publish_async_ref_codes(info) + if publish_ref_code and isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: ref_code_list.append(ref_code) has_ref_code = True + if async_chunk: + _mark_async_ref_codes_published(info) else: ref_code_list.append(torch.empty(0, dtype=torch.long)) + if async_chunk: + # The streaming consumer uses codes.ref.shape[0] directly and + # does not read ref_code_len or codec_streaming. Avoid creating + # GPU control tensors that would only be snapshotted and copied + # back to the CPU on every decode step. + continue ref_len = meta.get("ref_code_len") if ref_len is None: continue @@ -680,7 +790,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if ref_len.numel() == 0: raise ValueError("ref_code_len is an empty tensor") ref_len_tail = ref_len.reshape(-1)[-1:].to(dtype=torch.int32, device=ac.device) - ref_code_len_list.append(ref_len_tail.expand(span_len).contiguous()) + ref_code_len_segments.append((ref_len_tail, span_len, ac.device)) continue if isinstance(ref_len, list): if len(ref_len) != 1: @@ -688,7 +798,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): ref_len_val = int(ref_len[0]) else: ref_len_val = int(ref_len) - ref_code_len_list.append(torch.full((span_len,), ref_len_val, dtype=torch.int32, device=ac.device)) + ref_code_len_segments.append((ref_len_val, span_len, ac.device)) if not audio_codes_list: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) @@ -696,15 +806,61 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): audio_codes = torch.cat(audio_codes_list, dim=0) span_len = int(audio_codes.shape[0]) mm: OmniPayload = {"codes": {"audio": audio_codes}} - if ref_code_len_list: - mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len] + if ref_code_len_segments: + if all(isinstance(value, int) for value, _length, _device in ref_code_len_segments): + ref_code_len = _materialize_span_scalars( + [int(value) for value, _length, _device in ref_code_len_segments], + [length for _value, length, _device in ref_code_len_segments], + dtype=torch.int32, + device=ref_code_len_segments[0][2], + ) + if ref_code_len is None: + ref_code_len = torch.cat( + [ + torch.full((length,), int(value), dtype=torch.int32, device=device) + for value, length, device in ref_code_len_segments + ], + dim=0, + ) + else: + ref_code_len_parts = [ + value.expand(length) + if isinstance(value, torch.Tensor) + else torch.full((length,), value, dtype=torch.int32, device=device) + for value, length, device in ref_code_len_segments + ] + ref_code_len = torch.cat(ref_code_len_parts, dim=0) + mm.setdefault("meta", {})["ref_code_len"] = ref_code_len[:span_len] if has_ref_code: # Batch-aligned, one entry per request: ``to_payload_element`` # indexes ``element[idx]`` per request, and a shorter list would # silently broadcast one request's reference codes to the others. mm.setdefault("codes", {})["ref"] = ref_code_list - if codec_streaming_list: - mm.setdefault("meta", {})["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len] + if codec_streaming_values: + assert codec_streaming_device is not None + codec_streaming = _materialize_span_scalars( + codec_streaming_values, + codec_streaming_span_lengths, + dtype=torch.int8, + device=codec_streaming_device, + ) + if codec_streaming is None: + codec_streaming = torch.cat( + [ + torch.full((length,), value, dtype=torch.int8, device=codec_streaming_device) + for value, length in zip( + codec_streaming_values, + codec_streaming_span_lengths, + strict=True, + ) + ], + dim=0, + ) + mm.setdefault("meta", {})["codec_streaming"] = codec_streaming[:span_len] + if codec_frame_valid_parts: + assert codec_frame_valid_device is not None + codec_frame_valid = torch.cat(codec_frame_valid_parts, dim=0) + mm.setdefault("meta", {})["codec_frame_valid"] = codec_frame_valid[:span_len] return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) # -------------------- preprocess / postprocess -------------------- @@ -763,27 +919,32 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if is_prefill: # Prefill (prompt embeddings) - prompt_embeds_cpu = embed.get("prefill") - # First prefill round: prompt_embeds_cpu is not yet populated. - # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round. - is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 + prompt_embeds_buffer = embed.get("prefill") + # First prefill round: the retained prompt buffer is not populated. + # Subsequent chunks slice the buffer stored by the first round. + is_first_prefill = not isinstance(prompt_embeds_buffer, torch.Tensor) or prompt_embeds_buffer.ndim != 2 if is_first_prefill: full_prompt_embeds, tailing_text_hidden, ref_code_len, ref_code = ( self._prompt_builder.build_prompt_embeds(task_type=task_type, info_dict=info_dict) ) - # Store full prompt embeddings on CPU (large, prefill-only). - # tailing_text_hidden stays on GPU (gpu_resident_buffer_keys). - prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous() + prompt_embeds_buffer = full_prompt_embeds.detach() + if not bool(getattr(self, "_use_v2_model_runner", False)): + prompt_embeds_buffer = prompt_embeds_buffer.to("cpu") + prompt_embeds_buffer = prompt_embeds_buffer.contiguous() info_update: OmniPayload = { - "embed": {"prefill": prompt_embeds_cpu}, + "embed": {"prefill": prompt_embeds_buffer}, "hidden_states": {"trailing_text": tailing_text_hidden.detach()}, "meta": { "talker_text_offset": 0, "codec_streaming": codec_streaming, + "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device), }, } if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: - info_update.setdefault("codes", {})["ref"] = ref_code.detach().to("cpu").contiguous() + ref_code_buffer = ref_code.detach() + if not bool(getattr(self, "_use_v2_model_runner", False)): + ref_code_buffer = ref_code_buffer.to("cpu") + info_update.setdefault("codes", {})["ref"] = ref_code_buffer.contiguous() if ref_code_len is not None: info_update["meta"]["ref_code_len"] = int(ref_code_len) # First prefill: source the slice offset from `_omni_num_computed_tokens` @@ -794,16 +955,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): else: # Subsequent prefill chunk: slice from stored embeddings at running offset. offset = max(0, int(meta.get("talker_prefill_offset", 0) or 0)) - info_update = {"meta": {"codec_streaming": codec_streaming}} + info_update = { + "meta": { + "codec_streaming": codec_streaming, + "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device), + } + } # Always return a span_len slice; if the scheduled placeholder is longer than what # the prompt actually fills, pad with tts_pad_embed (preserves placeholder/embedding alignment). - s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) - e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) - take = prompt_embeds_cpu[s:e] + s = max(0, min(offset, int(prompt_embeds_buffer.shape[0]))) + e = max(0, min(offset + span_len, int(prompt_embeds_buffer.shape[0]))) + take = prompt_embeds_buffer[s:e] if int(take.shape[0]) < span_len: pad_n = int(span_len - int(take.shape[0])) - pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1) + pad_rows = tts_pad_embed.to(device=take.device).reshape(1, -1).expand(pad_n, -1) take = torch.cat([take, pad_rows], dim=0) prompt_embeds = take.to(device=input_ids.device, dtype=dtype) info_update["meta"]["talker_prefill_offset"] = int(offset + span_len) @@ -826,7 +992,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): .to(device=input_ids.device, dtype=dtype) .reshape(span_len, -1) ) - return input_ids, inputs_embeds_out, {"meta": {"codec_streaming": codec_streaming}} + return ( + input_ids, + inputs_embeds_out, + { + "meta": { + "codec_streaming": codec_streaming, + "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device), + } + }, + ) # Decode: span_len == 1 # Pop one text-step vector from tailing_text_hidden queue. @@ -884,6 +1059,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): "meta": { "talker_text_offset": int(next_text_offset), "codec_streaming": codec_streaming, + "codec_frame_valid": ( + (input_ids.reshape(-1)[-1] >= 0) & (input_ids.reshape(-1)[-1] < self._codebook_vocab_size) + ).reshape(()), }, } if trailing_text_update is not None: @@ -895,11 +1073,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): *, input_ids: torch.Tensor, req_infos: list[dict[str, Any]], + input_embeds: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[dict[str, Any]]]: """Batch the decode-only preprocess path for Qwen3-TTS. This mirrors the scalar decode branch in ``preprocess()``, but performs - the token embedding lookup once for the whole decode batch. + the token embedding lookup once for the whole decode batch, unless + the runner has already prepared embeddings for these input ids. """ input_ids_flat = input_ids.reshape(-1) if int(input_ids_flat.numel()) != len(req_infos): @@ -907,6 +1087,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): f"preprocess_decode_batch expected {len(req_infos)} input ids, got {int(input_ids_flat.numel())}" ) + codec_frame_valid_batch = (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size) device = input_ids_flat.device dtype = self._embedding_dtype # Request-independent constant (see :meth:`_init_runtime_buffers`) — @@ -979,17 +1160,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): "meta": { "talker_text_offset": int(next_text_offset), "codec_streaming": codec_streaming, + "codec_frame_valid": codec_frame_valid_batch[len(updates)].reshape(()), }, } if trailing_text_update is not None: info_update["hidden_states"] = {"trailing_text": trailing_text_update.detach()} updates.append(info_update) - inputs_embeds_out = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long)).to( - device=device, - dtype=dtype, - ) - inputs_embeds_out = inputs_embeds_out.reshape(len(req_infos), -1) + if input_embeds is None: + input_embeds = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long)) + inputs_embeds_out = input_embeds.to(device=device, dtype=dtype).reshape(len(req_infos), -1) return ( input_ids_flat, inputs_embeds_out, @@ -998,6 +1178,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): updates, ) + def preprocess_decode_batch_mrv2( + self, + *, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[dict[str, Any]]]: + """Expose the existing TTS decode batcher through the MRv2 contract.""" + return self.preprocess_decode_batch(input_ids=input_ids, req_infos=req_infos, input_embeds=input_embeds) + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: # Keep the last token hidden for the next decode step's code predictor. # Stays on GPU - gpu_resident_buffer_keys avoids the CPU round-trip. @@ -1006,6 +1196,23 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): last = hidden_states[-1, :].detach() return {"hidden_states": {"last": last}} + def postprocess_batch_mrv2( + self, + *, + hidden_states: torch.Tensor, + last_token_indices: torch.Tensor, + ) -> tuple[tuple[str, str], OwnedBatchTensor]: + """Gather every request's Talker tail in one GPU operation. + + ``index_select`` freshly allocates the result, so the batch is marked + as owned and the intermediate buffer stores row views without taking + a second snapshot. + """ + return ( + ("hidden_states", "last"), + OwnedBatchTensor(hidden_states.index_select(0, last_token_indices)), + ) + @torch.inference_mode() def preprocess_batch( self, @@ -1026,6 +1233,39 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device=device, ) + @torch.inference_mode() + def preprocess_batch_mrv2( + self, + *, + req_infos: list[dict[str, Any]], + device: torch.device, + ) -> None: + """Batch new MRv2 prefill slots without rebuilding a req-id map.""" + if self._silence_ban_frames > 0: + for info in req_infos: + self._req_x_vector_only[info["req_id"]] = resolve_x_vector_only(info) + self._prompt_builder.preprocess_infos_batch( + req_infos=req_infos, + device=device, + ) + + @contextmanager + def mrv2_sampling_context(self, *, req_ids: list[str], num_output_tokens: torch.Tensor) -> Iterator[None]: + """Scope the silence gate to sampling, never prompt-logprob computation.""" + self.set_batch_req_ids(req_ids) + self._mrv2_silence_ban_mask = None + if self._silence_ban_frames > 0: + in_scope = torch.tensor( + [self._req_x_vector_only.get(req_id) is True for req_id in req_ids], + device=num_output_tokens.device, + dtype=torch.bool, + ) + self._mrv2_silence_ban_mask = (in_scope & (num_output_tokens < self._silence_ban_frames)).unsqueeze(1) + try: + yield + finally: + self._mrv2_silence_ban_mask = None + def set_batch_req_ids(self, req_ids: Sequence[str]) -> None: """Record the current batch's req_ids, in the order that indexes logits. @@ -1051,10 +1291,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): sampling_rate=target_sr, return_tensors="pt", ) - inputs = inputs.to(device).to(torch.bfloat16) - - input_values = inputs["input_values"].squeeze(1) - padding_mask = inputs["padding_mask"].squeeze(1) + input_values = inputs["input_values"].squeeze(1).to(device=device, dtype=torch.bfloat16) + downsample = self._encoder_downsample_rate + valid_code_lengths = [(int(wav.size) + downsample - 1) // downsample for wav in wavs] with torch.inference_mode(): encoded = self.encoder.encode( @@ -1063,13 +1302,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): ) audio_codes = encoded.audio_codes[:, : self._encoder_valid_num_quantizers] - downsample = self._encoder_downsample_rate return [ - code[..., : -(-mask.sum() // downsample)].transpose(0, 1).to(device=device, dtype=torch.long) - for code, mask in zip(audio_codes, padding_mask) + code[:, : int(valid_len)].transpose(0, 1).to(dtype=torch.long).contiguous() + for code, valid_len in zip(audio_codes, valid_code_lengths, strict=True) ] def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + self._prompt_builder._projected_token_cache.clear() # Consume talker weights, and conditionally consume speaker encoder # weights only if they are present in the checkpoint. speaker_weights: list[tuple[str, torch.Tensor]] = [] @@ -1256,6 +1495,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): top_p: float | None = None, generator: torch.Generator | None = None, generators: Sequence[torch.Generator | None] | None = None, + sample_uniforms: torch.Tensor | None = None, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1). @@ -1296,6 +1536,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): top_p=top_p, generator=generator, generators=generators, + sample_uniforms=sample_uniforms, ) # [B, Q] # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes. @@ -1313,3 +1554,20 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): summed = (last_id_hidden.squeeze(1) + gathered.sum(dim=1)).unsqueeze(1) inputs_embeds_out = (summed + text_step).reshape(bsz, -1) return inputs_embeds_out, audio_codes.to(dtype=torch.long) + + # MRV2 capability names read the V1 canonical values so platform patches + # such as the NPU 310P graph-safety override apply to both runners. + mtp = talker_mtp + get_mtp_seed = staticmethod(get_tts_local_seed) + + @property + def mtp_output_key(self) -> tuple[str, str]: + return self.talker_mtp_output_key + + @property + def mtp_graph_safe(self) -> bool: + return self.talker_mtp_graph_safe + + @property + def mtp_accepts_per_row_generators(self) -> bool: + return self.talker_mtp_accepts_per_row_generators diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 9572d9ee2..d01eae654 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +31,7 @@ from transformers import MimiConfig, MimiModel from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.integrations import use_kernel_forward_from_hub -from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask +from transformers.masking_utils import create_sliding_window_causal_mask from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_outputs import BaseModelOutputWithPast @@ -521,15 +524,63 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config) self.gradient_checkpointing = False - self.has_sliding_layers = "sliding_attention" in self.config.layer_types self.window_size = config.sliding_window self.input_proj = nn.Linear(config.latent_dim, config.hidden_size) self.output_proj = nn.Linear(config.hidden_size, config.latent_dim) + # Stateless calls with implicit positions share one broadcastable mask + # per warmup shape. Incremental or explicit-position calls build a live mask. + self._sliding_attention_mask_cache: dict[tuple[int, torch.dtype, torch.device, str], torch.Tensor | None] = {} # Initialize weights and apply final processing self.post_init() + def _apply(self, fn, recurse=True): + # A placement or dtype change invalidates every cached mask. + self._sliding_attention_mask_cache.clear() + return super()._apply(fn, recurse=recurse) + + def _get_sliding_attention_mask( + self, + inputs_embeds: torch.Tensor, + ) -> torch.Tensor | None: + sequence_length = int(inputs_embeds.shape[1]) + max_position_embeddings = int(self.config.max_position_embeddings) + if sequence_length > max_position_embeddings: + raise ValueError( + f"Input length {sequence_length} exceeds max_position_embeddings {max_position_embeddings}" + ) + + cache_key = ( + sequence_length, + inputs_embeds.dtype, + inputs_embeds.device, + str(self.config._attn_implementation), + ) + if cache_key not in self._sliding_attention_mask_cache: + mask_inputs = inputs_embeds[:1] + mask_kwargs = { + "config": self.config, + "attention_mask": None, + "past_key_values": None, + "position_ids": None, + } + sig = inspect.signature(create_sliding_window_causal_mask) + if "input_embeds" in sig.parameters: + mask_kwargs["input_embeds"] = mask_inputs + mask_kwargs["cache_position"] = torch.arange( + sequence_length, + device=inputs_embeds.device, + ) + else: + mask_kwargs["inputs_embeds"] = mask_inputs + sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs) + if sliding_attention_mask is not None: + sliding_attention_mask = sliding_attention_mask.contiguous() + self._sliding_attention_mask_cache[cache_key] = sliding_attention_mask + + return self._sliding_attention_mask_cache[cache_key] + # Note: @check_model_inputs decorator removed for vLLM compatibility # The decorator causes "unexpected keyword argument 'inputs_embeds'" error @auto_docstring @@ -558,6 +609,17 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr inputs_embeds = self.input_proj(inputs_embeds) + sequence_length = inputs_embeds.shape[1] + # Only internally generated positions are known to be contiguous and + # zero-based without reading CUDA tensor values back to the host. + # Explicit positions use the live mask path, including during capture. + uses_prepared_or_incremental_mask = ( + attention_mask is not None + or bool(use_cache) + or past_key_values is not None + or cache_position is not None + or position_ids is not None + ) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) @@ -565,7 +627,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, - past_seen_tokens + inputs_embeds.shape[1], + past_seen_tokens + sequence_length, device=inputs_embeds.device, ) @@ -574,22 +636,24 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr hidden_states = inputs_embeds - if not isinstance(causal_mask_mapping := attention_mask, dict): + if isinstance(attention_mask, dict): + sliding_attention_mask = attention_mask["sliding_attention"] + elif uses_prepared_or_incremental_mask: mask_kwargs = { "config": self.config, "attention_mask": attention_mask, "past_key_values": past_key_values, "position_ids": position_ids, } - sig = inspect.signature(create_causal_mask) + sig = inspect.signature(create_sliding_window_causal_mask) if "input_embeds" in sig.parameters: mask_kwargs["input_embeds"] = inputs_embeds mask_kwargs["cache_position"] = cache_position else: mask_kwargs["inputs_embeds"] = inputs_embeds - causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)} - if self.has_sliding_layers: - causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) + sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs) + else: + sliding_attention_mask = self._get_sliding_attention_mask(inputs_embeds) # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) @@ -597,7 +661,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr for decoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, - attention_mask=causal_mask_mapping[decoder_layer.attention_type], + attention_mask=sliding_attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, @@ -1400,13 +1464,15 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): @staticmethod def _slice_dynamic_cache(cache: DynamicCache, row: int) -> DynamicCache: - request_cache = copy.deepcopy(cache) - for layer in request_cache.layers: - if layer.keys is not None: - layer.keys = layer.keys[row : row + 1].clone() - if layer.values is not None: - layer.values = layer.values[row : row + 1].clone() - return request_cache + # Preserve cache metadata and independent ownership without cloning the + # full batch for every request. Seed deepcopy with only the selected KV + # rows; it will copy the remaining cache structure normally. + memo: dict[int, Any] = {} + for layer in cache.layers: + for tensor in (layer.keys, layer.values): + if tensor is not None and id(tensor) not in memo: + memo[id(tensor)] = tensor[row : row + 1].clone() + return copy.deepcopy(cache, memo) @staticmethod def _cache_tensors_are_batchable(request_caches: list[dict[str, Any]], keys: tuple[str, ...]) -> bool: 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 431569626..79a6e782a 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -70,7 +70,7 @@ def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torc return None if audio_codes.ndim == 2: frame = audio_codes[-1] - if frame.numel() == 0 or not bool(frame.any().item()): + if frame.numel() == 0: return None return frame.to(torch.long).reshape(-1) if audio_codes.ndim == 1: @@ -78,6 +78,31 @@ def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torc raise ValueError(f"Invalid audio_codes shape for Qwen3-TTS async_chunk: {tuple(audio_codes.shape)}") +def _should_append_async_frame( + multimodal_output: OmniPayload | dict[str, Any], + frame: torch.Tensor, +) -> bool: + meta = multimodal_output.get("meta", {}) + frame_valid = meta.get("codec_frame_valid") if isinstance(meta, Mapping) else None + if isinstance(frame_valid, torch.Tensor): + if frame_valid.numel() == 0: + raise ValueError("codec_frame_valid must not be empty") + # A single request can carry a multi-token prefill span. The model + # emits one validity entry per codec row, while _extract_last_frame() + # selects the final row from that span. + return bool(frame_valid.reshape(-1)[-1].item()) + if frame_valid is not None: + return bool(frame_valid) + # Compatibility for payload producers predating the explicit validity + # contract. Real Qwen3-TTS outputs carry codec_frame_valid and avoid this + # value-dependent synchronization. + return bool(frame.any().item()) + + +def _append_tensor_frame(transfer_manager: Any, request_id: str, frame: torch.Tensor) -> None: + transfer_manager.code_prompt_token_ids[request_id].append(frame.detach().to(torch.long).reshape(-1).contiguous()) + + def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any] | None, @@ -93,12 +118,11 @@ def talker2code2wav_async_chunk( if isinstance(multimodal_output, Mapping): frame = _extract_last_frame(multimodal_output) - if frame is not None: - codec_codes = frame.cpu().tolist() - transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + if frame is not None and _should_append_async_frame(multimodal_output, frame): + _append_tensor_frame(transfer_manager, request_id, frame) ref_code = multimodal_output.get("codes", {}).get("ref") if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None: - request_payload[request_id] = ref_code.to(torch.long).cpu().contiguous() + request_payload[request_id] = ref_code.detach().to(torch.long).contiguous() elif not finished: return None @@ -183,6 +207,17 @@ def talker2code2wav_async_chunk( ) initial_chunk_size = chunk_size length = len(transfer_manager.code_prompt_token_ids[request_id]) + emitted_frames = getattr(transfer_manager, "_qwen3_tts_emitted_frames", None) + if emitted_frames is None: + emitted_frames = {} + transfer_manager._qwen3_tts_emitted_frames = emitted_frames + if request_id in emitted_frames and length <= emitted_frames[request_id]: + if finished: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(request_id=request_id, left_context_size=0, finished=torch.tensor(True)), + ) + return None # Full-utterance Code2Wav is gated by ``full_utterance_decode`` (explicit # opt-in via additional_information). Do not reuse ``non_streaming_mode``: @@ -305,6 +340,8 @@ def talker2code2wav_async_chunk( if finished and first_chunk: context_length = length + if request_id in emitted_frames: + context_length = min(context_length, length - emitted_frames[request_id]) if adaptive_enabled: ctrl = transfer_manager._adaptive_states.get(request_id) @@ -343,7 +380,7 @@ def talker2code2wav_async_chunk( if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: ref_context = ref_code if ref_code_context_frames > 0 and int(ref_context.shape[0]) > ref_code_context_frames: - logger.info_once( + logger.debug( "Qwen3-TTS async chunk uses the last %d/%d ref_code frames as bounded Code2Wav context.", ref_code_context_frames, int(ref_context.shape[0]), @@ -375,6 +412,7 @@ def talker2code2wav_async_chunk( meta.ref_context_request_id = ref_context_request_id meta.ref_context_included = ref_context_included + emitted_frames[request_id] = length return OmniPayloadStruct( codes=CodesStruct(audio=code_predictor_codes), meta=meta, @@ -383,6 +421,91 @@ def talker2code2wav_async_chunk( ) +def _copy_async_chunk_tensors_to_cpu( + tensors: list[torch.Tensor], +) -> list[torch.Tensor]: + """Materialize one producer cohort with at most one D2H per device.""" + if not tensors: + return [] + + outputs: list[torch.Tensor | None] = [None] * len(tensors) + by_device: dict[torch.device, list[tuple[int, torch.Tensor]]] = {} + for index, tensor in enumerate(tensors): + value = tensor.detach().to(dtype=torch.long).reshape(-1).contiguous() + if value.device.type == "cpu": + outputs[index] = value + else: + by_device.setdefault(value.device, []).append((index, value)) + + for entries in by_device.values(): + lengths = [int(value.numel()) for _, value in entries] + combined = torch.cat([value for _, value in entries], dim=0).to(device="cpu") + offset = 0 + for (index, source), length in zip(entries, lengths): + outputs[index] = combined[offset : offset + length].reshape(source.shape).contiguous() + offset += length + + return [value for value in outputs if value is not None] + + +def talker2code2wav_async_chunk_batch( + transfer_manager: Any, + pooling_outputs: list[OmniPayload | dict[str, Any] | None], + requests: list[Any], + is_finished: list[bool], +) -> list[OmniPayloadStruct | None]: + """Build one Talker model-step worth of Code2Wav payloads. + + MRv2 invokes this hook from its ordered native output worker. Codec frames + and first-chunk reference codes from the same model step are copied as a + cohort, while chunk boundaries and request metadata remain delegated to the + scalar builder below. + """ + if not (len(pooling_outputs) == len(requests) == len(is_finished)): + raise ValueError("batch codec inputs must have identical lengths") + + request_payload = getattr(transfer_manager, "request_payload", None) + if request_payload is None: + request_payload = {} + transfer_manager.request_payload = request_payload + + selected: list[torch.Tensor] = [] + destinations: list[tuple[str, str, tuple[int, ...]]] = [] + for pooling_output, request in zip(pooling_outputs, requests): + if not isinstance(pooling_output, Mapping): + continue + request_id = request.external_req_id + frame = _extract_last_frame(pooling_output) + if frame is not None and _should_append_async_frame(pooling_output, frame): + selected.append(frame) + destinations.append(("frame", request_id, tuple(frame.shape))) + + ref_code = pooling_output.get("codes", {}).get("ref") + if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None: + selected.append(ref_code) + destinations.append(("ref", request_id, tuple(ref_code.shape))) + + copied = _copy_async_chunk_tensors_to_cpu(selected) + for value, (kind, request_id, shape) in zip(copied, destinations): + value = value.reshape(shape) + if kind == "frame": + _append_tensor_frame(transfer_manager, request_id, value) + else: + request_payload[request_id] = value.contiguous() + + payloads: list[OmniPayloadStruct | None] = [] + for request, finished in zip(requests, is_finished): + payloads.append( + talker2code2wav_async_chunk( + transfer_manager=transfer_manager, + multimodal_output={}, + request=request, + is_finished=finished, + ) + ) + return payloads + + # ============================================================================ # Worker-connector data plane (non-async-chunk path). # AR runner's `flatten_payload` converts the model emit diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index e74931a20..ff0a21d45 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -60,6 +60,10 @@ class OmniModelRunnerOutput(ModelRunnerOutput): # The Scheduler can safely free the block tables for these requests. kv_extracted_req_ids: list[str] | None = None omni_connector_output: OmniConnectorOutput | None = None + # True when sampled_token_ids has already been materialized on the host. + # MRv2 consumers must treat the value as immutable and must not perform + # another device-to-host conversion or rebuild it from sampler tensors. + sampled_token_ids_materialized: bool = False @classmethod def with_kv_conn_output_only(cls, kv_connector_output: Any) -> "OmniModelRunnerOutput": diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 5d2ede43d..cac8d9990 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -684,11 +684,19 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): # Accumulate multimodal tensors regardless of path. if isinstance(req_state, OmniRequestState): + mm_type = getattr(eco, "output_type", None) or default_mm_type + mm_output = getattr(eco, "multimodal_output", None) if mm_output is not None: - mm_type = getattr(eco, "output_type", None) or default_mm_type req_state.add_multimodal_tensor(mm_output, mm_type) + # Omni AR stages carry model payloads alongside text tokens in + # pooling_output. Preserve the payload, then clear the field so + # upstream still runs the normal detokenization path. + if eco.pooling_output is not None and req_state.detokenizer is not None: + req_state.add_multimodal_tensor(eco.pooling_output, mm_type) + eco.pooling_output = None + # Route: if no detokenizer and no pooling output, handle locally # to avoid upstream's assert on detokenizer. if req_state.detokenizer is None and eco.pooling_output is None: @@ -697,20 +705,22 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): upstream_outputs.append(eco) # Handle multimodal-only outputs (generation stages) locally. - mm_request_outputs = self._process_mm_only_outputs( + mm_request_outputs, mm_reqs_to_abort = self._process_mm_only_outputs( mm_only_outputs, engine_core_timestamp=engine_core_timestamp, iteration_stats=iteration_stats, ) # Delegate text/pooling outputs to upstream. - processed = super().process_outputs( + upstream_processed = super().process_outputs( upstream_outputs, engine_core_timestamp=engine_core_timestamp, iteration_stats=iteration_stats, ) - processed.request_outputs.extend(mm_request_outputs) - return processed + return OutputProcessorOutput( + request_outputs=mm_request_outputs + upstream_processed.request_outputs, + reqs_to_abort=mm_reqs_to_abort + upstream_processed.reqs_to_abort, + ) def _process_mm_only_outputs( self, @@ -718,30 +728,32 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): *, engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, - ) -> list[OmniRequestOutput | PoolingRequestOutput]: + ) -> tuple[list[OmniRequestOutput | PoolingRequestOutput], list[str]]: """Handle outputs from generation stages that have no detokenizer. These cannot go through upstream process_outputs because it asserts detokenizer is not None when pooling_output is None. """ request_outputs: list[OmniRequestOutput | PoolingRequestOutput] = [] + reqs_to_abort: list[str] = [] for eco in engine_core_outputs: req_state = self.request_states.get(eco.request_id) if req_state is None or not isinstance(req_state, OmniRequestState): continue - new_token_ids = eco.new_token_ids - finish_reason = eco.finish_reason - stop_reason = eco.stop_reason - kv_transfer_params = eco.kv_transfer_params - ec_transfer_params = getattr(eco, "ec_transfer_params", None) - routed_experts = eco.routed_experts self._update_stats_from_output( req_state, eco, engine_core_timestamp, iteration_stats, ) + + new_token_ids = eco.new_token_ids + finish_reason = eco.finish_reason + stop_reason = eco.stop_reason + kv_transfer_params = eco.kv_transfer_params + ec_transfer_params = eco.ec_transfer_params + routed_experts = eco.routed_experts prefill_stats = getattr(eco, "prefill_stats", None) if prefill_stats is not None: req_state.num_cached_tokens = prefill_stats.num_cached_tokens @@ -770,12 +782,17 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): is_segment_finished = bool(getattr(eco, "is_segment_finished", False)) if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk: self._finish_request(req_state) + if not getattr(eco, "finished", True): + reqs_to_abort.append(eco.request_id) self._update_stats_from_finished( req_state, finish_reason, iteration_stats, ) - return request_outputs + if self.tracing_enabled: + self.do_tracing(eco, req_state, iteration_stats) + + return request_outputs, reqs_to_abort def _update_stats_from_output( self, @@ -804,6 +821,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): None, ) + reported_count = getattr(engine_core_output, "num_generation_tokens", None) + if isinstance(reported_count, int) and not isinstance(reported_count, bool) and reported_count >= 0: + self._native_text_metric_record(req_state.external_req_id)["num_generation_tokens"] = reported_count + if iteration_stats is None or engine_core_timestamp is None or native_stats is None: return @@ -816,7 +837,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): req_state.lora_name, ) record = self._native_text_metric_record(req_state.external_req_id) - record["num_generation_tokens"] = int(native_stats.num_generation_tokens) + record["num_generation_tokens"] = ( + reported_count + if isinstance(reported_count, int) and not isinstance(reported_count, bool) and reported_count >= 0 + else int(native_stats.num_generation_tokens) + ) if was_prefilling: record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) return diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 44ec3dad5..1d0fcd5dc 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Base worker class for vLLM-Omni with device-level GPU memory profiling.""" from __future__ import annotations diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 654cd4dc5..7ee43a05b 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import gc import os @@ -75,7 +78,6 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): f"DP adjusted local rank {self.local_rank} is out of " f"bounds for {torch.accelerator.device_count()} devices." ) - visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank) self.device = current_omni_platform.get_torch_device(visible_device_index) torch.accelerator.set_device_index(self.device) @@ -113,13 +115,24 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) - if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False + model_stage = getattr(self.model_config, "model_stage", None) + stage_id = getattr(self.model_config, "stage_id", None) - # Construct the model runner - self.model_runner = self.model_runner_cls(self.vllm_config, self.device) + self.use_v2_model_runner = bool(getattr(self.model_config, "use_v2_model_runner", False)) + if self.use_v2_model_runner: + from vllm_omni.worker_v2.omni_ar_model_runner import ( + OmniARModelRunner, + ) + + logger.info( + "Using MR v2 OmniARModelRunner for omni AR stage (stage_id=%s model_stage=%s).", + stage_id, + model_stage, + ) + self.model_runner = OmniARModelRunner(self.vllm_config, self.device) + else: + self.use_v2_model_runner = False + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index d6a8d6c8e..0e939957e 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import gc import os @@ -8,7 +11,10 @@ from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import init_worker_distributed_environment +from vllm.v1.worker.gpu_worker import ( + CompilationTimes, + init_worker_distributed_environment, +) from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.platforms import current_omni_platform @@ -96,13 +102,28 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) + self.use_v2_model_runner = bool(getattr(self.model_config, "use_v2_model_runner", False)) if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False + from vllm_omni.worker_v2.omni_generation_model_runner import ( + OmniGenerationModelRunner, + ) - self.model_runner = self.model_runner_cls(self.vllm_config, self.device) + logger.info("Using MR v2 OmniGenerationModelRunner for generation stage.") + self.model_runner = OmniGenerationModelRunner(self.vllm_config, self.device) + else: + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + + @instrument(span_name="Compile/warmup") + def compile_or_warm_up_model(self) -> CompilationTimes: + """Generation stages have no KV cache or sampler — skip warmup_kernels.""" + if not self.use_v2_model_runner: + return super().compile_or_warm_up_model() + import time + + start = time.perf_counter() + self.model_runner.profile_run() + return CompilationTimes(language_model=time.perf_counter() - start, encoder=0.0) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 48e9764f9..1664df24b 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -2042,10 +2042,20 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): if req_state is not None: existing = self.model_intermediate_buffer.setdefault(req_id, {}) existing.setdefault(out_key[0], {})[out_key[1]] = row + validity_key = getattr(self.model, "talker_mtp_validity_key", None) + if validity_key is not None: + existing.setdefault(validity_key[0], {})[validity_key[1]] = torch.ones( + (), dtype=torch.bool, device=code_predictor_codes.device + ) req_state.additional_information_cpu = existing else: for idx, req_id in enumerate(decode_req_ids): update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} + validity_key = getattr(self.model, "talker_mtp_validity_key", None) + if validity_key is not None: + update_dict.setdefault(validity_key[0], {})[validity_key[1]] = torch.ones( + (), dtype=torch.bool, device=code_predictor_codes.device + ) self._update_intermediate_buffer(req_id, update_dict) def _model_forward( diff --git a/vllm_omni/worker/payload_span.py b/vllm_omni/worker/payload_span.py new file mode 100644 index 000000000..c16126468 --- /dev/null +++ b/vllm_omni/worker/payload_span.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Helpers for explicit thinker decode span metadata.""" + +from collections.abc import Mapping +from typing import Any + +import torch + +TensorSpan = tuple[torch.Tensor, int, int] + + +def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None: + tensor = payload.get(tensor_key) + start = payload.get(start_key) + end = payload.get(end_key) + if not isinstance(tensor, torch.Tensor): + return None + if not isinstance(start, int) or not isinstance(end, int): + return None + if start < 0 or end < start or (end - start) != int(tensor.shape[0]): + return None + return tensor, start, end + + +def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None: + if existing_span is None or incoming_span is None: + return None + + existing_tensor, existing_start, existing_end = existing_span + incoming_tensor, incoming_start, incoming_end = incoming_span + if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype: + incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype) + if incoming_start == existing_end: + if existing_start == existing_end: + return incoming_tensor, incoming_start, incoming_end + return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end + if incoming_start < existing_end: + overlap = existing_end - incoming_start + if overlap >= int(incoming_tensor.shape[0]): + return existing_tensor, existing_start, existing_end + trimmed_tensor = incoming_tensor[overlap:] + return ( + torch.cat([existing_tensor, trimmed_tensor], dim=0), + existing_start, + existing_end + int(trimmed_tensor.shape[0]), + ) + return None diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index 8bcea3e6d..00e19024d 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/vllm_omni/worker/sampling_utils.py @@ -6,18 +6,72 @@ from typing import Any import torch from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor logger = init_logger(__name__) __all__ = [ - "build_model_sampler_extra_args", - "call_model_sampler", + "apply_fixed_seed_to_sampling_params", + "get_tts_local_seed", "clamp_prompt_ids_to_penalty_padding", "sanitize_min_tokens_stop_ids", + "sanitize_sampling_params_min_tokens_stop_ids", + "build_model_sampler_extra_args", + "call_model_sampler", ] +def get_tts_local_seed(sampling_params: Any) -> int | None: + """Return the explicit per-request Talker MTP seed, if configured.""" + extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None + seed = extra_args.get("tts_local_seed") if isinstance(extra_args, dict) else None + return int(seed) if seed is not None else None + + +def apply_fixed_seed_to_sampling_params( + sampling_params: Any, + seed: int, + *, + seed_talker_mtp: bool, +) -> None: + """Apply one request seed to a stage sampler and, for Talkers, its MTP RNG.""" + seed = int(seed) + sampling_params.seed = seed + if not seed_talker_mtp: + return + extra_args = dict(getattr(sampling_params, "extra_args", None) or {}) + extra_args["tts_local_seed"] = seed + sampling_params.extra_args = extra_args + + +def sanitize_sampling_params_min_tokens_stop_ids( + sampling_params: SamplingParams, + logits_vocab: int, +) -> None: + """Remove unreachable ids before MRv2 builds persistent min-token state. + + MRv2 copies ``SamplingParams.all_stop_token_ids`` into GPU-resident + ``LogitBiasState`` when a request is registered. Narrow codec heads cannot + safely carry the text tokenizer EOS in that state. The engine-facing EOS + and explicit ``stop_token_ids`` fields remain unchanged; only the set used + by the sampler's min-token mask is normalized. + """ + if sampling_params.min_tokens <= 0: + return + stop_token_ids = sampling_params.all_stop_token_ids + unreachable = {token_id for token_id in stop_token_ids if token_id < 0 or token_id >= logits_vocab} + if not unreachable: + return + stop_token_ids.difference_update(unreachable) + logger.warning_once( + "min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); " + "the model head cannot emit them.", + str(sorted(unreachable)), + logits_vocab, + ) + + def build_model_sampler_extra_args(input_batch: Any, requests: Any) -> list[dict | None]: """Return per-request ``SamplingParams.extra_args`` in batch-row order.""" request_states = requests or {} diff --git a/vllm_omni/worker_v2/__init__.py b/vllm_omni/worker_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker_v2/delivery.py b/vllm_omni/worker_v2/delivery.py new file mode 100644 index 000000000..e84ed28c0 --- /dev/null +++ b/vllm_omni/worker_v2/delivery.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import threading +import time +from enum import Enum, auto + + +class DeliveryState(Enum): + QUEUED = auto() + IN_FLIGHT = auto() + DELIVERED = auto() + FAILED = auto() + CANCELLED = auto() + + +class DeliveryTimeoutError(TimeoutError): + pass + + +class DeliveryCancelledError(RuntimeError): + pass + + +_TERMINAL_STATES = { + DeliveryState.DELIVERED, + DeliveryState.FAILED, + DeliveryState.CANCELLED, +} + + +class DeliveryTicket: + """Exactly-once completion for one MRv2 connector delivery.""" + + def __init__( + self, + manager: OmniDeliveryManager, + *, + request_id: str, + put_key: str, + deadline: float, + ) -> None: + self._manager = manager + self.request_id = request_id + self.put_key = put_key + self.deadline = deadline + self._state = DeliveryState.QUEUED + self._error: BaseException | None = None + self._event = threading.Event() + + @property + def state(self) -> DeliveryState: + with self._manager._lock: + return self._state + + def mark_in_flight(self) -> bool: + return self._manager._mark_in_flight(self) + + def set_delivered(self) -> bool: + return self._manager._set_delivered(self) + + def set_failed(self, error: BaseException) -> bool: + return self._manager._set_failed(self, error) + + # Compatibility with the connector mixin's completion interface. + def set_result(self) -> bool: + return self.set_delivered() + + def set_error(self, error: BaseException) -> bool: + return self.set_failed(error) + + def wait(self) -> None: + while True: + with self._manager._lock: + state = self._state + error = self._error + remaining = self.deadline - time.monotonic() + if state is DeliveryState.DELIVERED: + return + if state in (DeliveryState.FAILED, DeliveryState.CANCELLED): + assert error is not None + raise error + if remaining <= 0: + self._manager._timeout(self) + continue + self._event.wait(timeout=remaining) + + +class OmniDeliveryManager: + """MRv2-only connector delivery state machine. + + A timeout or permanent delivery failure quarantines the whole manager. + This is a stage-fatal condition: no new payload may be accepted after the + transport can no longer prove ordered, exactly-once completion. + """ + + def __init__( + self, + *, + delivery_timeout_s: float, + shutdown_timeout_s: float = 5.0, + ) -> None: + if delivery_timeout_s <= 0: + raise ValueError("delivery_timeout_s must be positive") + if shutdown_timeout_s <= 0: + raise ValueError("shutdown_timeout_s must be positive") + self.delivery_timeout_s = float(delivery_timeout_s) + self.shutdown_timeout_s = float(shutdown_timeout_s) + self._lock = threading.Lock() + self._tickets: set[DeliveryTicket] = set() + self._quarantine_error: BaseException | None = None + self._closed = False + + @property + def is_quarantined(self) -> bool: + with self._lock: + return self._quarantine_error is not None + + def create_ticket(self, *, request_id: str, put_key: str) -> DeliveryTicket: + with self._lock: + if self._closed: + raise RuntimeError("MRv2 delivery manager is closed") + if self._quarantine_error is not None: + raise RuntimeError("MRv2 delivery manager is quarantined") from self._quarantine_error + ticket = DeliveryTicket( + self, + request_id=request_id, + put_key=put_key, + deadline=time.monotonic() + self.delivery_timeout_s, + ) + self._tickets.add(ticket) + return ticket + + def quarantine(self, error: BaseException) -> None: + with self._lock: + self._quarantine_locked(error, primary_ticket=None) + + def shutdown(self, cause: BaseException | None = None) -> None: + with self._lock: + if self._closed: + return + self._closed = True + if self._quarantine_error is None: + self._quarantine_error = cause or RuntimeError("MRv2 stage shutdown") + for ticket in tuple(self._tickets): + if ticket._state in _TERMINAL_STATES: + continue + message = f"connector delivery cancelled for request {ticket.request_id!r}, key {ticket.put_key!r}" + if cause is not None: + message += f": {cause}" + self._finish_locked( + ticket, + DeliveryState.CANCELLED, + DeliveryCancelledError(message), + ) + + def _mark_in_flight(self, ticket: DeliveryTicket) -> bool: + with self._lock: + if ticket._state is not DeliveryState.QUEUED: + return False + ticket._state = DeliveryState.IN_FLIGHT + return True + + def _set_delivered(self, ticket: DeliveryTicket) -> bool: + with self._lock: + if ticket._state in _TERMINAL_STATES: + return False + return self._finish_locked(ticket, DeliveryState.DELIVERED, None) + + def _set_failed(self, ticket: DeliveryTicket, error: BaseException) -> bool: + with self._lock: + if ticket._state in _TERMINAL_STATES: + return False + self._quarantine_locked(error, primary_ticket=ticket) + return True + + def _timeout(self, ticket: DeliveryTicket) -> bool: + error = DeliveryTimeoutError( + "connector delivery timed out after " + f"{self.delivery_timeout_s:.3f}s for request " + f"{ticket.request_id!r}, key {ticket.put_key!r}" + ) + with self._lock: + if ticket._state in _TERMINAL_STATES: + return False + self._quarantine_locked(error, primary_ticket=ticket) + return True + + def _quarantine_locked( + self, + error: BaseException, + *, + primary_ticket: DeliveryTicket | None, + ) -> None: + if self._quarantine_error is None: + self._quarantine_error = error + stage_error = self._quarantine_error + for ticket in tuple(self._tickets): + if ticket._state in _TERMINAL_STATES: + continue + ticket_error = error if ticket is primary_ticket else stage_error + self._finish_locked(ticket, DeliveryState.FAILED, ticket_error) + + def _finish_locked( + self, + ticket: DeliveryTicket, + state: DeliveryState, + error: BaseException | None, + ) -> bool: + if ticket._state in _TERMINAL_STATES: + return False + ticket._state = state + ticket._error = error + self._tickets.discard(ticket) + ticket._event.set() + return True diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py new file mode 100644 index 000000000..df327a129 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Omni-aware ``init_model_state`` factory. + +Extends the upstream v2 factory with Omni architecture dispatch. +Non-Omni architectures fall through to the upstream ``init_model_state``. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states import ( + init_model_state as _upstream_init_model_state, +) +from vllm.v1.worker.gpu.model_states.interface import ModelState + + +def init_omni_model_state( + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, +) -> ModelState: + """Create the appropriate ``ModelState`` for *model*. + + Returns an ``OmniModelState`` when the model declares Omni lifecycle + capabilities; otherwise + delegates to the upstream v2 factory. + """ + uses_omni_lifecycle = any( + getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs") + ) + if uses_omni_lifecycle: + from vllm_omni.worker_v2.model_states.omni_model_state import ( + OmniModelState, + ) + + return OmniModelState(vllm_config, model, encoder_cache, device) + return _upstream_init_model_state(vllm_config, model, encoder_cache, device) diff --git a/vllm_omni/worker_v2/model_states/intermediate_buffer.py b/vllm_omni/worker_v2/model_states/intermediate_buffer.py new file mode 100644 index 000000000..b9614a729 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/intermediate_buffer.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniIntermediateBuffer — per-request cross-stage state for Omni pipelines. + +Uses ``req_index`` (not ``req_id``) for O(1) access, aligned with v2's +``RequestState`` slot management. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.input_batch import InputBatch + +from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans + + +def _resolve_prompt_embeds(pe: Any) -> torch.Tensor | None: + """Convert a prompt_embeds payload to a contiguous CPU tensor.""" + if pe is None: + return None + if isinstance(pe, torch.Tensor): + return pe.detach().cpu().contiguous() + try: + data = getattr(pe, "data", None) + shape = getattr(pe, "shape", None) + if data is not None and shape is not None: + dt = np.dtype(getattr(pe, "dtype", "float32")) + arr = np.frombuffer(data, dtype=dt).reshape(shape) + return torch.from_numpy(arr.copy()) + except Exception as exc: + raise ValueError("Failed to decode prompt_embeds payload") from exc + raise TypeError(f"Unsupported prompt_embeds payload type: {type(pe).__name__}") + + +def _resolve_additional_information(payload: Any) -> dict[str, Any]: + """Convert an additional_information payload to a plain dict.""" + if payload is None: + return {} + try: + from vllm_omni.data_entry_keys import unflatten_payload + from vllm_omni.engine.serialization import deserialize_additional_information + + info = deserialize_additional_information(payload) + if isinstance(info, dict) and any(isinstance(key, str) and "." in key for key in info): + info = unflatten_payload(info) + except Exception as exc: + raise ValueError("Failed to decode additional_information payload") from exc + if not isinstance(info, dict): + raise TypeError(f"Decoded additional_information must be a dict, got {type(info).__name__}") + return info + + +class OmniIntermediateBuffer: + """Per-request intermediate state for multi-stage Omni pipelines. + + Stores ``prompt_embeds``, ``additional_information``, ``mm_features``, + ``req_id`` and any runtime updates written back by model postprocess. + """ + + def __init__(self, max_num_reqs: int): + self.buffers: list[dict[str, Any]] = [{} for _ in range(max_num_reqs)] + self.req_id_to_index: dict[str, int] = {} + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + info: dict[str, Any] = {} + + pe = getattr(new_req_data, "prompt_embeds", None) + if pe is not None: + pe_cpu = _resolve_prompt_embeds(pe) + if pe_cpu is not None: + info["prompt_embeds_cpu"] = pe_cpu + + ai = getattr(new_req_data, "additional_information", None) + if ai is not None: + info.update(_resolve_additional_information(ai)) + + if new_req_data.mm_features: + info["mm_features"] = new_req_data.mm_features + + sampling_params = getattr(new_req_data, "sampling_params", None) + if sampling_params is not None: + info["sampling_params"] = sampling_params + + info["req_id"] = new_req_data.req_id + previous_id = self.buffers[req_index].get("req_id") + if previous_id is not None and self.req_id_to_index.get(previous_id) == req_index: + self.req_id_to_index.pop(previous_id, None) + self.buffers[req_index] = info + self.req_id_to_index[new_req_data.req_id] = req_index + + def remove_request(self, req_index: int) -> None: + req_id = self.buffers[req_index].get("req_id") + if req_id is not None and self.req_id_to_index.get(req_id) == req_index: + self.req_id_to_index.pop(req_id) + self.buffers[req_index] = {} + + def gather(self, input_batch: InputBatch) -> list[dict[str, Any]]: + """Return buffer dicts in current batch order (via ``idx_mapping_np``).""" + return [self.buffers[idx] for idx in input_batch.idx_mapping_np] + + @staticmethod + def _split_gpu_keys(gpu_resident_keys: set[Any] | None) -> tuple[set[Any], dict[Any, set[Any]]]: + top_level: set[Any] = set() + nested: dict[Any, set[Any]] = {} + for key in gpu_resident_keys or set(): + if isinstance(key, tuple) and len(key) == 2: + nested.setdefault(key[0], set()).add(key[1]) + else: + top_level.add(key) + return top_level, nested + + @staticmethod + def _store_value(dest: dict[Any, Any], key: Any, value: Any, gpu_keys: set[Any]) -> None: + if isinstance(value, torch.Tensor): + if key in gpu_keys: + dest[key] = value.detach().clone() + else: + dest[key] = value.detach().cpu().contiguous() + elif isinstance(value, list): + dest[key] = [ + ( + item.detach().clone() + if key in gpu_keys and isinstance(item, torch.Tensor) + else item.detach().cpu().contiguous() + if isinstance(item, torch.Tensor) + else item + ) + for item in value + ] + else: + dest[key] = value + + def update( + self, + req_index: int, + updates: dict[Any, Any], + gpu_resident_keys: set[Any] | None = None, + ) -> None: + """Merge *updates* into the buffer at *req_index*. + + Tensors are detached; those whose key is **not** in + *gpu_resident_keys* are moved to CPU. + """ + if not updates: + return + top_gpu_keys, nested_gpu_keys = self._split_gpu_keys(gpu_resident_keys) + existing = self.buffers[req_index] + for k, v in updates.items(): + if isinstance(k, tuple) and len(k) == 2: + type_key, qualifier = k + existing_sub = existing.setdefault(type_key, {}) + if not isinstance(existing_sub, dict): + existing_sub = {} + existing[type_key] = existing_sub + self._store_value(existing_sub, qualifier, v, nested_gpu_keys.get(type_key, set())) + elif isinstance(v, dict): + existing_sub = existing.setdefault(k, {}) + if not isinstance(existing_sub, dict): + existing_sub = {} + existing[k] = existing_sub + merged_qualifiers: set[str] = set() + if k == "embed" and isinstance(v.get("decode"), torch.Tensor): + span_keys = {"decode_token_start", "decode_token_end"} + if span_keys.intersection(v): + incoming_span = get_tensor_span( + v, + tensor_key="decode", + start_key="decode_token_start", + end_key="decode_token_end", + ) + if incoming_span is None: + raise RuntimeError("Incoming Thinker decode delta has an invalid absolute span") + existing_span = get_tensor_span( + existing_sub, + tensor_key="decode", + start_key="decode_token_start", + end_key="decode_token_end", + ) + merged_span: tuple[Any, int, int] | None = incoming_span + if existing_span is not None: + merged_span = merge_tensor_spans(existing_span, incoming_span) + if merged_span is None: + raise RuntimeError( + "Non-contiguous pending Thinker decode spans in model buffer: " + f"existing=[{existing_span[1]}, {existing_span[2]}) " + f"incoming=[{incoming_span[1]}, {incoming_span[2]})" + ) + assert merged_span is not None + merged_tensor, merged_start, merged_end = merged_span + self._store_value( + existing_sub, + "decode", + merged_tensor, + nested_gpu_keys.get(k, set()), + ) + existing_sub["decode_token_start"] = merged_start + existing_sub["decode_token_end"] = merged_end + merged_qualifiers.update({"decode", "decode_token_start", "decode_token_end"}) + for qualifier, value in v.items(): + if qualifier in merged_qualifiers: + continue + self._store_value(existing_sub, qualifier, value, nested_gpu_keys.get(k, set())) + else: + self._store_value(existing, k, v, top_gpu_keys) + + def update_gpu_tensor_rows( + self, + req_indices: list[int], + key: Any, + values: torch.Tensor, + *, + keepdim: bool = True, + ) -> None: + """Snapshot a batch-first tensor once and retain owned row views.""" + self._store_gpu_tensor_rows(req_indices, key, values.detach().clone(), keepdim=keepdim) + + def update_owned_gpu_tensor_rows( + self, + req_indices: list[int], + key: Any, + owned_values: torch.Tensor, + *, + keepdim: bool = True, + ) -> None: + """Store row views of an already-owned batch tensor without cloning. + + The caller must guarantee ``owned_values`` is not written after this + call (the model output ownership contract). Row views keep the owned storage + alive until each request row is replaced or the slot is freed. + """ + self._store_gpu_tensor_rows(req_indices, key, owned_values, keepdim=keepdim) + + def _store_gpu_tensor_rows( + self, + req_indices: list[int], + key: Any, + owned: torch.Tensor, + *, + keepdim: bool, + ) -> None: + num_rows = len(req_indices) + if owned.ndim == 0 or owned.shape[0] != num_rows: + raise ValueError( + "Batched GPU state update changed the request axis: " + f"key={key!r} expected={num_rows} " + f"actual={owned.shape[0] if owned.ndim else 0}" + ) + + for row, req_index in enumerate(req_indices): + existing = self.buffers[req_index] + row_value = owned[row : row + 1] if keepdim else owned[row] + if isinstance(key, tuple) and len(key) == 2: + type_key, qualifier = key + existing_sub = existing.setdefault(type_key, {}) + if not isinstance(existing_sub, dict): + existing_sub = {} + existing[type_key] = existing_sub + existing_sub[qualifier] = row_value + else: + existing[key] = row_value diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py new file mode 100644 index 000000000..48515b4f5 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -0,0 +1,1223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniModelState — generic ModelState base for all Omni model stages. + +Extends ``DefaultModelState`` with: + +* Cross-stage intermediate buffer (``OmniIntermediateBuffer``) +* ``model_intermediate_buffer`` / ``runtime_additional_information`` injection + into ``model_inputs`` via ``prepare_inputs()`` +* ``OmniOutput`` → ``(text_hidden, multimodal_outputs)`` post-processing +""" + +from __future__ import annotations + +import inspect +import threading +import types +from collections.abc import Callable +from typing import Any, cast + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.states import RequestState + +from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor +from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + OmniIntermediateBuffer, +) + +logger = init_logger(__name__) + + +_rope_patch_lock = threading.Lock() + + +def _default_mrope_positions( + self_model: Any, + input_tokens: list[int], + mm_features: list, +) -> tuple[torch.Tensor, int]: + """Return 3D sequential positions with zero delta. + + For non-vision Omni models (e.g. TTS Talker), all 3 M-RoPE + dimensions use the same sequential positions. Delta=0 keeps decode + positions sequential, identical to the 1D case but broadcast to 3 dims. + """ + n = len(input_tokens) + pos = torch.arange(n, dtype=torch.long) + return pos.unsqueeze(0).expand(3, -1), 0 + + +def _make_safe_get_rope(orig_get_rope): + from vllm.v1.worker.gpu.mm.rope import RopeState + + def _safe_get_rope(model_config: Any, mdl: Any, **kwargs: Any) -> Any: + try: + result = orig_get_rope(model_config, mdl, **kwargs) + except AssertionError: + result = None + + needs_mrope = bool(getattr(model_config, "uses_mrope", False)) + if result is not None and (not needs_mrope or getattr(result, "num_dims", 0) >= 3): + return result + if not needs_mrope: + return None + if not hasattr(mdl, "get_mrope_input_positions"): + mdl.get_mrope_input_positions = types.MethodType(_default_mrope_positions, mdl) + return RopeState(num_dims=3, has_delta=True, **kwargs) + + return _safe_get_rope + + +class OmniModelState(DefaultModelState): + """Generic Omni ``ModelState`` — works for **all** Omni model stages. + + Model-owned preprocess and postprocess hooks supply stage-specific behavior. + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + # DefaultModelState.__init__ calls get_rope_state() which asserts + # isinstance(model, SupportsMRoPE). Two categories of Omni models: + # + # 1. Models that implement SupportsMRoPE (e.g. Qwen3-Omni Thinker): + # get_rope_state() succeeds normally, _safe_get_rope is a no-op. + # These models get correct 3D M-RoPE positions from the runner. + # + # 2. Models that do NOT implement SupportsMRoPE (e.g. Qwen3-TTS + # Talker, Code2Wav, FishSpeech): get_rope_state() would assert. + # These models compute their own position encoding internally + # (via model.forward kwargs or fixed 1D positions from + # InputBatch.positions), so rope_state = None is correct — + # DefaultModelState.prepare_inputs returns {} when rope_state + # is None, and upstream execute_model falls back to + # InputBatch.positions (1D sequential). + # Patch get_rope_state to handle Omni models that declare + # M-RoPE in config (mrope_section) but do not implement the + # SupportsMRoPE interface. For these models we create a + # RopeState with 3D sequential positions (matching V1 MR). + # + # The patch is applied via a class-level lock to prevent + # concurrent OmniModelState instances (e.g. different stages + # in a thread pool) from overwriting each other's patch. + from vllm.v1.worker.gpu.model_states import default as _default_mod + + with _rope_patch_lock: + orig_get_rope = _default_mod.get_rope_state + _default_mod.get_rope_state = _make_safe_get_rope(orig_get_rope) + try: + super().__init__(vllm_config, model, encoder_cache, device) + finally: + _default_mod.get_rope_state = orig_get_rope + max_num_reqs = self.scheduler_config.max_num_seqs + self.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) + self.has_preprocess: bool = getattr(model, "has_preprocess", False) + self.has_postprocess: bool = getattr(model, "has_postprocess", False) + self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False) + self._decode_preprocess = self._resolve_decode_preprocess(model) + self._mtp_generators: dict[str, torch.Generator] = {} + # 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 + + # Static inputs_embeds buffer for FULL CUDA graph — preprocess fills it in-place each step. + self._static_inputs_embeds: torch.Tensor | None = None + if self._embed_dim > 0: + self._static_inputs_embeds = torch.zeros( + (self.max_num_tokens, self._embed_dim), + dtype=self.dtype, + device=device, + ) + + # Static MTP buffers so _run_batched_mtp uses .copy_() instead of torch.cat(). + self._mtp_input_ids: torch.Tensor | None = None + self._mtp_input_embeds: torch.Tensor | None = None + self._mtp_hidden: torch.Tensor | None = None + self._mtp_text_step: torch.Tensor | None = None + self._mtp_offsets: torch.Tensor | None = None + self._mtp_sample_uniforms: torch.Tensor | None = None + self._mtp_runner: Any | 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) + self._mtp_input_embeds = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + self._mtp_hidden = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + self._mtp_text_step = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + self._mtp_offsets = torch.zeros(max_bs, dtype=torch.long, device=device) + sample_steps = int(getattr(model, "mtp_sample_steps", 0)) + sample_vocab_size = int(getattr(model, "mtp_sample_vocab_size", 0)) + if bool(getattr(model, "mtp_sample_uniforms", False)) and sample_steps > 0 and sample_vocab_size > 0: + self._mtp_sample_uniforms = torch.empty( + (max_bs, sample_steps, sample_vocab_size), + dtype=torch.float32, + device=device, + ) + self._mtp_runner = self._init_mtp_runner(model) + + @staticmethod + def _resolve_decode_preprocess(model: nn.Module) -> Callable | None: + """Adapt existing V1 hooks once, outside the per-step hot path. + + Keep the V2 hook precedence for models with an optimized implementation. + Both hooks return the same five values; only the embedding argument is + optional in the original contract. Errors inside a hook propagate. + """ + hook = getattr(model, "preprocess_decode_batch_mrv2", None) + if not callable(hook): + hook = getattr(model, "preprocess_decode_batch", None) + if not callable(hook): + return None + parameters = inspect.signature(hook).parameters + if "input_embeds" in parameters or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in parameters.values()): + return hook + + def decode_preprocess(*, input_ids, input_embeds, req_infos): + return hook(input_ids=input_ids, req_infos=req_infos) + + return decode_preprocess + + def _init_mtp_runner(self, model: nn.Module) -> Any: + mtp = getattr(model, "mtp", None) + if mtp is None: + return None + output_key = getattr(model, "mtp_output_key", None) + if not (isinstance(output_key, str) or (isinstance(output_key, tuple) and len(output_key) == 2)): + raise TypeError("Models with an mtp hook must declare mtp_output_key as a string or 2-tuple") + + compilation_config = self.vllm_config.compilation_config + cudagraph_mode = getattr(compilation_config, "cudagraph_mode", CUDAGraphMode.NONE) + if bool(getattr(model, "mtp_disable_graph", False)): + logger.info("Skipping mtp graph wrapper because the model marks it graph-unsafe.") + return mtp + graph_safe = bool(getattr(model, "mtp_graph_safe", False)) + if cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() and graph_safe: + graph_wrapper_cls = current_omni_platform.get_graph_wrapper_cls() + return graph_wrapper_cls(mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL) + return mtp + + def _is_mtp_graph_runner(self) -> bool: + runner = getattr(self, "_mtp_runner", None) + if runner is None: + return False + graph_wrapper_cls = current_omni_platform.get_graph_wrapper_cls() + return isinstance(runner, graph_wrapper_cls) + + def capture_mtp_graphs(self, dispatch_batch_descriptor: Callable[[int], Any]) -> None: + if not self._is_mtp_graph_runner(): + return + if ( + self._mtp_input_ids is None + or self._mtp_input_embeds is None + or self._mtp_hidden is None + or self._mtp_text_step is None + ): + return + if getattr(self.model, "mtp_accepts_req_infos", False): + logger.warning("Skipping mtp graph capture because this model requires per-request req_infos.") + return + + from vllm.compilation.monitor import set_cudagraph_capturing_enabled + from vllm.distributed.parallel_state import graph_capture + from vllm.forward_context import set_forward_context + + compilation_config = self.vllm_config.compilation_config + capture_sizes = self._get_mtp_capture_sizes() + num_warmups = compilation_config.cudagraph_num_of_warmups + capture_kwargs = self._get_mtp_base_sampling_kwargs() + logger.info("Capturing MRv2 mtp graphs for sizes %s", capture_sizes) + + set_cudagraph_capturing_enabled(True) + try: + with torch.inference_mode(), graph_capture(device=self.device): + for bsz in capture_sizes: + batch_descriptor = dispatch_batch_descriptor(int(bsz)) + num_tokens = int(getattr(batch_descriptor, "num_tokens", bsz)) + ids = self._mtp_input_ids[:num_tokens] + emb = self._mtp_input_embeds[:num_tokens] + hidden = self._mtp_hidden[:num_tokens] + text_step = self._mtp_text_step[:num_tokens] + + for _ in range(num_warmups): + with set_forward_context( + None, + self.vllm_config, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + batch_descriptor=batch_descriptor, + ): + self._call_mtp_runner(ids, emb, hidden, text_step, **capture_kwargs) + + with set_forward_context( + None, + self.vllm_config, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + batch_descriptor=batch_descriptor, + ): + sample_uniforms = None + if self._mtp_sample_uniforms is not None: + sample_uniforms = self._mtp_sample_uniforms[:num_tokens] + sample_uniforms.fill_(0.5) + self._call_mtp_runner( + ids, + emb, + hidden, + text_step, + sample_uniforms=sample_uniforms, + **capture_kwargs, + ) + torch.accelerator.synchronize() + + logger.info("Captured MRv2 mtp graphs for %d sizes", len(capture_sizes)) + except RuntimeError as e: + raise RuntimeError(f"MRv2 mtp graph capture failed: {e}") from e + finally: + set_cudagraph_capturing_enabled(False) + + def _get_mtp_capture_sizes(self) -> list[int]: + """Return graph buckets that a request-batched Talker can reach.""" + max_num_reqs = int(self.scheduler_config.max_num_seqs) + return sorted( + { + int(size) + for size in self.vllm_config.compilation_config.cudagraph_capture_sizes + if 0 < int(size) <= max_num_reqs + }, + reverse=True, + ) + + @staticmethod + def _get_embed_dim(model: nn.Module, device: torch.device) -> int: + """Return the embedding dim that ``embed_input_ids`` produces (may differ from hf_text_config).""" + if hasattr(model, "embed_input_ids"): + dummy = torch.zeros(1, dtype=torch.long, device=device) + with torch.no_grad(): + out = model.embed_input_ids(dummy) + return out.shape[-1] + return 0 + + # ------------------------------------------------------------------ + # Request lifecycle + # ------------------------------------------------------------------ + + 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._initialize_upstream_warmup_buffer(req_index, new_req_data.req_id) + + def _initialize_upstream_warmup_buffer(self, req_index: int, req_id: str) -> None: + """Satisfy model-declared output contracts for vLLM warmup requests.""" + if not str(req_id).startswith("_warmup_"): + return + + validity_key = getattr(self.model, "mtp_validity_key", None) + if validity_key is None: + return + + validity = torch.zeros((), dtype=torch.bool) + buffer = self.intermediate_buffer.buffers[req_index] + if isinstance(validity_key, tuple) and len(validity_key) == 2: + buffer.setdefault(validity_key[0], {})[validity_key[1]] = validity + elif isinstance(validity_key, str): + buffer[validity_key] = validity + else: + raise TypeError( + f"mtp_validity_key must be a string or 2-tuple, got {type(validity_key).__name__}: {validity_key!r}" + ) + + def _resolve_req_index(self, req_index_or_id: int | str) -> int | None: + if isinstance(req_index_or_id, int): + return req_index_or_id + + return self.intermediate_buffer.req_id_to_index.get(req_index_or_id) + + def remove_request(self, req_index_or_id: int | str) -> None: + req_index = self._resolve_req_index(req_index_or_id) + if req_index is None: + return + 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.intermediate_buffer.remove_request(req_index) + + # ------------------------------------------------------------------ + # Input preparation + # ------------------------------------------------------------------ + + def prepare_inputs(self, input_batch: InputBatch, req_states: RequestState) -> dict[str, Any]: + # Forward-only stages have no preprocess hook. Honor their declared + # GPU inputs before gathering the model's runtime view, so payload + # tensors do not fall back to per-request blocking device copies. + if not self.has_preprocess and getattr(self.model, "batched_gpu_staging_keys", None): + self._stage_batched_preprocess_inputs( + [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)], + input_batch.input_ids.device, + ) + base = super().prepare_inputs(input_batch, req_states) + buffer_list = self.intermediate_buffer.gather(input_batch) + base["model_intermediate_buffer"] = buffer_list + base["runtime_additional_information"] = buffer_list + base["seq_token_counts"] = [int(input_batch.num_scheduled_tokens[i]) for i in range(input_batch.num_reqs)] + # Return static inputs_embeds so FULL graph replay uses the same + # tensor address that was captured. Preprocess fills it in-place. + if self._static_inputs_embeds is not None: + base["inputs_embeds"] = self._static_inputs_embeds[: input_batch.num_tokens_after_padding] + return base + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + base = super().prepare_dummy_inputs(num_reqs, num_tokens) + dummy_buffer: list[dict[str, Any]] = [{} for _ in range(num_reqs)] + base["model_intermediate_buffer"] = dummy_buffer + base["runtime_additional_information"] = dummy_buffer + if num_reqs > 0: + per_req = num_tokens // num_reqs + remainder = num_tokens % num_reqs + counts = [per_req] * num_reqs + counts[-1] += remainder + base["seq_token_counts"] = counts + else: + base["seq_token_counts"] = [] + # Return static inputs_embeds for FULL graph capture so the graph + # captures this tensor's address. + if self._static_inputs_embeds is not None: + base["inputs_embeds"] = self._static_inputs_embeds[:num_tokens] + return base + + # ------------------------------------------------------------------ + # Pre-forward: per-request preprocess + batched MTP + # ------------------------------------------------------------------ + + @staticmethod + def _get_req_state_value(field: Any, req_idx: int) -> int | None: + values = getattr(field, "np", field) + try: + return int(values[req_idx]) + except Exception: + return None + + @staticmethod + def _get_input_batch_num_computed(input_batch: InputBatch, req_idx: int, batch_idx: int) -> int | None: + for attr in ("num_computed_tokens_cpu", "num_computed_tokens_np"): + values = getattr(input_batch, attr, None) + if values is None: + continue + try: + if len(values) == getattr(input_batch, "num_reqs", len(values)): + return int(values[batch_idx]) + return int(values[req_idx]) + except Exception: + continue + return None + + @staticmethod + def _preprocess_result_needs_writeback(original: torch.Tensor, updated: torch.Tensor) -> bool: + return updated is not original + + @staticmethod + def _batch_move_tensor_rows( + tensors: list[torch.Tensor], + device: torch.device, + ) -> list[torch.Tensor]: + """Move compatible variable-length rows with one device transfer.""" + lengths = [int(tensor.shape[0]) for tensor in tensors] + staged = torch.cat(tensors, dim=0).to(device=device, non_blocking=True) + return list(staged.split(lengths, dim=0)) + + def _stage_batched_preprocess_inputs( + self, + req_indices: list[int], + device: torch.device, + ) -> None: + staging_keys: set[tuple[str, str]] = getattr(self.model, "batched_gpu_staging_keys", set()) + for key in staging_keys: + if not isinstance(key, tuple) or len(key) != 2: + raise TypeError(f"batched_gpu_staging_keys entries must be 2-tuples, got {key!r}") + type_key, qualifier = key + groups: dict[ + tuple[torch.device, torch.dtype, tuple[int, ...]], list[tuple[dict[str, Any], torch.Tensor]] + ] = {} + for req_idx in req_indices: + nested = self.intermediate_buffer.buffers[req_idx].get(type_key) + if not isinstance(nested, dict): + continue + tensor = nested.get(qualifier) + if not isinstance(tensor, torch.Tensor) or tensor.device == device: + continue + group_key = tensor.device, tensor.dtype, tuple(tensor.shape[1:]) + groups.setdefault(group_key, []).append((nested, tensor)) + + for group in groups.values(): + staged_rows = self._batch_move_tensor_rows([tensor for _nested, tensor in group], device) + for (nested, _tensor), staged in zip(group, staged_rows, strict=True): + nested[qualifier] = staged + + def _write_batched_gpu_preprocess_updates( + self, + req_indices: list[int], + updates_by_req: list[dict[str, Any]], + gpu_keys: set[Any], + ) -> None: + """Snapshot compatible resident preprocess values once per cohort.""" + if not req_indices or len(req_indices) != len(updates_by_req): + return + + for key in gpu_keys: + nested_key = isinstance(key, tuple) and len(key) == 2 + values: list[torch.Tensor] = [] + for updates in updates_by_req: + if nested_key: + type_key, qualifier = key + nested = updates.get(type_key) + value = nested.get(qualifier) if isinstance(nested, dict) else None + else: + value = updates.get(key) + if not isinstance(value, torch.Tensor): + break + values.append(value) + if len(values) != len(updates_by_req): + continue + + first = values[0] + if any( + value.shape != first.shape or value.dtype != first.dtype or value.device != first.device + for value in values[1:] + ): + continue + + stacked = torch.stack(values, dim=0) + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + key, + stacked, + keepdim=False, + ) + + for updates in updates_by_req: + if nested_key: + type_key, qualifier = key + nested = updates[type_key] + del nested[qualifier] + if not nested: + del updates[type_key] + else: + del updates[key] + + def run_preprocess( + self, + input_batch: InputBatch, + model_inputs: dict[str, Any], + req_states: RequestState | None = None, + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None, + ) -> None: + """Per-request preprocess + MTP before model forward. + + Modifies ``model_inputs["input_ids"]`` and ``model_inputs["inputs_embeds"]`` + in-place. Collects decode-step MTP inputs and runs a single batched MTP + forward at the end. + """ + if not self.has_preprocess: + return + + input_ids = model_inputs.get("input_ids") + if input_ids is None: + input_ids = input_batch.input_ids + + embeds = model_inputs.get("inputs_embeds") + if embeds is None: + embeds = self.model.embed_input_ids(input_batch.input_ids[: input_batch.num_tokens]) + model_inputs["inputs_embeds"] = embeds + elif self._static_inputs_embeds is not None and embeds.data_ptr() == self._static_inputs_embeds.data_ptr(): + # FULL graph replay requires a stable inputs_embeds address. Refresh + # the active rows from the current token ids before model-specific + # preprocessing; otherwise decode reuses embeddings left by the + # previous step at the same static address. + num_active_tokens = int(getattr(input_batch, "num_tokens", input_ids.shape[0])) + token_embeds = self.model.embed_input_ids(input_ids[:num_active_tokens]) + embeds[:num_active_tokens].copy_(token_embeds) + + gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set()) + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]] = [] + prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None = None + + req_indices = [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)] + self._stage_batched_preprocess_inputs(req_indices, embeds.device) + + preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = [] + for i, req_idx in enumerate(req_indices): + buf = self.intermediate_buffer.buffers[req_idx] + if not buf or "req_id" not in buf: + continue + if str(buf["req_id"]).startswith("_warmup_"): + continue + + start = int(input_batch.query_start_loc_np[i]) + n_tok = int(input_batch.num_scheduled_tokens[i]) + + info = {key: value for key, value in buf.items() if isinstance(key, str)} + prompt_len = None + num_computed_tokens = None + if req_states is not None: + prompt_len = self._get_req_state_value(getattr(req_states, "prompt_len", None), req_idx) + num_computed_tokens = self._get_req_state_value( + getattr(req_states, "num_computed_tokens", None), + req_idx, + ) + if num_computed_tokens is None: + num_computed_tokens = self._get_input_batch_num_computed(input_batch, req_idx, i) + if prompt_len is not None: + info["_omni_prompt_len"] = int(prompt_len) + if num_computed_tokens is not None: + info["_omni_num_computed_tokens"] = int(num_computed_tokens) + if prompt_len is not None and num_computed_tokens is not None: + info["_omni_is_prefill"] = int(num_computed_tokens) < int(prompt_len) + + is_prefill = ( + int(num_computed_tokens) < int(prompt_len) + if prompt_len is not None and num_computed_tokens is not None + else n_tok > 1 + ) + preprocess_entries.append((i, req_idx, start, n_tok, info, is_prefill)) + + preprocess_batch_mrv2 = getattr(self.model, "preprocess_batch_mrv2", None) + if callable(preprocess_batch_mrv2): + prefill_infos = [ + self.intermediate_buffer.buffers[req_idx] + for _i, req_idx, _start, _n_tok, _info, is_prefill in preprocess_entries + if is_prefill + ] + if prefill_infos: + preprocess_batch_mrv2(req_infos=prefill_infos, device=embeds.device) + for _i, req_idx, _start, _n_tok, info, is_prefill in preprocess_entries: + if not is_prefill: + continue + runtime_info = {key: value for key, value in info.items() if key.startswith("_omni_")} + info.update(self.intermediate_buffer.buffers[req_idx]) + info.update(runtime_info) + + if not hasattr(self, "_decode_preprocess"): + self._decode_preprocess = self._resolve_decode_preprocess(self.model) + batch_decode_preprocess = self._decode_preprocess + decode_entries = [entry for entry in preprocess_entries if entry[3] == 1 and not entry[5]] + batched_decode_indices: set[int] = set() + if callable(batch_decode_preprocess) and decode_entries: + batch_size = len(decode_entries) + starts = [entry[2] for entry in decode_entries] + is_contiguous_cohort = all(entry[0] == row and entry[2] == row for row, entry in enumerate(decode_entries)) + if is_contiguous_cohort: + batch_ids = input_ids[:batch_size] + batch_embeds = embeds[:batch_size] + batch_offsets = None + else: + batch_offsets = torch.as_tensor(starts, device=input_ids.device, dtype=torch.long) + batch_ids = input_ids.index_select(0, batch_offsets) + batch_embeds = embeds.index_select(0, batch_offsets) + + new_ids, new_embeds, batch_hidden, batch_text_step, updates_by_req = batch_decode_preprocess( + input_ids=batch_ids, + input_embeds=batch_embeds, + req_infos=[entry[4] for entry in decode_entries], + ) + if len(updates_by_req) != batch_size: + raise RuntimeError( + "Batched Omni preprocess returned the wrong update count: " + f"expected={batch_size} actual={len(updates_by_req)}" + ) + if batch_hidden.shape[0] != batch_size or batch_text_step.shape[0] != batch_size: + raise RuntimeError( + "Batched Omni preprocess changed the request axis: " + f"expected={batch_size} hidden={batch_hidden.shape[0]} text_step={batch_text_step.shape[0]}" + ) + + if self._preprocess_result_needs_writeback(batch_ids, new_ids): + if batch_offsets is None: + input_ids[:batch_size].copy_(new_ids.reshape(-1)[:batch_size]) + else: + input_ids.index_copy_(0, batch_offsets, new_ids.reshape(-1)[:batch_size]) + if self._preprocess_result_needs_writeback(batch_embeds, new_embeds): + reshaped_embeds = new_embeds.reshape(batch_size, -1) + if batch_offsets is None: + embeds[:batch_size].copy_(reshaped_embeds) + else: + embeds.index_copy_(0, batch_offsets, reshaped_embeds) + + self._write_batched_gpu_preprocess_updates( + [entry[1] for entry in decode_entries], + updates_by_req, + gpu_keys, + ) + for row, (entry, updates) in enumerate(zip(decode_entries, updates_by_req, strict=True)): + i, req_idx, start, _n_tok, _info, _is_prefill = entry + mtp_batches.append( + ( + i, + start, + (batch_hidden[row : row + 1], batch_text_step[row : row + 1]), + ) + ) + batched_decode_indices.add(i) + state_updates = [(entry[1], updates) for entry, updates in zip(decode_entries, updates_by_req, strict=True)] + for req_idx, updates in state_updates: + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + prepacked_mtp_inputs = batch_hidden, batch_text_step + + for i, req_idx, start, n_tok, info, _is_prefill in preprocess_entries: + if i in batched_decode_indices: + continue + + ids_slice = input_ids[start : start + n_tok] + emb_slice = embeds[start : start + n_tok] + new_ids, new_emb, updates = self.model.preprocess(ids_slice, emb_slice, **info) + + # Write back in-place + seg = min(n_tok, new_ids.shape[0]) + if self._preprocess_result_needs_writeback(ids_slice, new_ids): + input_ids[start : start + seg] = new_ids[:seg] + if self._preprocess_result_needs_writeback(emb_slice, new_emb): + embeds[start : start + seg] = new_emb[:seg] + + # Collect MTP inputs for decode steps (n_tok == 1 with mtp_inputs) + mtp_inputs = updates.pop("mtp_inputs", None) + if mtp_inputs is not None and n_tok == 1: + mtp_batches.append((i, start, mtp_inputs)) + + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + if mtp_batches and hasattr(self.model, "mtp"): + if prepacked_mtp_inputs is None: + self._run_batched_mtp( + mtp_batches, + input_ids, + embeds, + input_batch, + gpu_keys, + mtp_batch_descriptor_dispatcher, + ) + else: + self._run_batched_mtp( + mtp_batches, + input_ids, + embeds, + input_batch, + gpu_keys, + mtp_batch_descriptor_dispatcher, + prepacked_mtp_inputs=prepacked_mtp_inputs, + ) + + def _pack_mtp_batch( + self, + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], + input_ids: torch.Tensor, + embeds: torch.Tensor, + offsets: torch.Tensor | None = None, + prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack per-request Talker state with a bounded number of GPU ops.""" + bsz = len(mtp_batches) + if offsets is None: + offsets = torch.as_tensor( + [start for _i, start, _mtp in mtp_batches], + device=input_ids.device, + dtype=torch.long, + ) + if prepacked_mtp_inputs is None: + hidden_rows = [past_hidden.reshape(1, -1) for _i, _start, (past_hidden, _step) in mtp_batches] + text_step_rows = [text_step.reshape(1, -1) for _i, _start, (_hidden, text_step) in mtp_batches] + packed_hidden = None + packed_text_step = None + else: + packed_hidden, packed_text_step = prepacked_mtp_inputs + packed_hidden = packed_hidden.reshape(bsz, -1) + packed_text_step = packed_text_step.reshape(bsz, -1) + + if self._mtp_input_ids is not None and bsz <= self._mtp_input_ids.shape[0]: + assert self._mtp_input_embeds is not None + assert self._mtp_hidden is not None + assert self._mtp_text_step is not None + batch_ids = self._mtp_input_ids[:bsz] + batch_embeds = self._mtp_input_embeds[:bsz] + batch_hidden = self._mtp_hidden[:bsz] + batch_text_step = self._mtp_text_step[:bsz] + if input_ids.dtype == batch_ids.dtype: + torch.index_select(input_ids, 0, offsets, out=batch_ids) + else: + batch_ids.copy_(input_ids.index_select(0, offsets)) + torch.index_select(embeds, 0, offsets, out=batch_embeds) + if prepacked_mtp_inputs is None: + torch.cat(hidden_rows, dim=0, out=batch_hidden) + torch.cat(text_step_rows, dim=0, out=batch_text_step) + else: + batch_hidden.copy_(packed_hidden) + batch_text_step.copy_(packed_text_step) + return batch_ids, batch_embeds, batch_hidden, batch_text_step, offsets + + if prepacked_mtp_inputs is None: + packed_hidden = torch.cat(hidden_rows, dim=0) + packed_text_step = torch.cat(text_step_rows, dim=0) + return ( + input_ids.index_select(0, offsets), + embeds.index_select(0, offsets), + packed_hidden, + packed_text_step, + offsets, + ) + + def _mtp_batch_offsets( + self, + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], + input_batch: InputBatch, + device: torch.device, + ) -> torch.Tensor: + """Reuse MRv2's device query offsets for a contiguous decode prefix. + + The steady-state Talker batch contains one decode token from every + scheduled request, in scheduler order. ``InputBatch.query_start_loc`` + already contains the exact device-side gather offsets for that shape; + rebuilding it from a Python list creates a small synchronous H2D + transfer on every decode step. + """ + bsz = len(mtp_batches) + query_start_loc = getattr(input_batch, "query_start_loc", None) + static_offsets = self._mtp_offsets + if ( + isinstance(query_start_loc, torch.Tensor) + and query_start_loc.device == device + and isinstance(static_offsets, torch.Tensor) + and static_offsets.device == device + and bsz <= static_offsets.shape[0] + and all(batch_index == expected for expected, (batch_index, _start, _mtp) in enumerate(mtp_batches)) + ): + offsets = static_offsets[:bsz] + 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, + ) + + def _run_batched_mtp( + self, + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], + input_ids: torch.Tensor, + embeds: torch.Tensor, + input_batch: InputBatch, + gpu_keys: set[str], + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None, + prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + """Batch MTP forward for all decode-step requests. + + Uses pre-allocated static buffers to avoid per-step torch.cat + memory allocations. + """ + from vllm.forward_context import set_forward_context + + bsz = len(mtp_batches) + batch_offsets = self._mtp_batch_offsets( + mtp_batches, + input_batch, + input_ids.device, + ) + batch_ids, batch_emb, batch_hidden, batch_step, batch_offsets = self._pack_mtp_batch( + mtp_batches, + input_ids, + embeds, + batch_offsets, + prepacked_mtp_inputs, + ) + + req_indices = [int(input_batch.idx_mapping_np[i]) for i, _start, _mtp in mtp_batches] + buffers = [self.intermediate_buffer.buffers[req_idx] for req_idx in req_indices] + req_ids = [str(buffer.get("req_id")) for buffer in buffers] + generators = [ + self._get_mtp_generator( + req_id, + buffer.get("sampling_params"), + batch_ids.device, + ) + for req_id, buffer in zip(req_ids, buffers, strict=True) + ] + + use_graph_runner = self._is_mtp_graph_runner() + has_explicit_generator = any(generator is not None for generator in generators) + sample_uniform_buffer = getattr(self, "_mtp_sample_uniforms", None) + can_use_graph = use_graph_runner and (not has_explicit_generator or sample_uniform_buffer is not None) + sample_uniforms = None + batch_descriptor = None + cudagraph_mode = CUDAGraphMode.NONE + num_tokens = bsz + if can_use_graph and mtp_batch_descriptor_dispatcher is not None: + batch_descriptor = mtp_batch_descriptor_dispatcher(bsz) + cudagraph_mode = getattr(batch_descriptor, "cg_mode", CUDAGraphMode.FULL) + num_tokens = int(getattr(batch_descriptor, "num_tokens", bsz)) + assert self._mtp_input_ids is not None + assert self._mtp_input_embeds is not None + assert self._mtp_hidden is not None + assert self._mtp_text_step is not None + batch_ids = self._mtp_input_ids[:num_tokens] + batch_emb = self._mtp_input_embeds[:num_tokens] + batch_hidden = self._mtp_hidden[:num_tokens] + batch_step = self._mtp_text_step[:num_tokens] + if sample_uniform_buffer is not None: + sample_uniforms = self._prepare_mtp_sample_uniforms(generators, num_tokens) + elif can_use_graph: + cudagraph_mode = CUDAGraphMode.FULL + if sample_uniform_buffer is not None: + sample_uniforms = self._prepare_mtp_sample_uniforms(generators, bsz) + with set_forward_context( + None, + self.vllm_config, + num_tokens=num_tokens, + cudagraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_descriptor, + ): + new_emb, codes = self._call_mtp_with_sampling( + batch_ids, + batch_emb, + batch_hidden, + batch_step, + buffers=buffers, + req_ids=req_ids, + generators=generators, + sample_uniforms=sample_uniforms, + ) + + embeds.index_copy_(0, batch_offsets, new_emb[:bsz].reshape(bsz, -1)) + audio_key = getattr(self.model, "mtp_output_key", None) + validity_key = getattr(self.model, "mtp_validity_key", None) + valid_rows = None + if codes is not None and validity_key is not None: + valid_rows = torch.ones((bsz,), dtype=torch.bool, device=codes.device) + if codes is not None and audio_key in gpu_keys: + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + audio_key, + codes[:bsz], + ) + if valid_rows is not None: + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + validity_key, + valid_rows, + keepdim=False, + ) + return + for j, (i, _start, _) in enumerate(mtp_batches): + if codes is None: + continue + req_idx = int(input_batch.idx_mapping_np[i]) + if isinstance(audio_key, tuple) and len(audio_key) == 2: + updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}} + elif isinstance(audio_key, str): + updates = {audio_key: codes[j : j + 1]} + else: + raise TypeError( + f"mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}" + ) + if valid_rows is not None: + if isinstance(validity_key, tuple) and len(validity_key) == 2: + updates.setdefault(validity_key[0], {})[validity_key[1]] = valid_rows[j] + elif isinstance(validity_key, str): + updates[validity_key] = valid_rows[j] + else: + raise TypeError( + "mtp_validity_key must be a string or 2-tuple, " + f"got {type(validity_key).__name__}: {validity_key!r}" + ) + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + def _get_mtp_base_sampling_kwargs(self) -> dict[str, Any]: + sampling_params = getattr(self.model, "mtp_sampling_params", None) + if not isinstance(sampling_params, dict): + sampling_params = {} + return { + "do_sample": sampling_params.get("do_sample"), + "temperature": sampling_params.get("temperature"), + "top_k": sampling_params.get("top_k"), + "top_p": sampling_params.get("top_p"), + } + + def _get_mtp_sampling_kwargs( + self, + buffers: list[dict[str, Any]], + req_ids: list[str], + generators: list[torch.Generator | None], + ) -> dict[str, Any]: + kwargs = self._get_mtp_base_sampling_kwargs() + + if len(generators) == 1: + if generators[0] is not None: + kwargs["generator"] = generators[0] + elif any(generator is not None for generator in generators): + kwargs["generators"] = generators + + if getattr(self.model, "mtp_accepts_req_infos", False): + kwargs["req_ids"] = req_ids + kwargs["req_infos"] = buffers + + return kwargs + + def _get_mtp_generator( + self, + req_id: str, + 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) + if generator is None or generator.device != device: + generator = torch.Generator(device=device) + generator.manual_seed(seed) + cache[req_id] = generator + return generator + + def _prepare_mtp_sample_uniforms( + self, + generators: list[torch.Generator | None], + num_tokens: int, + ) -> torch.Tensor | None: + sample_uniforms = self._mtp_sample_uniforms + if sample_uniforms is None: + return None + bsz = len(generators) + if not 0 < bsz <= num_tokens <= sample_uniforms.shape[0]: + raise ValueError( + f"invalid Talker-MTP sampling buffer shape: bsz={bsz} num_tokens={num_tokens} " + f"capacity={sample_uniforms.shape[0]}" + ) + + active = sample_uniforms[:bsz] + if any(generator is not None for generator in generators): + # Consume one request-local stream for the whole residual span so the + # outer MTP graph can treat sampling noise as a stable tensor input. + # The draw granularity differs from the legacy per-step path, but + # request-local determinism and batch-order independence are kept. + for row, generator in enumerate(generators): + active[row : row + 1].uniform_(1e-20, 1.0 - 1e-20, generator=generator) + else: + active.uniform_(1e-20, 1.0 - 1e-20) + if num_tokens > bsz: + sample_uniforms[bsz:num_tokens].uniform_(1e-20, 1.0 - 1e-20) + return sample_uniforms[:num_tokens] + + def _call_mtp_with_sampling( + self, + batch_ids: torch.Tensor, + batch_emb: torch.Tensor, + batch_hidden: torch.Tensor, + batch_step: torch.Tensor, + *, + buffers: list[dict[str, Any]], + req_ids: list[str], + generators: list[torch.Generator | None], + sample_uniforms: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if sample_uniforms is None: + kwargs = self._get_mtp_sampling_kwargs(buffers, req_ids, generators) + else: + kwargs = self._get_mtp_base_sampling_kwargs() + kwargs["sample_uniforms"] = sample_uniforms + if getattr(self.model, "mtp_accepts_req_infos", False): + kwargs["req_ids"] = req_ids + kwargs["req_infos"] = buffers + + if ( + len(req_ids) > 1 + and "generators" in kwargs + and "sample_uniforms" not in kwargs + and not getattr(self.model, "mtp_accepts_per_row_generators", False) + ): + emb_chunks = [] + code_chunks = [] + for row, (req_id, buffer, generator) in enumerate(zip(req_ids, buffers, generators, strict=True)): + row_kwargs = self._get_mtp_sampling_kwargs( + [buffer], + [req_id], + [generator], + ) + row_emb, row_codes = self._call_mtp_runner( + batch_ids[row : row + 1], + batch_emb[row : row + 1], + batch_hidden[row : row + 1], + batch_step[row : row + 1], + **row_kwargs, + ) + emb_chunks.append(row_emb) + if row_codes is not None: + code_chunks.append(row_codes) + new_emb = torch.cat(emb_chunks) + codes = torch.cat(code_chunks) if code_chunks else None + return new_emb, codes + + return self._call_mtp_runner( + batch_ids, + batch_emb, + batch_hidden, + batch_step, + **kwargs, + ) + + def _call_mtp_runner( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_hidden: torch.Tensor, + text_step: torch.Tensor, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Explicit generators are request state, while the outer whole-MTP + # graph owns one captured RNG stream. Keep the transformer/code- + # predictor graphs, but call the raw MTP function so per-row streams + # remain deterministic and independent of scheduler batch makeup. + has_explicit_generator = kwargs.get("generator") is not None or kwargs.get("generators") is not None + runner = self.model.mtp if has_explicit_generator else getattr(self, "_mtp_runner", None) + runner = runner or self.model.mtp + return runner(input_ids, input_embeds, last_hidden, text_step, **kwargs) + + # ------------------------------------------------------------------ + # Post-forward: per-request postprocess + # ------------------------------------------------------------------ + + def run_postprocess(self, hidden_states: torch.Tensor, input_batch: InputBatch) -> None: + """Per-request postprocess after model forward. + + Extracts per-request updates from hidden_states and writes them + back to the intermediate buffer (e.g. ``last_talker_hidden``). + """ + if not self.has_postprocess: + return + gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set()) + batch_postprocess = getattr(self.model, "postprocess_batch_mrv2", None) + if callable(batch_postprocess) and input_batch.num_reqs: + req_indices = [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)] + buffers = [self.intermediate_buffer.buffers[req_idx] for req_idx in req_indices] + if all(buf and "req_id" in buf and not str(buf["req_id"]).startswith("_warmup_") for buf in buffers): + query_start_loc = getattr(input_batch, "query_start_loc", None) + if ( + isinstance(query_start_loc, torch.Tensor) + and query_start_loc.device == hidden_states.device + and query_start_loc.shape[0] >= input_batch.num_reqs + 1 + ): + last_token_indices = query_start_loc[1 : input_batch.num_reqs + 1] - 1 + else: + last_token_indices = torch.as_tensor( + [ + int(input_batch.query_start_loc_np[i]) + int(input_batch.num_scheduled_tokens[i]) - 1 + for i in range(input_batch.num_reqs) + ], + device=hidden_states.device, + dtype=torch.long, + ) + output_key, values = batch_postprocess( + hidden_states=hidden_states, + last_token_indices=last_token_indices, + ) + if output_key not in gpu_keys: + raise RuntimeError(f"Batched Omni postprocess output must be GPU-resident: key={output_key!r}") + owned_rows = isinstance(values, OwnedBatchTensor) + raw_values = values.tensor if owned_rows else values + if ( + not isinstance(raw_values, torch.Tensor) + or raw_values.ndim == 0 + or raw_values.shape[0] != input_batch.num_reqs + ): + actual = raw_values.shape[0] if isinstance(raw_values, torch.Tensor) and raw_values.ndim else 0 + raise RuntimeError( + "Batched Omni postprocess changed the request axis: " + f"expected={input_batch.num_reqs} actual={actual}" + ) + if owned_rows: + # Owned batch: share row views, no second snapshot. + self.intermediate_buffer.update_owned_gpu_tensor_rows( + req_indices, + output_key, + raw_values, + keepdim=False, + ) + else: + # Borrowed output (graph buffer or shared scratch): snapshot once. + self.intermediate_buffer.update_gpu_tensor_rows( + req_indices, + output_key, + raw_values, + keepdim=False, + ) + return + for i in range(input_batch.num_reqs): + req_idx = int(input_batch.idx_mapping_np[i]) + buf = self.intermediate_buffer.buffers[req_idx] + if not buf or "req_id" not in buf: + continue + if str(buf["req_id"]).startswith("_warmup_"): + continue + start = int(input_batch.query_start_loc_np[i]) + n_tok = int(input_batch.num_scheduled_tokens[i]) + h_slice = hidden_states[start : start + n_tok] + info = {key: value for key, value in buf.items() if isinstance(key, str) and key != "hidden_states"} + updates = self.model.postprocess(h_slice, **info) + if updates: + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + # ------------------------------------------------------------------ + # Output post-processing + # ------------------------------------------------------------------ + + def postprocess_model_output( + self, + model_output: Any, + input_batch: InputBatch, + req_states: RequestState, + ) -> tuple[torch.Tensor, dict]: + """Convert raw model output to ``(text_hidden, multimodal_outputs)``. + + Handles ``OmniOutput`` unwrapping and ``make_omni_output`` + conversion. + """ + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + if isinstance(model_output, (list, tuple)) or self.have_multimodal_outputs: + buffer_list = self.intermediate_buffer.gather(input_batch) + make_output_kwargs = { + "model_intermediate_buffer": buffer_list, + "request_token_spans": [ + ( + int(input_batch.query_start_loc_np[i]), + int(input_batch.query_start_loc_np[i]) + int(input_batch.num_scheduled_tokens[i]), + ) + for i in range(input_batch.num_reqs) + ], + } + make_output_kwargs["runtime_additional_information"] = buffer_list + model_output = self.model.make_omni_output(model_output, **make_output_kwargs) + + if isinstance(model_output, OmniOutput): + text_hidden = model_output.text_hidden_states + multimodal_outputs = cast(dict[str, Any], model_output.multimodal_outputs or {}) + elif isinstance(model_output, (list, tuple)): + text_hidden = model_output[0] + multimodal_outputs = {} + else: + text_hidden = model_output + multimodal_outputs = {} + + return text_hidden, multimodal_outputs diff --git a/vllm_omni/worker_v2/native_output_worker.py b/vllm_omni/worker_v2/native_output_worker.py new file mode 100644 index 000000000..3d6b7ad47 --- /dev/null +++ b/vllm_omni/worker_v2/native_output_worker.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Materialize native producer output independently of scheduler consumption.""" + +from concurrent.futures import Future, ThreadPoolExecutor +from threading import BoundedSemaphore, Lock +from typing import Any + +from vllm.v1.outputs import AsyncModelRunnerOutput + + +class NativeAsyncOutput(AsyncModelRunnerOutput): + """Consume an ordered publication on the engine's owner thread.""" + + def __init__(self, future: Future, plane: Any, copy_event: Any) -> None: + self._future: Future | None = future + self._plane = plane + self.copy_event = copy_event + self._resolved: Future = Future() + + def get_output(self) -> Any: + # Like the upstream lazy AsyncOutputFuture, consumption belongs to one + # engine thread. Cache finalization and errors for repeated consumption. + if not self._resolved.done(): + try: + assert self._future is not None + output = self._future.result() + output.omni_connector_output = self._plane.get_omni_connector_output() + self._resolved.set_result(output) + except Exception as error: + self._resolved.set_exception(error) + finally: + if self._resolved.done(): + # Do not retain raw snapshots or a callback bound to self. + self._future = None + self._plane = None + return self._resolved.result() + + +class NativeOutputWorker: + """Bounded FIFO CPU materialization and native-queue publication. + + Input outputs must have no owner-thread finalizer installed. Reservations + and metadata writes must be complete before submit. This worker never + updates scheduler state or drains connector notifications. + """ + + def __init__(self, capacity: int, device: Any | None = None) -> None: + if capacity < 1: + raise ValueError("Native output capacity must be positive") + self._slots = BoundedSemaphore(capacity) + self._lock = Lock() + self._closed = False + self._executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="omni-native-materialize", + initializer=self._initialize_device, + initargs=(device,), + ) + + @staticmethod + def _initialize_device(device: Any | None) -> None: + if device is not None: + from vllm.platforms import current_platform + + current_platform.set_device(device) + + @staticmethod + def _materialize(output: AsyncModelRunnerOutput, plane: Any) -> Any: + result = output.get_output() + plane.enqueue_outputs( + req_ids=list(result.req_ids), + inter_stage_outputs=getattr(result, "inter_stage_outputs", None), + sampled_token_ids=getattr(result, "sampled_token_ids", None), + ) + result.inter_stage_outputs = None + return result + + def submit(self, output: AsyncModelRunnerOutput, plane: Any) -> NativeAsyncOutput: + with self._lock: + if self._closed: + raise RuntimeError("Native output worker is closed") + self._slots.acquire() + try: + with self._lock: + if self._closed: + raise RuntimeError("Native output worker is closed") + future = self._executor.submit(self._materialize, output, plane) + except BaseException: + self._slots.release() + raise + future.add_done_callback(lambda _: self._slots.release()) + return NativeAsyncOutput(future, plane, output.copy_event) + + def close(self) -> None: + with self._lock: + self._closed = True + # Drain submitted publications before the native data plane is closed. + self._executor.shutdown(wait=True) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py new file mode 100644 index 000000000..23cc84f85 --- /dev/null +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -0,0 +1,818 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniARModelRunner — autoregressive stage runner on MR V2. + +Extends ``OmniGPUModelRunner`` with: + +* ``OmniOutput`` post-processing in ``sample_tokens`` +* Per-request ``pooler_output`` construction (hidden + multimodal slices) +* Async D2H copy via ``OmniAsyncOutput`` for non-blocking output transfer +* Cross-stage KV extraction before state cleanup +""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any, cast + +import numpy as np +import torch +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + get_ep_all2all_manager, +) +from vllm.utils.torch_utils import PIN_MEMORY +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput, RoutedExpertsTensors +from vllm.v1.worker.gpu.eplb_utils import step_eplb_after + +from vllm_omni.data_entry_keys import OmniPayload, flatten_payload +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVTransferManager, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput +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.output_snapshot import PackedOutputSnapshot, pack_output_snapshot + +logger = init_logger(__name__) +_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64 + + +def _copy_mm_to_snapshot_slot(value: Any, slot: dict[tuple[Any, ...], torch.Tensor], path: tuple[Any, ...] = ()) -> Any: + if not path and isinstance(value, dict): + packed = pack_output_snapshot(value, slot, max_buckets=_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT) + if packed is not None: + return packed + if isinstance(value, torch.Tensor): + bucket_key = path + (tuple(value.shape), value.dtype, value.device) + cached = slot.get(bucket_key) + if cached is None: + if len(slot) >= _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT: + return value.detach().clone() + cached = torch.empty_like(value) + slot[bucket_key] = cached + cached.copy_(value) + return cached + if isinstance(value, dict): + return {key: _copy_mm_to_snapshot_slot(item, slot, path + (key,)) for key, item in value.items()} + if isinstance(value, list): + return [_copy_mm_to_snapshot_slot(item, slot, path + (index,)) for index, item in enumerate(value)] + if isinstance(value, tuple): + return tuple(_copy_mm_to_snapshot_slot(item, slot, path + (index,)) for index, item in enumerate(value)) + return value + + +def _has_cuda_tensor(value: Any) -> bool: + if isinstance(value, torch.Tensor): + return value.device.type == "cuda" + if isinstance(value, dict): + return any(_has_cuda_tensor(item) for item in value.values()) + if isinstance(value, (list, tuple)): + return any(_has_cuda_tensor(item) for item in value) + return False + + +def _guard_graph_replay_for_pooler_copy( + main_stream: Any, + copy_event: Any, + *, + need_pooler: bool, + async_chunk: bool, +) -> None: + """Keep static graph outputs alive until non-snapshotted D2H completes.""" + if need_pooler and not async_chunk: + main_stream.wait_event(copy_event) + + +class OmniARModelRunner(OmniGPUModelRunner): + """AR stage runner. Produces per-request hidden states + multimodal outputs.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.kv_transfer_manager: OmniKVTransferManager | None = None + self._kv_extracted_req_ids: list[str] | None = None + self._async_mm_snapshot_slots: list[dict[tuple[Any, ...], torch.Tensor]] = [{} for _ in range(4)] + self._async_mm_snapshot_events: list[torch.cuda.Event | None] = [None] * 4 + self._async_mm_snapshot_pending = [False] * 4 + self._async_mm_snapshot_cursor = 0 + self._last_multimodal_snapshot_slot: int | None = None + + def _ensure_kv_transfer_manager(self) -> OmniKVTransferManager: + if self.kv_transfer_manager is None: + self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + return self.kv_transfer_manager + + # ------------------------------------------------------------------ + # execute_model: KV transfer pre-hook + delegate to super + # ------------------------------------------------------------------ + + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: Any | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + context_len: int = 0, + ) -> Any: + if not dummy_run: + self._handle_kv_transfer_pre(scheduler_output) + return super().execute_model( + scheduler_output, + intermediate_tensors, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + is_profile=is_profile, + context_len=context_len, + ) + + # ------------------------------------------------------------------ + # sample_tokens: OmniOutput handling + pooler_output + async D2H + # ------------------------------------------------------------------ + + @torch.inference_mode() + @step_eplb_after() + def sample_tokens( + self, grammar_output: GrammarOutput | None + ) -> OmniAsyncOutput | OmniModelRunnerOutput | ModelRunnerOutput | None: + kv_extracted = self._kv_extracted_req_ids + self._kv_extracted_req_ids = None + + if self.execute_model_state is None: + return None + + input_batch = self.execute_model_state.input_batch + hidden_states = self.execute_model_state.hidden_states + finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output + routed_experts = self.execute_model_state.routed_experts + self.execute_model_state = None + + # --- Omni: reconstruct raw model output and post-process --- + aux = self._last_aux_output + self._last_aux_output = None + multimodal_outputs = self._last_multimodal_outputs + self._last_multimodal_outputs = None + snapshot_slot = None + self._last_multimodal_snapshot_slot = None + raw_output = self._reconstruct_raw_model_output( + hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + aux=aux, + ) + text_hidden, multimodal_outputs = self.model_state.postprocess_model_output( + raw_output, input_batch, self.req_states + ) + if bool(getattr(self.model_config, "async_chunk", False)) and multimodal_outputs: + multimodal_outputs = self._retain_multimodal_outputs(multimodal_outputs) + snapshot_slot = self._last_multimodal_snapshot_slot + self._last_multimodal_snapshot_slot = None + + # --- Standard v2 sampling --- + sampling_context = nullcontext() + model_sampling_context = getattr(self.model, "mrv2_sampling_context", None) + if callable(model_sampling_context) and getattr(self.model, "logitsprocs_need_output_token_ids", False): + # seq_lens includes this forward's input, but not its sampled token: + # final prefill has zero prior outputs; the first decode has one. + sampling_context = model_sampling_context( + req_ids=input_batch.req_ids, + num_output_tokens=( + input_batch.seq_lens[: input_batch.num_reqs] + - self.req_states.prompt_len.gpu[input_batch.idx_mapping] + ), + ) + with sampling_context: + sampler_output, num_sampled, num_rejected = self.sample( + text_hidden, + input_batch, + grammar_output, + ) + if self.pp_handler is not None: + self.pp_handler.broadcast( + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch, + ) + + # --- Omni: prompt logprobs --- + assert self.prompt_logprobs_worker is not None + # Mirror the current parent GPUModelRunner call exactly. + prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs( + self.model.compute_logits, + text_hidden, + input_batch, + self.req_states.all_token_ids.gpu, + self.req_states.num_computed_tokens.gpu, + self.req_states.prompt_len.np, + ) + + # --- Omni: pooler_output --- + engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", "text") + need_pooler = engine_output_type != "text" + + # --- Build base output --- + model_runner_output = OmniModelRunnerOutput( + req_ids=input_batch.req_ids, + req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)}, + sampled_token_ids=None, # type: ignore[arg-type] + prompt_logprobs_dict=prompt_logprobs_dict, + kv_connector_output=None, + ) + model_runner_output.kv_extracted_req_ids = kv_extracted + model_runner_output._async_chunk = bool(getattr(self.model_config, "async_chunk", False)) + + # --- Async D2H via OmniAsyncOutput --- + materialize_native = self._uses_native_output_materializer() + async_output = OmniAsyncOutput( + model_runner_output=model_runner_output, + sampler_output=sampler_output, + num_sampled_tokens=num_sampled, + main_stream=self.main_stream, + copy_stream=self.output_copy_stream, + text_hidden=text_hidden if need_pooler else None, + multimodal_outputs=multimodal_outputs if need_pooler else None, + input_batch=input_batch if need_pooler else None, + async_chunk=bool(getattr(self.model_config, "async_chunk", False)), + finalize_output=(None if materialize_native else self._finalize_native_data_plane_output), + check_ep_fault=self.check_ep_fault, + routed_experts=routed_experts, + ) + self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event) + _guard_graph_replay_for_pooler_copy( + self.main_stream, + async_output.copy_event, + need_pooler=need_pooler, + async_chunk=bool(getattr(self.model_config, "async_chunk", False)), + ) + + # Postprocess AFTER creating async output (so copy_event is + # recorded before postprocess, matching upstream pattern). + self.postprocess_sampled( + input_batch.idx_mapping, + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch.query_start_loc, + ) + model_runner_output.kv_connector_output = self.kv_connector.post_forward(finished_req_ids) + model_runner_output.ec_connector_output = ec_connector_output + + self._reserve_native_data_plane_outputs(list(model_runner_output.req_ids)) + if materialize_native: + return self._materialize_native_output(async_output) + return async_output + + 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 + 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().") + event = self._async_mm_snapshot_events[slot_index] + if event is not None: + self.main_stream.wait_event(event) + if _has_cuda_tensor(outputs): + with torch.cuda.stream(self.main_stream): + snapshot = _copy_mm_to_snapshot_slot(outputs, self._async_mm_snapshot_slots[slot_index]) + else: + snapshot = _copy_mm_to_snapshot_slot(outputs, self._async_mm_snapshot_slots[slot_index]) + self._async_mm_snapshot_pending[slot_index] = True + self._last_multimodal_snapshot_slot = slot_index + self._async_mm_snapshot_cursor = (slot_index + 1) % len(self._async_mm_snapshot_slots) + return snapshot + + def _release_multimodal_snapshot(self, slot_index: int | None, copy_event: torch.cuda.Event) -> None: + if slot_index is None: + return + self._async_mm_snapshot_events[slot_index] = copy_event + self._async_mm_snapshot_pending[slot_index] = False + + # ------------------------------------------------------------------ + # pooler_output construction + # ------------------------------------------------------------------ + + @staticmethod + def _reconstruct_raw_model_output( + *, + hidden_states: torch.Tensor, + multimodal_outputs: dict[str, Any] | None, + aux: Any | None, + ) -> Any: + if multimodal_outputs: + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs=cast(OmniPayload, multimodal_outputs), + ) + if aux is not None: + return (hidden_states, aux) + return hidden_states + + @staticmethod + def _build_pooler_output_from_cpu( + hidden_cpu: torch.Tensor, + mm_cpu: dict[str, Any], + query_start_loc_np: np.ndarray, + num_scheduled_tokens: np.ndarray, + num_reqs: int, + ) -> list[dict[str, Any]]: + """Build pooler_output from already-CPU tensors.""" + total = hidden_cpu.shape[0] + pooler: list[dict[str, Any]] = [] + for i in range(num_reqs): + start = int(query_start_loc_np[i]) + end = start + int(num_scheduled_tokens[i]) + payload: dict[str, Any] = {"hidden": hidden_cpu[start:end].clone()} + for k, v in mm_cpu.items(): + payload[k] = _slice_pooler_value( + v, + req_index=i, + start=start, + end=end, + total_tokens=total, + ) + pooler.append(flatten_payload(payload)) + return pooler + + @staticmethod + def _build_async_chunk_outputs_from_mm( + mm_outputs: dict[str, Any], + query_start_loc_np: np.ndarray, + num_scheduled_tokens: np.ndarray, + num_reqs: int, + total_tokens: int, + padded_total_tokens: int | None = None, + ) -> tuple[list[dict[str, Any] | None] | None, list[dict[str, Any] | None] | None]: + """Build async-chunk payloads without materializing hidden on CPU.""" + if not mm_outputs: + return None, None + + inter_stage_list: list[dict[str, Any] | None] = [] + client_mm_list: list[dict[str, Any] | None] = [] + for i in range(num_reqs): + start = int(query_start_loc_np[i]) + end = start + int(num_scheduled_tokens[i]) + payload = { + key: _slice_pooler_value( + value, + req_index=i, + start=start, + end=end, + total_tokens=total_tokens, + padded_total_tokens=padded_total_tokens, + ) + for key, value in mm_outputs.items() + } + inter_stage, client_mm = partition_flat_payload(flatten_payload(payload)) + inter_stage_list.append(inter_stage or None) + client_mm_list.append(client_mm or None) + + return ( + None if all(item is None for item in inter_stage_list) else inter_stage_list, + None if all(item is None for item in client_mm_list) else client_mm_list, + ) + + # ------------------------------------------------------------------ + # KV transfer + # ------------------------------------------------------------------ + + def _handle_kv_transfer_pre(self, scheduler_output: SchedulerOutput) -> None: + finished: dict = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) + if not finished: + return + + kv_caches = getattr(self, "kv_caches", None) + if kv_caches is None: + return + + if hasattr(self.model, "get_kv_transfer_metadata"): + for req_id, data in finished.items(): + try: + meta = self.model.get_kv_transfer_metadata(req_id) + if meta: + existing = data.get("custom_metadata") or {} + existing.update(meta) + data["custom_metadata"] = existing + except Exception: + logger.warning( + "Failed to get KV transfer metadata for %s", + req_id, + exc_info=True, + ) + + mgr = self._ensure_kv_transfer_manager() + self._kv_extracted_req_ids = mgr.handle_finished_requests_kv_transfer( + finished_reqs=finished, + kv_caches=kv_caches, + block_size=self.cache_config.block_size, + cache_dtype=str(self.cache_config.cache_dtype), + request_id_resolver=self._resolve_global_request_id, + ) + + def _resolve_global_request_id(self, req_id: str) -> str: + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + return req_id + info = self.model_state.intermediate_buffer.buffers[req_idx] + global_id = info.get("global_request_id") + if isinstance(global_id, list): + global_id = global_id[0] if global_id else None + if isinstance(global_id, bytes): + return global_id.decode("utf-8") + return str(global_id) if global_id else req_id + + +# ====================================================================== +# OmniAsyncOutput — async D2H for Omni AR outputs +# ====================================================================== + + +def _async_copy_to_np( + x: torch.Tensor, + *, + copy_stream: torch.cuda.Stream | None = None, + pin_memory: bool | None = None, +) -> np.ndarray: + return _async_copy_tensor( + x, + copy_stream=copy_stream, + pin_memory=pin_memory, + ).numpy() + + +def _async_copy_tensor( + x: torch.Tensor, + *, + copy_stream: torch.cuda.Stream | None = None, + pin_memory: bool | None = None, +) -> torch.Tensor: + x = x.detach() + if x.device.type == "cpu": + return x.clone() + if pin_memory is None: + pin_memory = PIN_MEMORY + cpu = torch.empty_like(x, device="cpu", pin_memory=pin_memory) + cpu.copy_(x, non_blocking=pin_memory) + if copy_stream is None: + copy_stream = torch.cuda.current_stream() + x.record_stream(copy_stream) + return cpu + + +def _async_copy_mm_value( + value: Any, + *, + copy_stream: torch.cuda.Stream | None = None, + pin_memory: bool | None = None, +) -> Any: + if isinstance(value, torch.Tensor): + return _async_copy_tensor( + value, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + if isinstance(value, dict): + return { + key: _async_copy_mm_value( + val, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + for key, val in value.items() + } + if isinstance(value, list): + return [ + _async_copy_mm_value( + val, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + for val in value + ] + if isinstance(value, tuple): + return tuple( + _async_copy_mm_value( + val, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + for val in value + ) + return value + + +def _async_copy_mm( + mm_outputs: dict | None, + total_tokens: int, + *, + copy_stream: torch.cuda.Stream | None = None, + pin_memory: bool | None = None, +) -> dict[str, Any]: + """Non-blocking D2H copy of multimodal output tensors.""" + if not mm_outputs: + return {} + if isinstance(mm_outputs, PackedOutputSnapshot): + return mm_outputs.copy_to_cpu( + lambda tensor: _async_copy_tensor(tensor, copy_stream=copy_stream, pin_memory=pin_memory) + ) + return { + key: _async_copy_mm_value( + value, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + for key, value in mm_outputs.items() + } + + +def _slice_pooler_value( + value: Any, + *, + req_index: int, + start: int, + end: int, + total_tokens: int, + padded_total_tokens: int | None = None, + request_scoped: bool = False, +) -> Any: + if isinstance(value, torch.Tensor): + token_axis_sizes = {total_tokens} + if padded_total_tokens is not None: + token_axis_sizes.add(padded_total_tokens) + if not request_scoped and value.dim() > 0 and value.shape[0] in token_axis_sizes: + return value[start:end].contiguous() + return value.clone() + if isinstance(value, dict): + return { + key: _slice_pooler_value( + val, + req_index=req_index, + start=start, + end=end, + total_tokens=total_tokens, + padded_total_tokens=padded_total_tokens, + request_scoped=request_scoped, + ) + for key, val in value.items() + } + if isinstance(value, list): + if not value: + return [] + # Lists at the batch level contain request-owned payloads. Their + # tensors use local axes even when a size coincides with batch tokens. + values = value if request_scoped else [value[req_index] if req_index < len(value) else value[0]] + owned = [ + _slice_pooler_value( + item, + req_index=req_index, + start=start, + end=end, + total_tokens=total_tokens, + request_scoped=True, + ) + for item in values + ] + return owned if request_scoped else owned[0] + return value + + +def _ensure_tensor_values(payload: dict[str, Any]) -> dict[str, torch.Tensor]: + """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``. + + Non-tensor scalars (int/float/bool) are wrapped with ``torch.tensor()``; + values that cannot be safely converted are dropped. Enforces the tensor-only + invariant required by ``OmniEngineCoreOutput.multimodal_output`` (the channel + the async_chunk stage-input processor reads). Mirrors the V1 runner helper. + """ + result: dict[str, torch.Tensor] = {} + for key, val in payload.items(): + if isinstance(val, torch.Tensor): + result[key] = val + elif isinstance(val, (int, float, bool)): + result[key] = torch.tensor(val) + elif isinstance(val, (list, tuple)): + try: + result[key] = torch.tensor(val) + except (ValueError, TypeError, RuntimeError): + logger.warning( + "Dropping non-tensorizable multimodal output key '%s' (type=%s).", + key, + type(val).__name__, + ) + else: + logger.warning( + "Dropping non-tensor multimodal output key '%s' (type=%s).", + key, + type(val).__name__, + ) + return result + + +class OmniAsyncOutput(AsyncModelRunnerOutput): + """Async D2H copy for Omni AR model outputs. + + Mirrors upstream ``AsyncOutput`` but additionally handles + ``pooler_output`` (hidden states + multimodal outputs) via + non-blocking copies on the copy stream. + """ + + def __init__( + self, + model_runner_output: OmniModelRunnerOutput, + sampler_output: Any, + num_sampled_tokens: torch.Tensor, + main_stream: torch.cuda.Stream, + copy_stream: torch.cuda.Stream, + copy_event: torch.cuda.Event | None = None, + text_hidden: torch.Tensor | None = None, + multimodal_outputs: dict | None = None, + input_batch: Any | None = None, + async_chunk: bool = False, + finalize_output: Any | None = None, + check_ep_fault: bool = False, + routed_experts: RoutedExpertsTensors | None = None, + ): + self.model_runner_output = model_runner_output + self.sampler_output = sampler_output + self.num_sampled_tokens = num_sampled_tokens + self.routed_experts = routed_experts + self.copy_event = copy_event if copy_event is not None else torch.cuda.Event(blocking=True) + self._async_chunk = bool(async_chunk) + self._finalize_output = finalize_output + self._has_fault: torch.Tensor | None = None + + # Snapshot input_batch metadata needed for pooler_output slicing + self._need_pooler = text_hidden is not None or (self._async_chunk and bool(multimodal_outputs)) + self._query_start_loc_np: np.ndarray | None = None + self._num_scheduled_tokens: np.ndarray | None = None + self._num_reqs: int = 0 + self._total_tokens: int = 0 + self._padded_total_tokens: int = 0 + if self._need_pooler and input_batch is not None: + self._query_start_loc_np = input_batch.query_start_loc_np.copy() + self._num_scheduled_tokens = np.array(input_batch.num_scheduled_tokens, dtype=np.int32) + self._num_reqs = input_batch.num_reqs + if self._query_start_loc_np.shape[0] > self._num_reqs: + self._total_tokens = int(self._query_start_loc_np[self._num_reqs]) + else: + self._total_tokens = int(self._num_scheduled_tokens[: self._num_reqs].sum()) + self._padded_total_tokens = int(getattr(input_batch, "num_tokens_after_padding", self._total_tokens)) + + # Perform all D2H copies on the copy stream (non-blocking). + import contextlib + + @contextlib.contextmanager + def _stream(to_stream, from_stream): + try: + torch.cuda.set_stream(to_stream) + yield + finally: + torch.cuda.set_stream(from_stream) + + with _stream(copy_stream, main_stream): + copy_stream.wait_stream(main_stream) + pin_memory = PIN_MEMORY + + # Sampled token ids + self.sampled_token_ids_np = _async_copy_to_np( + sampler_output.sampled_token_ids, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + self.num_sampled_tokens_np = _async_copy_to_np( + num_sampled_tokens, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + self.sampling_mask_tensors = None + if sampler_output.sampling_mask_tensors is not None: + self.sampling_mask_tensors = sampler_output.sampling_mask_tensors.to_cpu_nonblocking() + self.routed_experts_cpu = None + if routed_experts is not None: + self.routed_experts_cpu = routed_experts.to_cpu_nonblocking() + + # Logprobs + self.logprobs_tensors = None + if sampler_output.logprobs_tensors is not None: + self.logprobs_tensors = sampler_output.logprobs_tensors.to_cpu_nonblocking() + self.num_nans: np.ndarray | None = None + if sampler_output.num_nans is not None: + self.num_nans = _async_copy_to_np( + sampler_output.num_nans, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + + # Prompt logprobs + self.prompt_logprobs_dict = { + k: v.to_cpu_nonblocking() if v is not None else None + for k, v in self.model_runner_output.prompt_logprobs_dict.items() + } + if check_ep_fault: + has_fault = get_ep_all2all_manager().query_fault() + self._has_fault = has_fault.to("cpu", non_blocking=True) + + # Pooler output (hidden + multimodal) — async D2H + self._hidden_cpu: torch.Tensor | None = None + self._mm_cpu: dict[str, Any] = {} + self._mm_snapshot: dict[str, Any] = {} + if self._need_pooler and self._async_chunk: + # 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, + ) + elif self._need_pooler and text_hidden is not None: + self._hidden_cpu = _async_copy_tensor( + text_hidden, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + total_tokens = text_hidden.shape[0] + self._mm_cpu = _async_copy_mm( + multimodal_outputs, + total_tokens, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) + + self.copy_event.record(copy_stream) + + def get_output(self) -> OmniModelRunnerOutput: + self.copy_event.synchronize() + + # Sampled token ids + sampled_token_ids: list[list[int]] = self.sampled_token_ids_np.tolist() + num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist() + for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens): + del token_ids[num_tokens:] + self.model_runner_output.sampled_token_ids = sampled_token_ids + if self.sampling_mask_tensors is not None: + self.model_runner_output.sampling_masks = self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np) + if self.routed_experts_cpu is not None: + self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists() + self.model_runner_output.sampled_token_ids_materialized = True + + # Logprobs + if self.num_nans is not None: + self.model_runner_output.num_nans_in_logits = dict( + zip(self.model_runner_output.req_ids, self.num_nans.tolist()) + ) + if self.logprobs_tensors is not None: + self.model_runner_output.logprobs = self.logprobs_tensors.tolists() + self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict + + if self._has_fault is not None and self._has_fault.item(): + mask = get_ep_all2all_manager().query_active_mask() + raise RuntimeError( + "Fault detected in EP all2all communication: one or more ranks " + f"timed out during dispatch/combine. Mask: {mask.cpu().tolist()}" + ) + + # Pooler output. Populate two channels from the same per-request payloads, + # mirroring the V1 runner: + # * pooler_output -> sync/full-payload path (inline pooling_output bridge) + # * multimodal_outputs -> wire multimodal_output, which the async_chunk + # stage-input processor (talker2code2wav_async_chunk) reads for codes. + if self._need_pooler and self._async_chunk: + pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm( + self._mm_snapshot, + self._query_start_loc_np, + self._num_scheduled_tokens, + self._num_reqs, + self._total_tokens, + self._padded_total_tokens, + ) + self.model_runner_output.pooler_output = None + self.model_runner_output.inter_stage_outputs = pooler_inter + self.model_runner_output.multimodal_outputs = ( + [_ensure_tensor_values(_async_copy_mm_value(p)) if p else {} for p in pooler_client] + if pooler_client + else None + ) + elif self._need_pooler and self._hidden_cpu is not None: + pooler_output = OmniARModelRunner._build_pooler_output_from_cpu( + self._hidden_cpu, + self._mm_cpu, + self._query_start_loc_np, + self._num_scheduled_tokens, + self._num_reqs, + ) + pooler_payload = cast(list[dict[str, Any] | None], pooler_output) if pooler_output else None + self.model_runner_output.pooler_output = pooler_payload + self.model_runner_output.inter_stage_outputs = pooler_payload + self.model_runner_output.multimodal_outputs = ( + [_ensure_tensor_values(p) if p else {} for p in pooler_payload] if pooler_payload else None + ) + + if self._finalize_output is not None: + return self._finalize_output(self.model_runner_output) + return self.model_runner_output diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py new file mode 100644 index 000000000..c8e5e57cb --- /dev/null +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -0,0 +1,604 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import threading +import time +from collections import defaultdict +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass, field +from queue import Full, Queue +from types import SimpleNamespace +from typing import Any, cast + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import OmniPayload, unflatten_payload +from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, +) +from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans +from vllm_omni.worker_v2.delivery import ( + DeliveryTimeoutError, + OmniDeliveryManager, +) + + +@dataclass +class _NativeRequestState: + request_id: str + external_req_id: str + prompt_token_ids: list[int] + additional_information: Any = None + sampling_params: Any = None + num_computed_tokens: int = 0 + resumable: bool = False + output_token_ids: list[int] = field(default_factory=list) + finished: bool = False + + def snapshot(self, *, include_token_history: bool) -> SimpleNamespace: + prompt = list(self.prompt_token_ids) if include_token_history else [] + output = list(self.output_token_ids) if include_token_history else [] + finished = self.finished + request = SimpleNamespace( + request_id=self.request_id, + req_id=self.request_id, + external_req_id=self.external_req_id, + prompt_token_ids=prompt, + output_token_ids=output, + all_token_ids=prompt + output, + output_token_count=len(self.output_token_ids), + additional_information=self.additional_information, + sampling_params=self.sampling_params, + num_computed_tokens=self.num_computed_tokens, + resumable=self.resumable, + ) + request.is_finished = lambda: finished + return request + + +@dataclass(frozen=True) +class _NativeOutputBatch: + req_ids: list[str] + inter_stage_outputs: list[Any | None] | None + sampled_token_ids: list[list[int]] | None + + +_NATIVE_OUTPUT_STOP = object() +_NATIVE_OUTPUT_QUEUE_DEPTH = 8 +_DEFAULT_DELIVERY_TIMEOUT_S = 30.0 +_DEFAULT_SHUTDOWN_TIMEOUT_S = 5.0 +logger = init_logger(__name__) + + +class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): + """MRv2-owned stage transport and request-side payload state.""" + + def __init__(self, vllm_config: Any, model_config: Any) -> None: + self._native_requests: dict[str, _NativeRequestState] = {} + self._native_output_lock = threading.RLock() + self._native_send_lock = threading.Lock() + self._native_outputs_in_flight: dict[str, int] = defaultdict(int) + self._native_terminal_pending: set[str] = set() + self.init_omni_connectors(model_config=model_config) + self._delivery_manager = OmniDeliveryManager( + delivery_timeout_s=self._connector_delivery_timeout(), + shutdown_timeout_s=_DEFAULT_SHUTDOWN_TIMEOUT_S, + ) + self._can_send = self._custom_process_func is not None + self._start_output_worker(max_pending_batches=_NATIVE_OUTPUT_QUEUE_DEPTH) + + def _connector_delivery_timeout(self) -> float: + config = getattr(getattr(self, "_omni_connector", None), "config", None) + extra = config.get("extra", {}) if isinstance(config, dict) else {} + value = extra.get("delivery_timeout_s", _DEFAULT_DELIVERY_TIMEOUT_S) + try: + timeout = float(value) + except (TypeError, ValueError) as error: + raise ValueError(f"invalid connector delivery_timeout_s: {value!r}") from error + if timeout <= 0: + raise ValueError("connector delivery_timeout_s must be positive") + return timeout + + def _create_send_completion(self, *, request_id: str, put_key: str) -> Any: + return self._delivery_manager.create_ticket( + request_id=request_id, + put_key=put_key, + ) + + def _start_output_worker(self, *, max_pending_batches: int) -> None: + if max_pending_batches <= 0: + raise ValueError("max_pending_batches must be positive") + self._native_output_queue: Queue[_NativeOutputBatch | object] = Queue(maxsize=max_pending_batches) + self._native_output_error_lock = threading.Lock() + self._native_output_error: BaseException | None = None + self._connector_send_error_sink = self._record_output_error + self._native_output_closed = False + stage_id = getattr(self, "_stage_id", "unknown") + self._native_output_worker: threading.Thread | None = threading.Thread( + target=self._output_worker_loop, + name=f"omni-native-output-{stage_id}", + daemon=True, + ) + self._native_output_worker.start() + + def _record_output_error(self, error: BaseException) -> None: + first_error = False + with self._native_output_error_lock: + if self._native_output_error is None: + self._native_output_error = error + first_error = True + if first_error: + manager = getattr(self, "_delivery_manager", None) + if manager is not None: + manager.quarantine(error) + + def _raise_output_error(self) -> None: + with self._native_output_error_lock: + error = self._native_output_error + if error is not None: + raise RuntimeError(f"native output worker failed: {error}") from error + + def _output_worker_loop(self) -> None: + while True: + batch = self._native_output_queue.get() + try: + if batch is _NATIVE_OUTPUT_STOP: + return + if self._native_output_error is not None: + continue + assert isinstance(batch, _NativeOutputBatch) + self.complete_outputs( + req_ids=batch.req_ids, + inter_stage_outputs=batch.inter_stage_outputs, + sampled_token_ids=batch.sampled_token_ids, + ) + except BaseException as error: + self._record_output_error(error) + finally: + self._native_output_queue.task_done() + + def enqueue_outputs( + self, + *, + req_ids: list[str], + inter_stage_outputs: list[Any | None] | None, + sampled_token_ids: list[list[int]] | None, + ) -> None: + """Transfer one completed model batch to the ordered output worker.""" + self._raise_output_error() + if self._native_output_closed: + raise RuntimeError("native output worker is closed") + batch = _NativeOutputBatch( + req_ids=req_ids, + inter_stage_outputs=inter_stage_outputs, + sampled_token_ids=sampled_token_ids, + ) + while True: + try: + self._native_output_queue.put(batch, timeout=0.1) + break + except Full: + self._raise_output_error() + self._raise_output_error() + + def drain_outputs(self) -> None: + queue = self._native_output_queue + pending_batches = max(1, queue.qsize() + 1) + timeout_s = self._delivery_manager.delivery_timeout_s * pending_batches + deadline = time.monotonic() + timeout_s + while queue.unfinished_tasks: + self._raise_output_error() + remaining = deadline - time.monotonic() + if remaining <= 0: + error = DeliveryTimeoutError( + "native output queue drain timed out after " + f"{timeout_s:.3f}s with {queue.unfinished_tasks} " + "unfinished batch(es)" + ) + self._record_output_error(error) + self._raise_output_error() + with queue.all_tasks_done: + if queue.unfinished_tasks: + queue.all_tasks_done.wait(timeout=min(0.05, remaining)) + self._raise_output_error() + + def _stop_output_worker(self) -> None: + worker = getattr(self, "_native_output_worker", None) + if worker is None: + return + self._native_output_closed = True + self._delivery_manager.shutdown(RuntimeError("MRv2 output worker shutdown")) + deadline = time.monotonic() + self._delivery_manager.shutdown_timeout_s + while worker.is_alive(): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + self._native_output_queue.put( + _NATIVE_OUTPUT_STOP, + timeout=min(0.05, remaining), + ) + break + except Full: + continue + worker.join(timeout=max(0.0, deadline - time.monotonic())) + if worker.is_alive(): + raise RuntimeError("MRv2 output worker did not stop within its shutdown deadline") + self._native_output_worker = None + + def register_request(self, request_data: Any) -> None: + self._raise_output_error() + req_id = str(request_data.req_id) + external_req_id = str(getattr(request_data, "external_req_id", None) or req_id) + with self._native_output_lock: + self._native_outputs_in_flight.pop(req_id, None) + self._native_terminal_pending.discard(req_id) + self._native_requests[req_id] = _NativeRequestState( + request_id=req_id, + external_req_id=external_req_id, + prompt_token_ids=list(getattr(request_data, "prompt_token_ids", None) or []), + additional_information=getattr(request_data, "additional_information", None), + sampling_params=getattr(request_data, "sampling_params", None), + num_computed_tokens=int(getattr(request_data, "num_computed_tokens", 0) or 0), + resumable=bool(getattr(request_data, "resumable", False)), + ) + + def reserve_outputs(self, req_ids: list[str]) -> None: + """Keep request state alive until deferred runner outputs are consumed.""" + self._raise_output_error() + with self._native_output_lock: + for req_id in dict.fromkeys(req_ids): + if req_id in self._native_requests: + self._native_outputs_in_flight[req_id] += 1 + + def _ready_terminal_requests(self) -> set[str]: + return { + req_id for req_id in self._native_terminal_pending if self._native_outputs_in_flight.get(req_id, 0) == 0 + } + + def _emit_ready_terminals(self) -> int: + ready = self._ready_terminal_requests() + if not ready: + return 0 + emitted = self.emit_chunks( + req_ids=[], + inter_stage_outputs=None, + sampled_token_ids=None, + terminal_req_ids=ready, + ) + self._native_terminal_pending.difference_update(ready) + for req_id in ready: + self._native_outputs_in_flight.pop(req_id, None) + return emitted + + def request_terminal(self, req_ids: set[str]) -> int: + """Emit terminal markers only after all deferred outputs are enqueued.""" + self._raise_output_error() + with self._native_output_lock: + self._native_terminal_pending.update(req_id for req_id in req_ids if req_id in self._native_requests) + return self._emit_ready_terminals() + + def abort_requests(self, req_ids: set[str]) -> int: + """Cancel deferred outputs and terminate each live request once.""" + with self._native_output_lock: + active_req_ids = {req_id for req_id in req_ids if req_id in self._native_requests} + # Receivers are registered before requests acquire model slots. + # Cancelling during that interval must still stop recv polling. + for req_id in req_ids - active_req_ids: + self.cleanup_finished_request(req_id) + if not active_req_ids: + return 0 + for req_id in active_req_ids: + self._native_outputs_in_flight.pop(req_id, None) + self._native_terminal_pending.discard(req_id) + return self.emit_chunks( + req_ids=[], + inter_stage_outputs=None, + sampled_token_ids=None, + terminal_req_ids=active_req_ids, + ) + + def complete_outputs( + self, + *, + req_ids: list[str], + inter_stage_outputs: list[Any | None] | None, + sampled_token_ids: list[list[int]] | None, + ) -> int: + """Commit one deferred output batch, then release its lifecycle holds.""" + send_lock_held = False + with self._native_output_lock: + entries, finished_req_ids = self._build_chunk_entries( + req_ids=req_ids, + inter_stage_outputs=inter_stage_outputs, + sampled_token_ids=sampled_token_ids, + terminal_req_ids=set(), + ) + # Claim connector ordering before publishing the updated request + # state. request_terminal() can still register a pending terminal + # while the enqueue runs, but abort cannot overtake this output. + if entries: + self._native_send_lock.acquire() + send_lock_held = True + try: + emitted = self._send_chunk_entries( + entries, + send_lock_held=send_lock_held, + ) + finally: + if send_lock_held: + self._native_send_lock.release() + with self._native_output_lock: + self._cleanup_finished_entries(finished_req_ids) + for req_id in dict.fromkeys(req_ids): + in_flight = self._native_outputs_in_flight.get(req_id, 0) + if in_flight > 1: + self._native_outputs_in_flight[req_id] = in_flight - 1 + elif in_flight == 1: + self._native_outputs_in_flight.pop(req_id, None) + return emitted + self._emit_ready_terminals() + + def register_receivers(self, handles: list[Any]) -> None: + for handle in handles: + self.register_chunk_recv(handle) + + @classmethod + def _copy_payload_structure(cls, value: Any) -> Any: + if isinstance(value, dict): + return {key: cls._copy_payload_structure(item) for key, item in value.items()} + if isinstance(value, list): + return list(value) + return value + + def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload: + """Accumulate chunk payloads (concat tensors, extend lists).""" + if req_id not in self._send_side_request_payload: + self._send_side_request_payload[req_id] = dict(payload_data) + return cast(OmniPayload, dict(self._send_side_request_payload[req_id])) + + origin = self._send_side_request_payload[req_id] + merged = dict(origin) + raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else [] + override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok} + + for key, value in payload_data.items(): + if isinstance(value, dict): + origin_sub = origin.get(key) + merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {} + span_handled: set[str] = set() + if key == "embed" and isinstance(origin_sub, dict): + for tk, sk, ek in (("decode", "decode_token_start", "decode_token_end"),): + if tk not in value or (key, tk) in override_keys: + continue + span = merge_tensor_spans( + get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek), + get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek), + ) + if span is None: + continue + t, s, e = span + merged_sub[tk] = t + merged_sub[sk] = s + merged_sub[ek] = e + span_handled |= {tk, sk, ek} + for qual, qval in value.items(): + if qual in span_handled: + continue + if key == "meta" and qual == "finished": + merged_sub[qual] = qval + continue + if (key, qual) in override_keys: + merged_sub[qual] = qval + continue + osv = merged_sub.get(qual) + if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor): + merged_sub[qual] = torch.cat([osv, qval], dim=0) + elif isinstance(qval, list) and isinstance(osv, list): + merged_sub[qual] = osv + qval + else: + merged_sub[qual] = qval + merged[key] = merged_sub + else: + if key in override_keys: + merged[key] = value + continue + ov = origin.get(key) + if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor): + merged[key] = torch.cat([ov, value], dim=0) + elif isinstance(value, list) and isinstance(ov, list): + merged[key] = ov + value + else: + merged[key] = value + + self._send_side_request_payload[req_id] = merged + return cast(OmniPayload, dict(merged)) + + def pop_local_stage_payload(self, req_id: str) -> Any: + """Hand one accumulated delta to the model and acknowledge its rows. + + Connector accumulation and ``model_intermediate_buffer`` are separate + ownership domains. Once decode rows are handed to the model, remove + them from connector accumulation so the next chunk cannot replay the + same rows into Talker's ``cached_decode``. + """ + with self._lock: + payload = self._local_stage_payload_cache.pop(req_id, None) + if payload is None: + return None + delivered = self._copy_payload_structure(payload) + external_req_id = self._request_ids_mapping.get(req_id, req_id) + accumulated = self._send_side_request_payload.get(external_req_id) + if isinstance(accumulated, dict): + embed = accumulated.get("embed") + if isinstance(embed, dict): + embed.pop("decode", None) + embed.pop("decode_token_start", None) + embed.pop("decode_token_end", None) + hidden_states = accumulated.get("hidden_states") + if isinstance(hidden_states, dict): + hidden_states.pop("output", None) + ids = accumulated.get("ids") + if isinstance(ids, dict): + ids.pop("output", None) + return delivered + + def emit_chunks( + self, + *, + req_ids: list[str], + inter_stage_outputs: list[Any | None] | None, + sampled_token_ids: list[list[int]] | None, + terminal_req_ids: set[str], + ) -> int: + with self._native_output_lock: + entries, finished_req_ids = self._build_chunk_entries( + req_ids=req_ids, + inter_stage_outputs=inter_stage_outputs, + sampled_token_ids=sampled_token_ids, + terminal_req_ids=terminal_req_ids, + ) + emitted = self._send_chunk_entries(entries) + self._cleanup_finished_entries(finished_req_ids) + return emitted + + def _build_chunk_entries( + self, + *, + req_ids: list[str], + inter_stage_outputs: list[Any | None] | None, + sampled_token_ids: list[list[int]] | None, + terminal_req_ids: set[str], + ) -> tuple[list[tuple[Any, Any | None]], list[str]]: + if not getattr(self, "_can_send", True): + receive_finished_ids = [req_id for req_id in sorted(terminal_req_ids) if req_id in self._native_requests] + return [], receive_finished_ids + payload_by_req = { + req_id: inter_stage_outputs[index] + for index, req_id in enumerate(req_ids) + if inter_stage_outputs is not None and index < len(inter_stage_outputs) + } + sampled_by_req = { + req_id: sampled_token_ids[index] + for index, req_id in enumerate(req_ids) + if sampled_token_ids is not None and index < len(sampled_token_ids) + } + entries: list[tuple[Any, Any | None]] = [] + finished_req_ids: list[str] = [] + for req_id in dict.fromkeys([*req_ids, *sorted(terminal_req_ids)]): + state = self._native_requests.get(req_id) + if state is None: + continue + state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, [])) + state.finished = req_id in terminal_req_ids + payload = payload_by_req.get(req_id) + if isinstance(payload, dict): + payload = unflatten_payload(payload) + if payload is None and not state.finished: + continue + include_token_history = bool(state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0) + entries.append( + ( + state.snapshot(include_token_history=include_token_history), + payload, + ) + ) + if state.finished: + finished_req_ids.append(req_id) + return entries, finished_req_ids + + def _send_chunk_entries( + self, + entries: list[tuple[Any, Any | None]], + *, + send_lock_held: bool = False, + ) -> int: + if not entries: + return 0 + context: AbstractContextManager + if send_lock_held: + context = nullcontext() + else: + context = self._native_send_lock + with context: + return self.send_chunks(entries, propagate_errors=True, wait_for_delivery=True) + + def _cleanup_finished_entries(self, finished_req_ids: list[str]) -> None: + for req_id in finished_req_ids: + self._native_requests.pop(req_id, None) + self.cleanup_finished_request(req_id) + + def shutdown_omni_connectors(self) -> None: + """Bound MRv2 shutdown even if a connector call cannot be cancelled. + + Connector ``close()`` runs concurrently with the I/O threads so a + backend that supports cancellation can release a blocked ``put()``. + Backends that do not support it are quarantined and left only on + daemon threads after the bounded deadline; V1 keeps its existing + shutdown implementation in the shared mixin. + """ + timeout_s = self._delivery_manager.shutdown_timeout_s + deadline = time.monotonic() + timeout_s + self._delivery_manager.shutdown(RuntimeError("MRv2 connector shutdown")) + self._stop_event.set() + self._work_available.set() + + close_errors: list[BaseException] = [] + connector = getattr(self, "_omni_connector", None) + + def close_connector() -> None: + if connector is None: + return + try: + connector.close() + except BaseException as error: + close_errors.append(error) + + close_thread = threading.Thread( + target=close_connector, + name=f"omni-native-connector-close-{getattr(self, '_stage_id', 'unknown')}", + daemon=True, + ) + close_thread.start() + + live_threads: list[str] = [] + for name, thread in ( + ("recv", getattr(self, "_recv_thread", None)), + ("save", getattr(self, "_save_thread", None)), + ("connector-close", close_thread), + ): + if thread is None: + continue + thread.join(timeout=max(0.0, deadline - time.monotonic())) + if thread.is_alive(): + live_threads.append(name) + + if close_errors: + raise RuntimeError("MRv2 connector close failed") from close_errors[0] + if live_threads: + raise RuntimeError( + f"MRv2 connector shutdown exceeded {timeout_s:.3f}s; blocked thread(s): {', '.join(live_threads)}" + ) + + def close(self) -> None: + error: BaseException | None = None + try: + self.drain_outputs() + except BaseException as exc: + error = exc + finally: + self._delivery_manager.shutdown(error or RuntimeError("MRv2 data plane close")) + try: + self._stop_output_worker() + except BaseException as exc: + if error is None: + error = exc + try: + self.shutdown_omni_connectors() + except BaseException as exc: + if error is None: + error = exc + if error is not None: + raise error diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py new file mode 100644 index 000000000..7c024f675 --- /dev/null +++ b/vllm_omni/worker_v2/omni_generation_model_runner.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniGenerationModelRunner — non-autoregressive stage runner on MR V2. + +Used for stages like Code2Wav that convert codec codes to audio waveforms. +No token sampling or logits computation — model output goes directly into +``pooler_output``. Inherits from ``OmniGPUModelRunner`` for intermediate +buffer and lifecycle hooks. +""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + get_ep_all2all_manager, +) +from vllm.utils.torch_utils import PIN_MEMORY +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput +from vllm.v1.worker.gpu.model_runner import ( + BatchDescriptor, + ExecuteModelState, + IntermediateTensors, +) + +from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData +from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker_v2.omni_ar_model_runner import ( + _async_copy_mm, + _ensure_tensor_values, +) +from vllm_omni.worker_v2.omni_model_runner import ( + OmniGPUModelRunner, +) + +logger = init_logger(__name__) + + +def _contains_cuda_tensor(value: Any) -> bool: + """Return whether a nested output payload still owns CUDA storage.""" + if isinstance(value, torch.Tensor): + return value.is_cuda + if isinstance(value, dict): + return any(_contains_cuda_tensor(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(_contains_cuda_tensor(v) for v in value) + return False + + +def _materialize_generation_value( + value: Any, + req_index: int, + num_reqs: int, + *, + top_level: bool = True, + copy_cpu: bool = True, +) -> Any: + """Recursively select and own one request's generation payload.""" + if isinstance(value, torch.Tensor): + if value.ndim > 0 and value.shape[0] == num_reqs: + value = value[req_index] + return value.detach().to(device="cpu", copy=copy_cpu and value.device.type == "cpu").contiguous() + if isinstance(value, dict): + return { + key: _materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu) + for key, item in value.items() + } + if isinstance(value, list): + # Tensor lists (e.g. variable-length waveforms) carry one item per + # request even inside a dict. Scalar lists remain payload metadata. + per_request = top_level or all(item is None or isinstance(item, torch.Tensor) for item in value) + if per_request and len(value) == num_reqs: + # The selected item no longer has a batch axis. + return _materialize_generation_value(value[req_index], 0, -1, top_level=False, copy_cpu=copy_cpu) + return [ + _materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu) + for item in value + ] + if isinstance(value, tuple): + return tuple( + _materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu) + for item in value + ) + return value + + +class OmniGenerationAsyncOutput(AsyncModelRunnerOutput): + """Overlap generation-stage D2H copies with the next model step.""" + + def __init__( + self, + *, + model_runner_output: OmniModelRunnerOutput, + multimodal_outputs: dict[str, Any] | None, + num_reqs: int, + main_stream: torch.cuda.Stream, + copy_stream: torch.cuda.Stream, + finalize_output: Any | None = None, + check_ep_fault: bool = False, + ) -> None: + self.model_runner_output = model_runner_output + self.num_reqs = num_reqs + self.finalize_output = finalize_output + self.copy_event = torch.cuda.Event(blocking=True) + self._has_fault: torch.Tensor | None = None + + with torch.cuda.stream(copy_stream): + copy_stream.wait_stream(main_stream) + self.multimodal_outputs_cpu = _async_copy_mm( + multimodal_outputs, + total_tokens=0, + copy_stream=copy_stream, + pin_memory=PIN_MEMORY, + ) + if check_ep_fault: + has_fault = get_ep_all2all_manager().query_fault() + self._has_fault = has_fault.to("cpu", non_blocking=True) + self.copy_event.record(copy_stream) + + def get_output(self) -> OmniModelRunnerOutput: + self.copy_event.synchronize() + if self._has_fault is not None and self._has_fault.item(): + mask = get_ep_all2all_manager().query_active_mask() + raise RuntimeError( + "Fault detected in EP all2all communication: one or more ranks " + f"timed out during dispatch/combine. Mask: {mask.cpu().tolist()}" + ) + payloads = OmniGenerationModelRunner._build_pooler_output_from_cpu( + self.multimodal_outputs_cpu, + self.num_reqs, + ) + self.model_runner_output.multimodal_outputs = [ + _ensure_tensor_values(payload) if payload else {} for payload in payloads + ] + if self.finalize_output is not None: + return self.finalize_output(self.model_runner_output) + return self.model_runner_output + + +class OmniGenerationModelRunner(OmniGPUModelRunner): + """Non-autoregressive generation runner (e.g. Code2Wav). + + Overrides ``execute_model`` to skip the tensor-only assertion and + ``sample_tokens`` to construct ``pooler_output`` from multimodal + model outputs without performing token sampling. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._gen_model_output: Any = None + self._gen_input_batch: Any = None + # Placeholder for ExecuteModelState.hidden_states — allocated + # once and reused every step to avoid per-forward allocation. + self._dummy_hidden = torch.zeros(1, dtype=self.dtype, device=self.device) + + # ------------------------------------------------------------------ + # Async-chunk support: replace prompt_token_ids for cached requests + # ------------------------------------------------------------------ + + def _handle_async_chunk_updates(self, scheduler_output: SchedulerOutput) -> None: + """In-place update cached requests whose prompt_token_ids changed. + + In async_chunk mode, the ``ChunkTransferAdapter`` replaces + ``Request.prompt_token_ids`` with new codec frames for each + chunk and resets ``num_computed_tokens`` to 0. The scheduler + propagates the new ``prompt_token_ids`` via + ``OmniCachedRequestData``. + + Instead of remove + re-add (which involves free_indices churn + and redundant model_state init), we update the existing slot + in-place. This is safe for Code2Wav because: + - No KV cache / rope state to reinitialize + - staged writes are applied once at the end + + The old intermediate buffer for this slot is cleared here; the + inherited ``OmniGPUModelRunner.update_requests`` (called right after + this method in ``execute_model``) writes the current chunk state. + """ + cached = scheduler_output.scheduled_cached_reqs + if not cached.req_ids: + return + + if not isinstance(cached, OmniCachedRequestData): + return + + new_prompt_ids = cached.prompt_token_ids + if not new_prompt_ids: + return + + updated = False + released_chunks: list[OmniNewRequestData] = [] + + for i, req_id in enumerate(cached.req_ids): + new_ids = new_prompt_ids.get(req_id) + if new_ids is None: + continue + + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + block_ids = cached.new_block_ids[i] + released_chunks.append( + OmniNewRequestData( + req_id=req_id, + prompt_token_ids=new_ids, + mm_features=[], + sampling_params=None, + pooling_params=None, + block_ids=block_ids if block_ids is not None else tuple(), + num_computed_tokens=0, + lora_request=None, + prompt_embeds=None, + prefill_token_ids=new_ids, + additional_information=cached.additional_information.get(req_id), + ) + ) + continue + + self.model_state.remove_request(req_idx) + + # In-place update token state — same slot, no remove/re-add. + # .np[] = direct write (no GPU buffer); stage_write = GPU-synced. + n = len(new_ids) + self.req_states.prompt_len.np[req_idx] = n + self.req_states.prefill_len.np[req_idx] = n + self.req_states.total_len.stage_write_elem(req_idx, n) + self.req_states.all_token_ids.stage_write(req_idx, 0, new_ids) + self.req_states.num_computed_tokens.stage_write_elem(req_idx, 0) + self.req_states.num_computed_prefill_tokens[req_idx] = 0 + + updated = True + + if released_chunks: + self.add_requests( + replace( + scheduler_output, + scheduled_new_reqs=released_chunks, + ) + ) + if updated: + self.req_states.apply_staged_writes() + + def _release_generation_slots(self, input_batch: Any) -> None: + if not getattr(self.model_config, "async_chunk", False): + return + model_state = getattr(self, "model_state", None) + remove_request = getattr(self, "_remove_request", None) + if model_state is None or remove_request is None: + return + for i in range(input_batch.num_reqs): + req_id = input_batch.req_ids[i] + req_idx = int(input_batch.idx_mapping_np[i]) + model_state.remove_request(req_idx) + remove_request(req_id) + + # ------------------------------------------------------------------ + # profile / warmup — skip sampler since there are no logits + # ------------------------------------------------------------------ + + @torch.inference_mode() + def profile_run(self) -> None: + """Generation models have no KV cache — skip profiling. + + Code2Wav shares GPU memory with the Talker stage (same device); + its memory footprint is managed via ``gpu_memory_utilization`` + config, not profiled dynamically. Running the real model with + random input_ids causes out-of-bounds indexing in codec lookup + tables. + """ + torch.accelerator.synchronize() + + # ------------------------------------------------------------------ + # execute_model — run the generation model, store raw output + # ------------------------------------------------------------------ + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + context_len: int = 0, + ) -> ModelRunnerOutput | IntermediateTensors | None: + if not dummy_run: + self._prepare_native_data_plane(scheduler_output) + self.finish_requests(scheduler_output) + self.free_states(scheduler_output) + if ( + scheduler_output.total_num_scheduled_tokens == 0 + and not scheduler_output.scheduled_new_reqs + and not scheduler_output.scheduled_cached_reqs.req_ids + ): + # Registration/terminal/abort work has been applied above. + # There is no new or cached request state to add or update. + self._apply_block_table_staged_writes_if_available() + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._attach_native_data_plane_signals( + self._merge_ec_connector_no_forward(scheduler_output, empty_output) + ) + # Handle async_chunk prompt_token_ids replacement for cached + # requests BEFORE add/update — update the existing slot + # in-place with the new chunk's tokens. + self._handle_async_chunk_updates(scheduler_output) + self.add_requests(scheduler_output) + self.update_requests(scheduler_output) + self._sync_native_data_plane_payloads(scheduler_output) + self._apply_block_table_staged_writes_if_available() + if scheduler_output.total_num_scheduled_tokens == 0: + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._attach_native_data_plane_signals( + self._merge_ec_connector_no_forward(scheduler_output, empty_output) + ) + + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + batch_req_state, uniform_tok_count = self.gather_batch_req_state(scheduler_output, dummy_run) + if batch_req_state is not None: + num_toks = batch_req_state.num_tokens + batch_desc, dp_sync = self._dispatch_batch_descriptor( + num_reqs=num_reqs, + num_toks=num_toks, + uniform_tok_count=uniform_tok_count, + num_active_loras=0, + # Generation stages return structured OmniOutput payloads and do + # not expose the tensor-only outer graph contract. Their model may + # still use an internal decoder graph, but the vLLM outer runner + # always uses eager dispatch here. + use_eager=True, + max_query_len=max_query_len, + ) + + if batch_desc.num_tokens == 0: + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._attach_native_data_plane_signals( + self._merge_ec_connector_no_forward(scheduler_output, empty_output) + ) + + if not dummy_run: + assert batch_req_state is not None + input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc) + else: + from vllm.v1.worker.gpu.input_batch import InputBatch + + input_batch = InputBatch.make_dummy( + batch_desc.num_reqs or num_reqs, + batch_desc.num_tokens, + self.input_buffers, + max_query_len=batch_desc.max_query_len, + ) + + attn_metadata = None + slot_mappings_by_layer = None + + input_ids, inputs_embeds, ec_connector_output = self._prepare_mm_inputs( + scheduler_output, + input_batch, + dummy_run=dummy_run, + ) + + model_inputs = { + "input_ids": input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": intermediate_tensors, + **self.model_state.prepare_inputs(input_batch, self.req_states), + } + self._add_legacy_forward_inputs(model_inputs, input_batch) + + eplb = getattr(self, "eplb", None) + if eplb is not None: + eplb.prepare_forward(self.model_config, input_batch.num_tokens) + + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=self.lora_config is not None, + num_active_loras=batch_desc.num_active_loras, + ) + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=batch_desc.cg_mode, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + is_padding=input_batch.is_padding, + ): + self.kv_connector.pre_forward(scheduler_output) + if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: + assert self.cudagraph_manager is not None + model_output = self.cudagraph_manager.run_pw_graph( + self.model, + model_inputs, + ) + else: + model_output = self.model(**model_inputs) + + # Convert raw model output to OmniOutput. + if not isinstance(model_output, OmniOutput): + buffer_list = self.model_state.intermediate_buffer.gather(input_batch) + make_output_kwargs = {"model_intermediate_buffer": buffer_list} + make_output_kwargs["runtime_additional_information"] = buffer_list + model_output = self.model.make_omni_output(model_output, **make_output_kwargs) + + self._gen_model_output = model_output + self._gen_input_batch = input_batch + + # ExecuteModelState is required by the upstream engine loop + # (EngineCore checks execute_model_state is not None before + # calling sample_tokens). + self.execute_model_state = ExecuteModelState( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=self._dummy_hidden, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + dp_sync=dp_sync, + ec_connector_output=ec_connector_output, + routed_experts=None, + ) + return None + + # ------------------------------------------------------------------ + # sample_tokens — build pooler_output, no actual sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() + def sample_tokens( + self, + grammar_output: GrammarOutput | None = None, + ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | None: + model_output = self._gen_model_output + input_batch = self._gen_input_batch + execute_model_state = self.execute_model_state + self._gen_model_output = None + self._gen_input_batch = None + self.execute_model_state = None + + if model_output is None or input_batch is None or execute_model_state is None: + return None + + num_reqs = input_batch.num_reqs + + # Mark all scheduled tokens as computed so the scheduler does + # not re-schedule them. Unlike AR stages we do NOT call + # self.postprocess() — that kernel advances num_computed_tokens + # by 1 and emits sampled tokens, which would cause check_stop + # to fire. Instead, set num_computed_tokens = prompt_len + # directly, matching V1's behavior. + for i in range(num_reqs): + req_idx = int(input_batch.idx_mapping_np[i]) + prompt_len = int(self.req_states.prompt_len.np[req_idx]) + self.req_states.num_computed_tokens.stage_write_elem(req_idx, prompt_len) + self.req_states.num_computed_tokens.apply_write() + kv_connector_output = self.kv_connector.post_forward(execute_model_state.finished_req_ids) + + # Async finalization can run after the input batch is reused for the + # next scheduler step, so output metadata must not alias the batch. + req_ids = list(input_batch.req_ids) + + # Generation models don't do token sampling. Return one empty + # list per request so the scheduler does NOT trigger check_stop + # (which would prematurely finish the request). The request + # stays RUNNING until the orchestrator marks it done via + # chunk_transfer_adapter.finished_requests. + sampled_token_ids: list[list[int]] = [[] for _ in range(len(req_ids))] + output = OmniModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={rid: i for i, rid in enumerate(req_ids)}, + sampled_token_ids=sampled_token_ids, + # Match the V1 generation runner contract: final generation stages + # publish audio only through multimodal_outputs. If the same + # payload is also exposed as pooling_output, OutputProcessor routes + # it through the pooling path and can duplicate accumulated audio. + pooler_output=None, + multimodal_outputs=None, + kv_connector_output=kv_connector_output, + ec_connector_output=execute_model_state.ec_connector_output, + sampled_token_ids_materialized=True, + ) + + raw_multimodal_outputs = model_output.multimodal_outputs + can_copy_async = getattr(getattr(self, "device", None), "type", "cpu") == "cuda" and ( + self.check_ep_fault or _contains_cuda_tensor(raw_multimodal_outputs) + ) + if can_copy_async: + async_output = OmniGenerationAsyncOutput( + model_runner_output=output, + multimodal_outputs=raw_multimodal_outputs, + num_reqs=num_reqs, + main_stream=self.main_stream, + copy_stream=self.output_copy_stream, + finalize_output=self._finalize_native_data_plane_output, + check_ep_fault=self.check_ep_fault, + ) + self._reserve_native_data_plane_outputs(list(req_ids)) + self._release_generation_slots(input_batch) + return async_output + + # CPU execution keeps the synchronous materialization path. Final + # generation stages publish audio only through multimodal_outputs, + # matching the V1 generation-runner contract. + self._reserve_native_data_plane_outputs(list(req_ids)) + multimodal_outputs = self._build_pooler_output(model_output, num_reqs) + output.multimodal_outputs = [ + _ensure_tensor_values(payload) if payload else {} for payload in multimodal_outputs + ] + self._release_generation_slots(input_batch) + return self._finalize_native_data_plane_output(output) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _apply_block_table_staged_writes_if_available(self) -> None: + """Flush block-table writes only when this stage has a writer. + + Generation stages such as Code2Wav do not use KV cache or attention + metadata. Newer vLLM BlockTable implementations assert that a fused + writer exists before applying staged writes; no-KV generation stages + legitimately do not have one. + """ + block_tables = getattr(self, "block_tables", None) + if block_tables is None: + return + if getattr(block_tables, "fused_writer", None) is None: + return + block_tables.apply_staged_writes() + + @staticmethod + def _build_pooler_output( + model_output: OmniOutput, + num_reqs: int, + ) -> list[dict[str, Any] | None]: + """Extract per-request pooler payloads from model output. + + Code2Wav's ``make_omni_output`` returns + ``{"model_outputs": [tensor_per_req, ...], "sr": [...]}``, + so each value is a ``list`` with ``len == num_reqs``. + """ + mm = model_output.multimodal_outputs + if not isinstance(mm, dict): + logger.warning( + "Unexpected multimodal_outputs type: %s; returning empty pooler_output", + type(mm).__name__ if mm is not None else "None", + ) + return [None] * num_reqs + + pooler: list[dict[str, Any] | None] = [] + for i in range(num_reqs): + payload = {key: _materialize_generation_value(value, i, num_reqs) for key, value in mm.items()} + pooler.append(flatten_payload(payload) if isinstance(payload, dict) else payload) + return pooler + + @staticmethod + def _build_pooler_output_from_cpu( + multimodal_outputs: dict[str, Any] | None, + num_reqs: int, + ) -> list[dict[str, Any] | None]: + """Build per-request payloads after an async whole-batch D2H copy.""" + if not isinstance(multimodal_outputs, dict): + return [None] * num_reqs + + pooler: list[dict[str, Any] | None] = [] + for i in range(num_reqs): + payload = { + key: _materialize_generation_value(value, i, num_reqs, copy_cpu=False) + for key, value in multimodal_outputs.items() + } + pooler.append(flatten_payload(payload) if isinstance(payload, dict) else payload) + return pooler diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py new file mode 100644 index 000000000..20337fc91 --- /dev/null +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -0,0 +1,666 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Omni v2 GPU model runner hooks.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any, cast + +import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import ( + set_forward_context, +) +from vllm.logger import init_logger +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.input_batch import set_dummy_context +from vllm.v1.worker.gpu.lora_utils import get_num_active_loras_for_dispatch +from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras +from vllm.v1.worker.gpu.model_runner import ( + BatchDescriptor, + BatchExecutionDescriptor, + ExecuteModelState, + GPUModelRunner, + IntermediateTensors, + build_slot_mappings_by_layer, +) +from vllm.v1.worker.gpu.sample.sampler import Sampler + +from vllm_omni.core.sched.omni_scheduling_coordinator import ( + uses_native_mrv2_data_plane, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker.sampling_utils import sanitize_sampling_params_min_tokens_stop_ids +from vllm_omni.worker_v2.model_states import init_omni_model_state +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + _resolve_additional_information, +) +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.omni_data_plane import OmniRunnerDataPlane + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + + from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker + +logger = init_logger(__name__) + +_model_state_patch_lock = threading.RLock() + + +def _needs_capture_tensor_unwrap(model: Any) -> bool: + return bool(getattr(model, "_returns_tuple", False)) + + +class OmniGPUModelRunner(GPUModelRunner): + """Thin layer over v2 ``GPUModelRunner`` for Omni lifecycle hooks.""" + + supports_mm_inputs: bool + encoder_cache: EncoderCache | None + _native_output_materializer: NativeOutputWorker | None + model_state: OmniModelState + _last_aux_output: Any + _last_multimodal_outputs: dict[str, Any] | None + _model_returns_tuple: bool + + def _configure_cudagraph_output_contract(self) -> None: + """Select the CUDA graph output contract declared by the model.""" + self._model_returns_tuple = _needs_capture_tensor_unwrap(self.model) + self._exclude_full_graph = self._model_returns_tuple or hasattr(self.model, "_last_captured_layers") + + @staticmethod + def _prepare_cudagraph_capture_output(model_output: Any) -> Any: + """Unwrap model outputs for the upstream PIECEWISE capture warmup.""" + if isinstance(model_output, OmniOutput): + return model_output.text_hidden_states + if isinstance(model_output, tuple) and len(model_output) == 2: + return model_output[0] + return model_output + + def _add_legacy_forward_inputs(self, model_inputs: dict[str, Any], input_batch: Any) -> None: + """Supply forward-only metadata for pipelines not on the native plane.""" + if self._omni_data_plane is not None: + return + model_inputs.setdefault("sampling_metadata", getattr(input_batch, "sampling_metadata", None)) + logits_index = getattr( + input_batch, + "logits_indices", + getattr(input_batch, "logits_index", None), + ) + model_inputs.setdefault("logits_index", logits_index) + model_inputs.setdefault("sampler", self.sampler) + + def _get_mm_embeddings( + self, + scheduled_encoder_inputs: dict[str, list[int]], + input_batch: Any, + ) -> torch.Tensor: + """Call the native vLLM multimodal embedding contract.""" + return self.model_state.prepare_inputs_embeds( + scheduled_encoder_inputs, + input_batch, + self.req_states, + ) + + def _prepare_mm_inputs( + self, + scheduler_output: SchedulerOutput, + input_batch: Any, + *, + dummy_run: bool, + ) -> tuple[Any, torch.Tensor | None, Any]: + """Prepare native vLLM multimodal model inputs.""" + input_ids = input_batch.input_ids + inputs_embeds = None + ec_connector_output = None + if not self.supports_mm_inputs or not self.is_first_pp_rank: + return input_ids, inputs_embeds, ec_connector_output + + if dummy_run: + inputs_embeds = self.model_state.dummy_inputs_embeds(input_batch.num_tokens_after_padding) + else: + scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs + if self.lora_config is not None: + set_active_mm_loras( + model=self.model, + lora_manager=self.lora_manager, + encoder_cache=self.encoder_cache, + req_id_to_index=self.req_states.req_id_to_index, + lora_state=self.lora_state, + scheduled_encoder_inputs=scheduled_encoder_inputs, + ) + with self.ec_connector.maybe_get_output(scheduler_output) as ec_connector_output: + inputs_embeds = self._get_mm_embeddings( + scheduled_encoder_inputs, + input_batch, + ) + if inputs_embeds is not None and not self.model.requires_raw_input_tokens: + input_ids = None + return input_ids, inputs_embeds, ec_connector_output + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._validate_parallel_support() + self._omni_data_plane = ( + OmniRunnerDataPlane(self.vllm_config, self.model_config) + if uses_native_mrv2_data_plane( + self.model_config, + use_v2_model_runner=True, + ) + else None + ) + + def _validate_parallel_support(self) -> None: + parallel_config = self.vllm_config.parallel_config + if parallel_config.pipeline_parallel_size > 1: + raise NotImplementedError( + "Omni Model Runner V2 does not support pipeline parallel; " + "use model_runner: v1 when pipeline_parallel_size > 1" + ) + if getattr(parallel_config, "prefill_context_parallel_size", 1) > 1: + raise NotImplementedError( + "Omni Model Runner V2 does not support prefill context parallelism; " + "use model_runner: v1 when prefill_context_parallel_size > 1" + ) + + def add_requests(self, scheduler_output: SchedulerOutput) -> None: + # The upstream Sampler stages static parameters only in add_request. + # With no new requests, its UVA snapshots remain valid and immutable. + # Keep the upstream path for custom/speculative samplers, whose staged + # writes may have a different lifecycle. + if not scheduler_output.scheduled_new_reqs and type(getattr(self, "sampler", None)) is Sampler: + return + logits_processor = getattr(self.model, "logits_processor", None) + logits_vocab = getattr(logits_processor, "vocab_size", None) + if isinstance(logits_vocab, int) and logits_vocab > 0: + for request_data in scheduler_output.scheduled_new_reqs: + sampling_params = request_data.sampling_params + if sampling_params is not None: + sanitize_sampling_params_min_tokens_stop_ids( + sampling_params, + logits_vocab, + ) + super().add_requests(scheduler_output) + + def shutdown(self) -> None: + # Every owner must release its resources even if an earlier drain fails. + try: + worker = getattr(self, "_native_output_materializer", None) + if worker is not None: + worker.close() + finally: + self._native_output_materializer = None + try: + if self._omni_data_plane is not None: + self._omni_data_plane.close() + finally: + super().shutdown() + + def _prepare_native_data_plane(self, scheduler_output: SchedulerOutput) -> None: + plane = getattr(self, "_omni_data_plane", None) + if plane is None: + return + for request_data in getattr(scheduler_output, "scheduled_new_reqs", []): + if str(getattr(request_data, "req_id", "")).startswith("_warmup_"): + continue + plane.register_request(request_data) + plane.register_receivers(list(getattr(scheduler_output, "pending_input_registrations", []))) + natural_terminal_req_ids = set(getattr(scheduler_output, "data_plane_terminal_req_ids", set())) + aborted_req_ids = set(getattr(scheduler_output, "finished_req_ids", set())).difference(natural_terminal_req_ids) + if natural_terminal_req_ids: + plane.request_terminal(natural_terminal_req_ids) + if aborted_req_ids: + plane.abort_requests(aborted_req_ids) + + def _sync_native_data_plane_payloads( + self, + scheduler_output: SchedulerOutput, + ) -> None: + plane = getattr(self, "_omni_data_plane", None) + if plane is None: + return + gpu_keys: set[tuple[str, str]] = getattr(self.model, "gpu_resident_buffer_keys", set()) + for req_id in scheduler_output.num_scheduled_tokens: + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + continue + payload = plane.pop_local_stage_payload(req_id) + if isinstance(payload, dict) and payload: + self.model_state.intermediate_buffer.update( + req_idx, + payload, + gpu_keys, + ) + + def _uses_native_output_materializer(self) -> bool: + parallel = getattr(getattr(self, "vllm_config", None), "parallel_config", None) + # Multi-rank executors already consume outputs asynchronously and choose + # the replying rank themselves. Preserve that ownership contract. + return getattr(parallel, "tensor_parallel_size", 1) == 1 and getattr(self, "_omni_data_plane", None) is not None + + def _materialize_native_output(self, output: Any) -> Any: + from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker + + worker = getattr(self, "_native_output_materializer", None) + if worker is None: + worker = NativeOutputWorker(getattr(self.vllm_config, "max_concurrent_batches", 2), self.device) + self._native_output_materializer = worker + return worker.submit(output, self._omni_data_plane) + + def _finalize_native_data_plane_output(self, output: Any) -> Any: + plane = getattr(self, "_omni_data_plane", None) + if plane is None or output is None: + return output + # MRv2's data plane consumes the host-materialized token lists from + # the runner output. Do not let a future connector path fall back to + # reading sampler tensors and synchronizing the device here. + if hasattr(output, "sampled_token_ids_materialized") and getattr(output, "sampled_token_ids", None) is not None: + if not output.sampled_token_ids_materialized: + raise RuntimeError("MRv2 native data plane received non-materialized sampled_token_ids") + inter_stage_outputs = getattr(output, "inter_stage_outputs", None) + plane.enqueue_outputs( + req_ids=list(getattr(output, "req_ids", [])), + inter_stage_outputs=inter_stage_outputs, + sampled_token_ids=getattr(output, "sampled_token_ids", None), + ) + output.inter_stage_outputs = None + output.omni_connector_output = plane.get_omni_connector_output() + return output + + def _reserve_native_data_plane_outputs(self, req_ids: list[str]) -> None: + plane = getattr(self, "_omni_data_plane", None) + if plane is not None: + plane.reserve_outputs(req_ids) + + def _attach_native_data_plane_signals(self, output: Any) -> Any: + plane = getattr(self, "_omni_data_plane", None) + if plane is None: + return output + return plane.attach_omni_connector_output(output) + + def load_model(self, *args: Any, **kwargs: Any) -> None: + import vllm.v1.worker.gpu.model_runner as _mr_module + + with _model_state_patch_lock: + _orig = _mr_module.init_model_state + _mr_module.init_model_state = init_omni_model_state + try: + super().load_model(*args, **kwargs) + finally: + _mr_module.init_model_state = _orig + self._last_aux_output = None + self._last_multimodal_outputs = None + self._configure_cudagraph_output_contract() + + # Preprocess models own embedding buffers; encoder_runner sizing would mismatch. + if getattr(self.model, "has_preprocess", False) and self.supports_mm_inputs: + self.supports_mm_inputs = False + self.encoder_cache = None + + # ------------------------------------------------------------------ + # CUDA Graph: conditionally exclude FULL mode + # ------------------------------------------------------------------ + + def _exclude_unsupported_full_graphs(self) -> None: + # vLLM 0.29 keys dispatch candidates by (uniform token count, LoRA count). + manager = self.cudagraph_manager + assert manager is not None + manager._capture_descs.pop(CUDAGraphMode.FULL, None) + manager._candidates = { + key: [desc for desc in candidates if desc.cg_mode != CUDAGraphMode.FULL] + for key, candidates in manager._candidates.items() + } + logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.") + + def capture_model(self) -> int: + """Handle CUDA graph capture for Omni models. + + Tuple-returning models use PIECEWISE graphs; FULL replay requires + a tensor output. + + For PIECEWISE capture, the warmup pass runs with + ``CUDAGraphMode.NONE`` which hits ``torch.empty_like(hidden_states)`` + in the cudagraph framework. If the model returns a tuple, that call + crashes. We temporarily wrap the model's forward to extract only the + tensor part during capture, then restore the original forward. + """ + if self._exclude_full_graph: + self._exclude_unsupported_full_graphs() + + # Wrap model forward during capture so tuple returns don't crash + # torch.empty_like() in the PIECEWISE warmup pass. + if self._model_returns_tuple: + original_forward = self.model.forward + + def _capture_forward(*args: Any, **kwargs: Any) -> Any: + output = original_forward(*args, **kwargs) + return self._prepare_cudagraph_capture_output(output) + + self.model.forward = _capture_forward # type: ignore[assignment] + try: + result = super().capture_model() + finally: + self.model.forward = original_forward # type: ignore[assignment] + else: + result = super().capture_model() + + capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None) + if callable(capture_mtp): + capture_mtp(self._dispatch_mtp_batch_descriptor) + return result + + def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any: + capture_sizes = self.model_state._get_mtp_capture_sizes() + captured_bucket = next( + (size for size in sorted(capture_sizes) if num_mtp_reqs <= size <= self.scheduler_config.max_num_seqs), + None, + ) + if captured_bucket is None: + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_mtp_reqs, + num_reqs=num_mtp_reqs, + ) + return self.cudagraph_manager.dispatch( + captured_bucket, + captured_bucket, + 1, + 0, + ) + + def _dispatch_batch_descriptor( + self, + *, + num_reqs: int, + num_toks: int, + uniform_tok_count: int | None, + num_active_loras: int, + use_eager: bool, + max_query_len: int, + ): + return dispatch_cg_and_sync_dp( + self.cudagraph_manager, + num_reqs, + num_toks, + uniform_tok_count, + self.dp_size, + self.dp_rank, + max_query_len=max_query_len, + need_eager=use_eager, + num_active_loras=num_active_loras, + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + context_len: int = 0, + ) -> Any: + if not dummy_run: + self._prepare_native_data_plane(scheduler_output) + self.finish_requests(scheduler_output) + self.free_states(scheduler_output) + self.add_requests(scheduler_output) + self.update_requests(scheduler_output) + self._sync_native_data_plane_payloads(scheduler_output) + self.block_tables.apply_staged_writes() + if scheduler_output.total_num_scheduled_tokens == 0: + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._attach_native_data_plane_signals( + self._merge_ec_connector_no_forward(scheduler_output, empty_output) + ) + + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + batch_req_state, uniform_tok_count = self.gather_batch_req_state(scheduler_output, dummy_run) + if batch_req_state is not None: + num_toks = batch_req_state.num_tokens + num_active_loras = 0 + if self.lora_config: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + num_active_loras = get_num_active_loras_for_dispatch( + self.lora_config, + self.lora_state, + req_ids, + dummy_run, + ) + # Encoder-decoder models: disable compilation when encoder inputs + # are scheduled (dynamic cross-attention cache updates). + skip_compiled = self.is_encoder_decoder and bool(scheduler_output.scheduled_encoder_inputs) + batch_desc, dp_sync = self._dispatch_batch_descriptor( + num_reqs=num_reqs, + num_toks=num_toks, + uniform_tok_count=uniform_tok_count, + num_active_loras=num_active_loras, + use_eager=is_profile or skip_compiled, + max_query_len=max_query_len, + ) + + if batch_desc.num_tokens == 0: + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._attach_native_data_plane_signals( + self._merge_ec_connector_no_forward(scheduler_output, empty_output) + ) + + if not dummy_run: + assert batch_req_state is not None + input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc) + block_tables, slot_mappings = self.prepare_attn(input_batch) + self.model_state.preprocess_state( + input_batch, + block_tables, + self.kv_cache_config, + self.req_states.num_computed_tokens.gpu, + ) + + if self.lora_config: + lora_inputs = self.lora_state.make_lora_inputs( + input_batch.req_ids, + input_batch.idx_mapping_np, + input_batch.num_scheduled_tokens, + ) + self._set_active_loras(*lora_inputs) + else: + from vllm.v1.worker.gpu.input_batch import InputBatch + + input_batch = InputBatch.make_dummy( + batch_desc.num_reqs or num_reqs, + batch_desc.num_tokens, + self.input_buffers, + max_query_len=batch_desc.max_query_len, + ) + if not skip_attn_for_dummy_run: + block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) + if context_len: + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + ) + else: + assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( + "Attention metadata is required for FULL CUDA graph dummy runs" + ) + block_tables = None + slot_mappings = None + + attn_metadata = None + slot_mappings_by_layer = None + if not (dummy_run and skip_attn_for_dummy_run): + assert slot_mappings is not None + slot_mappings_by_layer = build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config) + assert block_tables is not None + attn_metadata = self.model_state.prepare_attn( + input_batch, + batch_desc.cg_mode, + block_tables, + slot_mappings, + self.attn_groups, + self.kv_cache_config, + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, + ) + input_ids, inputs_embeds, ec_connector_output = self._prepare_mm_inputs( + scheduler_output, + input_batch, + dummy_run=dummy_run, + ) + + model_inputs: dict[str, Any] = { + "input_ids": input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": intermediate_tensors, + **self.model_state.prepare_inputs(input_batch, self.req_states), + } + self._add_legacy_forward_inputs(model_inputs, input_batch) + # ★ PRE-FORWARD: per-request preprocess + batched MTP. + # Runs for ALL graph modes (FULL, PIECEWISE, NONE). + # For FULL graph: OmniModelState provides a static inputs_embeds + # buffer that was captured by the graph. Preprocess writes + # in-place to this buffer, and FULL graph replay reads the + # updated values from the same tensor address. + if not dummy_run: + self.model_state.run_preprocess( + input_batch, + model_inputs, + self.req_states, + self._dispatch_mtp_batch_descriptor, + ) + + self.eplb.prepare_forward(self.model_config, input_batch.num_tokens) + + # --- Model forward --- + if batch_desc.cg_mode == CUDAGraphMode.FULL: + # FULL graph replay. Preprocess already wrote to the static + # inputs_embeds buffer above. + assert self.cudagraph_manager is not None + self.kv_connector.pre_forward(scheduler_output) + hidden_states = self.cudagraph_manager.run_fullgraph(batch_desc) + self._last_aux_output = None + self._last_multimodal_outputs = None + if hasattr(self.model, "_last_captured_layers"): + self.model._last_captured_layers = self._last_aux_output + else: + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=self.lora_config is not None, + num_active_loras=batch_desc.num_active_loras, + ) + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=batch_desc.cg_mode, + num_tokens_across_dp=dp_sync.num_tokens_across_dp if dp_sync is not None else None, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + skip_compiled=skip_compiled, + is_padding=input_batch.is_padding, + ): + self.kv_connector.pre_forward(scheduler_output) + if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: + assert self.cudagraph_manager is not None + model_output = self.cudagraph_manager.run_pw_graph( + self.model, + model_inputs, + ) + else: + model_output = self.model(**model_inputs) + + # Extract hidden_states from model output. + self._last_aux_output = None + self._last_multimodal_outputs = None + if isinstance(model_output, OmniOutput): + hidden_states = model_output.text_hidden_states + if model_output.multimodal_outputs: + self._last_multimodal_outputs = cast(dict[str, Any], model_output.multimodal_outputs) + elif isinstance(model_output, tuple) and len(model_output) == 2: + hidden_states, self._last_aux_output = model_output + if hasattr(self.model, "_last_captured_layers"): + self.model._last_captured_layers = self._last_aux_output + else: + hidden_states = model_output + if not dummy_run and isinstance(hidden_states, torch.Tensor): + self.model_state.run_postprocess(hidden_states, input_batch) + + routed_experts = None + if not dummy_run and (capturer := self.routed_experts_capturer) is not None: + assert slot_mappings is not None + routed_experts = capturer.get_routed_experts(slot_mappings, num_toks) + + self.execute_model_state = ExecuteModelState( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + hidden_states=hidden_states, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + dp_sync=dp_sync, + ec_connector_output=ec_connector_output, + routed_experts=routed_experts, + ) + + assert isinstance(hidden_states, torch.Tensor) + return None + + # ------------------------------------------------------------------ + # Request lifecycle: update intermediate buffer from cached requests + # ------------------------------------------------------------------ + + def update_requests(self, scheduler_output: SchedulerOutput) -> None: + """Merge updated additional_information into intermediate_buffer. + + In async_chunk mode, chunk_transfer_adapter attaches updated + additional_information (e.g. thinker_decode_embeddings) to + OmniCachedRequestData for cached requests every schedule step. + Upstream GPUModelRunner.update_requests does not handle this + field, so we merge it into the intermediate buffer here. + """ + super().update_requests(scheduler_output) + + cached = scheduler_output.scheduled_cached_reqs + addl_info = getattr(cached, "additional_information", None) + if not addl_info: + return + for req_id, info in addl_info.items(): + if info is None: + continue + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + continue + resolved = _resolve_additional_information(info) + if resolved: + gpu_keys: set[tuple[str, str]] = getattr(self.model, "gpu_resident_buffer_keys", set()) + self.model_state.intermediate_buffer.update(req_idx, resolved, gpu_keys) + + # ------------------------------------------------------------------ + # Request lifecycle: clean up intermediate buffer on finish + # ------------------------------------------------------------------ + + def finish_requests(self, scheduler_output: SchedulerOutput) -> None: + # IMPORTANT: Must query req_id_to_index BEFORE super().finish_requests() + # because super() calls req_states.remove_request(req_id) which pops the + # mapping and returns the slot index to free_indices. + finished = scheduler_output.finished_req_ids + on_finished = getattr(self.model, "on_requests_finished", None) + if finished and callable(on_finished): + # Request-owned state outlives generation slots and chunk boundaries. + # Notify even if the last chunk already released its runner slot. + on_finished(finished) + preempted = scheduler_output.preempted_req_ids + 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) + if idx is not None: + self.model_state.remove_request(idx) + super().finish_requests(scheduler_output) diff --git a/vllm_omni/worker_v2/output_snapshot.py b/vllm_omni/worker_v2/output_snapshot.py new file mode 100644 index 000000000..31187ddf7 --- /dev/null +++ b/vllm_omni/worker_v2/output_snapshot.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Owned, dtype-grouped snapshots for asynchronous runner outputs.""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable +from typing import Any + +import torch +from torch.utils._pytree import tree_flatten, tree_unflatten + + +class PackedOutputSnapshot(dict): + """A normal payload mapping with an internal batched-copy plan. + + The runner's snapshot-slot event protects the device slabs until D2H is + complete. Host slabs are allocated per output, so downstream consumers can + retain their views without depending on reuse of the device ring slot. + """ + + def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None: + super().__init__(tree_unflatten(leaves, spec)) + self._leaves = leaves + self._spec = spec + self._groups = groups + + def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]: + leaves = list(self._leaves) + for indices, slab in self._groups: + cpu = copy_tensor(slab) + offset = 0 + for index in indices: + value = leaves[index] + size = value.numel() + leaves[index] = cpu.narrow(0, offset, size).view(value.shape) + offset += size + return tree_unflatten(leaves, self._spec) + + +def pack_output_snapshot( + payload: dict[str, Any], slot: dict[tuple[Any, ...], torch.Tensor], *, max_buckets: int +) -> PackedOutputSnapshot | None: + """Copy tensor leaves once per dtype/device, retaining their exact values. + + No CUDA synchronization is introduced here. The caller must select the + producer stream and wait for the previous consumer before reusing a slot. + """ + leaves, spec = tree_flatten(payload) + grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list) + for index, value in enumerate(leaves): + if isinstance(value, torch.Tensor): + # Preserve the existing per-tensor path for unusual layouts. + if value.layout != torch.strided or value.is_quantized: + return None + grouped[(value.dtype, value.device)].append(index) + groups = [] + for (dtype, device), indices in grouped.items(): + values = [leaves[index].detach().reshape(-1) for index in indices] + size = sum(value.numel() for value in values) + key = ("packed-output", dtype, device, size) + slab = slot.get(key) + if slab is None: + slab = torch.empty(size, dtype=dtype, device=device) + if len(slot) < max_buckets: + slot[key] = slab + if len(values) == 1: + slab.copy_(values[0]) + else: + torch.cat(values, out=slab) + offset = 0 + for index, value in zip(indices, values): + shape = leaves[index].shape + leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) + offset += value.numel() + groups.append((indices, slab)) + return PackedOutputSnapshot(leaves, spec, groups)